"""
title: Autogen Studio Tool
author: matthewh
version: 1.0
required_open_webui_version: 0.3.9
Instructions:
1. Install and run AutoGen Studio:
```bash
pip install autogenstudio flaml[automl] matplotlib
autogenstudio ui --port=8081 --docs --host=0.0.0.0 --workers=2
```
2. Export the workflow and serve it as an endpoint:
```bash
autogenstudio serve --workflow=workflow.json --port=8082 --host=0.0.0.0 --docs --workers=2
```
3. Use http://localhost:8082/docs to interact with the endpoint.
4. This tool sends a query to the AutoGen API and displays the response correctly in the chat.
"""
import os
import requests
import urllib.parse
from typing import Optional, Dict # Ensure Dict is imported
from pydantic import BaseModel, Field
class Tools:
"""
Tools class to house the AutoGen tool functions, compatible with Open WebUI.
"""
class Valves(BaseModel):
"""
Configurable parameters (valves) for the AutoGen tool.
"""
AUTOGEN_BASE_URL: str = Field(
default=os.getenv(
"AUTOGEN_BASE_URL", "http://host.docker.internal:8082/predict"
),
description="Base URL for the AutoGen endpoint.",
)
request_timeout: int = Field(
default=300, description="Timeout for the HTTP request (in seconds)."
)
debug: bool = Field(
default=False, description="Enable or disable debug logging."
)
def __init__(self):
self.valves = self.Valves()
def autogen_tool(self, query: str) -> Dict[str, str]:
"""
AutoGen tool that performs a query using the AutoGen API.
:param query: User’s input query.
:return: A dictionary to return a formatted response for the LLM.
"""
if self.valves.debug:
print(f"[DEBUG] Starting tool with query: {query}")
encoded_prompt = urllib.parse.quote(query, safe="")
url = f"{self.valves.AUTOGEN_BASE_URL}/{encoded_prompt}"
headers = {"Accept": "application/json"}
timeout = self.valves.request_timeout
if self.valves.debug:
print(f"[DEBUG] Sending request to {url}")
try:
response = requests.get(url, headers=headers, timeout=timeout)
response.raise_for_status()
data = response.json()
if self.valves.debug:
print(f"[DEBUG] Response data: {data}")
# Format the response as an LLM-friendly prompt
if isinstance(data, dict) and "response" in data:
message = data["response"]
formatted_message = (
f"Provide the following response to the user:\n"
f"```autogen\n{message}\n```"
)
return {"content": formatted_message}
return {"content": "Unexpected response format."}
except Exception as e:
error_message = f"Error during request: {str(e)}"
if self.valves.debug:
print(f"[DEBUG] {error_message}")
return {"content": error_message}
# # Example usage for local testing
# if __name__ == "__main__":
# tools = Tools() # Initialize without custom valves
# query = "example query"
# result = tools.autogen_tool(query)
# print(result["content"])
