"""
title: OpenWebUI SSH Connection Manager (Async)
author: Wes Caldwell
email:
[email protected]
author_url: https://github.com/musicheardworldwide
version: 1.0.0
license: MIT
description: A structured Open WebUI tool that fetches the server's IP for testing purposes.
requirements:
"""
import json
import logging
import traceback
import socket
import requests
from typing import Optional, Callable, Any, Dict
from pydantic import BaseModel, Field
from fastapi import Request
from open_webui.utils.chat import generate_chat_completion
from open_webui.utils.misc import get_last_user_message
from open_webui.models.users import User
# Configure Logging
logger = logging.getLogger(__name__)
if not logger.handlers:
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
logger.addHandler(handler)
logger.setLevel(logging.INFO)
# Centralized Error Handling Function
def handle_error(exception: Exception, function_name: str, inputs: dict) -> dict:
"""
Handles errors and returns a structured response for OpenWebUI.
Args:
exception (Exception): The caught exception.
function_name (str): The name of the function where the error occurred.
inputs (dict): The original function inputs for debugging.
Returns:
dict: A structured error message to pass to OpenWebUI.
"""
error_message = str(exception)
stack_trace = traceback.format_exc()
logger.error(f"Error in {function_name}: {error_message}")
logger.debug(f"Stack Trace:\n{stack_trace}")
return {
"error": True,
"function": function_name,
"message": error_message,
"stack_trace": stack_trace,
"inputs": inputs,
"suggestion": "Check input values and ensure the correct API configurations.",
}
# Function to get IP addresses
def get_ip_addresses():
"""Retrieve both public and private IP addresses."""
try:
# Get Private IP
private_ip = socket.gethostbyname(socket.gethostname())
# Get Public IP
response = requests.get("https://api.ipify.org?format=json", timeout=5)
public_ip = response.json().get("ip", "Unknown")
return {"private_ip": private_ip, "public_ip": public_ip}
except Exception as e:
logger.error(f"IP Retrieval Error: {e}")
return {"private_ip": "Unknown", "public_ip": "Unknown"}
# Tool Definition
class Tools:
"""
OpenWebUI Tool with request handling, logging, and basic IP retrieval.
"""
class Config(BaseModel):
ENABLE_LOGGING: bool = Field(
default=True, description="Enable or disable logging for debugging"
)
def __init__(self):
self.config = self.Config()
async def process_request(
self, query: str, __event_emitter__: Optional[Callable[[Any], Any]] = None
) -> Dict:
"""
Simulated API request. Fetches server's IP instead.
Args:
query (str): The search query.
__event_emitter__ (Callable): Emits events back to OpenWebUI.
Returns:
Dict: API response containing IP addresses.
"""
try:
if __event_emitter__:
await __event_emitter__(
{
"type": "status",
"data": {"description": "Fetching server IP...", "done": False},
}
)
ip_info = get_ip_addresses()
logger.info(f"Fetched IP Info: {ip_info}")
response = {
"status": "success",
"query": query,
"ip_info": ip_info,
}
if __event_emitter__:
await __event_emitter__(
{
"type": "status",
"data": {"description": "IP retrieval complete.", "done": True},
}
)
return response
except Exception as e:
return handle_error(e, "process_request", {"query": query})
async def pipe(self, body: Dict, __user__: Dict, __request__: Request) -> Dict:
"""
Processes requests and injects IP information.
Args:
body (Dict): The request payload.
__user__ (Dict): User metadata.
__request__ (Request): FastAPI request object.
Returns:
Dict: Response with IP information.
"""
try:
logger.info("Processing request in Pipe...")
messages = body.get("messages", [])
if not messages:
return handle_error(ValueError("No input messages found"), "pipe", body)
# Get IP Information
ip_info = get_ip_addresses()
# Modify Response
body["messages"].append(
{"role": "assistant", "content": json.dumps({"ip_info": ip_info})}
)
return body
except Exception as e:
return handle_error(e, "pipe", body)
# Example Usage
if __name__ == "__main__":
tool = Tools()
async def test_tool():
response = await tool.process_request("Get IP")
print(json.dumps(response, indent=4))
import asyncio
asyncio.run(test_tool())
