""" title: OpenWebUI SSH Connection Manager (Async) author: Wes Caldwell email: musicheardworldwide@gmail.com 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())