mirror of
https://github.com/microsoft/autogen.git
synced 2026-04-20 03:02:16 -04:00
v1 of AutoGen Studio on AgentChat (#4097)
* add skeleton worflow manager * add test notebook * update test nb * add sample team spec * refactor requirements to agentchat and ext * add base provider to return agentchat agents from json spec * initial api refactor, update dbmanager * api refactor * refactor tests * ags api tutorial update * ui refactor * general refactor * minor refactor updates * backend api refaactor * ui refactor and update * implement v1 for streaming connection with ui updates * backend refactor * ui refactor * minor ui tweak * minor refactor and tweaks * general refactor * update tests * sync uv.lock with main * uv lock update
This commit is contained in:
Victor Dibia
@@ -1,93 +1,63 @@
|
||||
import asyncio
|
||||
# api/app.py
|
||||
import os
|
||||
import queue
|
||||
import threading
|
||||
import traceback
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, Union
|
||||
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||||
# import logging
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncGenerator
|
||||
from loguru import logger
|
||||
from openai import OpenAIError
|
||||
|
||||
from ..chatmanager import AutoGenChatManager
|
||||
from ..database import workflow_from_id
|
||||
from ..database.dbmanager import DBManager
|
||||
from ..datamodel import Agent, Message, Model, Response, Session, Skill, Workflow
|
||||
from ..profiler import Profiler
|
||||
from ..utils import check_and_cast_datetime_fields, init_app_folders, sha256_hash, test_model
|
||||
from .routes import sessions, runs, teams, agents, models, tools, ws
|
||||
from .deps import init_managers, cleanup_managers
|
||||
from .config import settings
|
||||
from .initialization import AppInitializer
|
||||
from ..version import VERSION
|
||||
from ..websocket_connection_manager import WebSocketConnectionManager
|
||||
|
||||
profiler = Profiler()
|
||||
managers = {"chat": None} # manage calls to autogen
|
||||
# Create thread-safe queue for messages between api thread and autogen threads
|
||||
message_queue = queue.Queue()
|
||||
active_connections = []
|
||||
active_connections_lock = asyncio.Lock()
|
||||
websocket_manager = WebSocketConnectionManager(
|
||||
active_connections=active_connections,
|
||||
active_connections_lock=active_connections_lock,
|
||||
)
|
||||
|
||||
|
||||
def message_handler():
|
||||
while True:
|
||||
message = message_queue.get()
|
||||
logger.info(
|
||||
"** Processing Agent Message on Queue: Active Connections: "
|
||||
+ str([client_id for _, client_id in websocket_manager.active_connections])
|
||||
+ " **"
|
||||
)
|
||||
for connection, socket_client_id in websocket_manager.active_connections:
|
||||
if message["connection_id"] == socket_client_id:
|
||||
logger.info(
|
||||
f"Sending message to connection_id: {message['connection_id']}. Connection ID: {socket_client_id}"
|
||||
)
|
||||
asyncio.run(websocket_manager.send_message(message, connection))
|
||||
else:
|
||||
logger.info(
|
||||
f"Skipping message for connection_id: {message['connection_id']}. Connection ID: {socket_client_id}"
|
||||
)
|
||||
message_queue.task_done()
|
||||
|
||||
|
||||
message_handler_thread = threading.Thread(target=message_handler, daemon=True)
|
||||
message_handler_thread.start()
|
||||
# Configure logging
|
||||
# logger = logging.getLogger(__name__)
|
||||
# logging.basicConfig(level=logging.INFO)
|
||||
|
||||
|
||||
# Initialize application
|
||||
app_file_path = os.path.dirname(os.path.abspath(__file__))
|
||||
folders = init_app_folders(app_file_path)
|
||||
ui_folder_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ui")
|
||||
|
||||
database_engine_uri = folders["database_engine_uri"]
|
||||
dbmanager = DBManager(engine_uri=database_engine_uri)
|
||||
|
||||
HUMAN_INPUT_TIMEOUT_SECONDS = 180
|
||||
initializer = AppInitializer(settings, app_file_path)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
print("***** App started *****")
|
||||
managers["chat"] = AutoGenChatManager(
|
||||
message_queue=message_queue,
|
||||
websocket_manager=websocket_manager,
|
||||
human_input_timeout=HUMAN_INPUT_TIMEOUT_SECONDS,
|
||||
)
|
||||
dbmanager.create_db_and_tables()
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""
|
||||
Lifecycle manager for the FastAPI application.
|
||||
Handles initialization and cleanup of application resources.
|
||||
"""
|
||||
# Startup
|
||||
logger.info("Initializing application...")
|
||||
try:
|
||||
# Initialize managers (DB, Connection, Team)
|
||||
await init_managers(initializer.database_uri, initializer.config_dir)
|
||||
logger.info("Managers initialized successfully")
|
||||
|
||||
yield
|
||||
# Close all active connections
|
||||
await websocket_manager.disconnect_all()
|
||||
print("***** App stopped *****")
|
||||
# Any other initialization code
|
||||
logger.info("Application startup complete")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize application: {str(e)}")
|
||||
raise
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
yield # Application runs here
|
||||
|
||||
# Shutdown
|
||||
try:
|
||||
logger.info("Cleaning up application resources...")
|
||||
await cleanup_managers()
|
||||
logger.info("Application shutdown complete")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during shutdown: {str(e)}")
|
||||
|
||||
# allow cross origin requests for testing on localhost:800* ports only
|
||||
# Create FastAPI application
|
||||
app = FastAPI(lifespan=lifespan, debug=True)
|
||||
|
||||
# CORS middleware configuration
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[
|
||||
@@ -101,412 +71,114 @@ app.add_middleware(
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
show_docs = os.environ.get("AUTOGENSTUDIO_API_DOCS", "False").lower() == "true"
|
||||
docs_url = "/docs" if show_docs else None
|
||||
# Create API router with version and documentation
|
||||
api = FastAPI(
|
||||
root_path="/api",
|
||||
title="AutoGen Studio API",
|
||||
version=VERSION,
|
||||
docs_url=docs_url,
|
||||
description="AutoGen Studio is a low-code tool for building and testing multi-agent workflows using AutoGen.",
|
||||
description="AutoGen Studio is a low-code tool for building and testing multi-agent workflows.",
|
||||
docs_url="/docs" if settings.API_DOCS else None,
|
||||
)
|
||||
# mount an api route such that the main route serves the ui and the /api
|
||||
app.mount("/api", api)
|
||||
|
||||
app.mount("/", StaticFiles(directory=ui_folder_path, html=True), name="ui")
|
||||
api.mount(
|
||||
"/files",
|
||||
StaticFiles(directory=folders["files_static_root"], html=True),
|
||||
name="files",
|
||||
# Include all routers with their prefixes
|
||||
api.include_router(
|
||||
sessions.router,
|
||||
prefix="/sessions",
|
||||
tags=["sessions"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
runs.router,
|
||||
prefix="/runs",
|
||||
tags=["runs"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
teams.router,
|
||||
prefix="/teams",
|
||||
tags=["teams"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
agents.router,
|
||||
prefix="/agents",
|
||||
tags=["agents"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
models.router,
|
||||
prefix="/models",
|
||||
tags=["models"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
tools.router,
|
||||
prefix="/tools",
|
||||
tags=["tools"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
api.include_router(
|
||||
ws.router,
|
||||
prefix="/ws",
|
||||
tags=["websocket"],
|
||||
responses={404: {"description": "Not found"}},
|
||||
)
|
||||
|
||||
|
||||
# manage websocket connections
|
||||
|
||||
|
||||
def create_entity(model: Any, model_class: Any, filters: dict = None):
|
||||
"""Create a new entity"""
|
||||
model = check_and_cast_datetime_fields(model)
|
||||
try:
|
||||
response: Response = dbmanager.upsert(model)
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
except Exception as ex_error:
|
||||
print(ex_error)
|
||||
return {
|
||||
"status": False,
|
||||
"message": f"Error occurred while creating {model_class.__name__}: " + str(ex_error),
|
||||
}
|
||||
|
||||
|
||||
def list_entity(
|
||||
model_class: Any,
|
||||
filters: dict = None,
|
||||
return_json: bool = True,
|
||||
order: str = "desc",
|
||||
):
|
||||
"""List all entities for a user"""
|
||||
return dbmanager.get(model_class, filters=filters, return_json=return_json, order=order)
|
||||
|
||||
|
||||
def delete_entity(model_class: Any, filters: dict = None):
|
||||
"""Delete an entity"""
|
||||
|
||||
return dbmanager.delete(filters=filters, model_class=model_class)
|
||||
|
||||
|
||||
@api.get("/skills")
|
||||
async def list_skills(user_id: str):
|
||||
"""List all skills for a user"""
|
||||
filters = {"user_id": user_id}
|
||||
return list_entity(Skill, filters=filters)
|
||||
|
||||
|
||||
@api.post("/skills")
|
||||
async def create_skill(skill: Skill):
|
||||
"""Create a new skill"""
|
||||
filters = {"user_id": skill.user_id}
|
||||
return create_entity(skill, Skill, filters=filters)
|
||||
|
||||
|
||||
@api.delete("/skills/delete")
|
||||
async def delete_skill(skill_id: int, user_id: str):
|
||||
"""Delete a skill"""
|
||||
filters = {"id": skill_id, "user_id": user_id}
|
||||
return delete_entity(Skill, filters=filters)
|
||||
|
||||
|
||||
@api.get("/models")
|
||||
async def list_models(user_id: str):
|
||||
"""List all models for a user"""
|
||||
filters = {"user_id": user_id}
|
||||
return list_entity(Model, filters=filters)
|
||||
|
||||
|
||||
@api.post("/models")
|
||||
async def create_model(model: Model):
|
||||
"""Create a new model"""
|
||||
return create_entity(model, Model)
|
||||
|
||||
|
||||
@api.post("/models/test")
|
||||
async def test_model_endpoint(model: Model):
|
||||
"""Test a model"""
|
||||
try:
|
||||
response = test_model(model)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Model tested successfully",
|
||||
"data": response,
|
||||
}
|
||||
except (OpenAIError, Exception) as ex_error:
|
||||
return {
|
||||
"status": False,
|
||||
"message": "Error occurred while testing model: " + str(ex_error),
|
||||
}
|
||||
|
||||
|
||||
@api.delete("/models/delete")
|
||||
async def delete_model(model_id: int, user_id: str):
|
||||
"""Delete a model"""
|
||||
filters = {"id": model_id, "user_id": user_id}
|
||||
return delete_entity(Model, filters=filters)
|
||||
|
||||
|
||||
@api.get("/agents")
|
||||
async def list_agents(user_id: str):
|
||||
"""List all agents for a user"""
|
||||
filters = {"user_id": user_id}
|
||||
return list_entity(Agent, filters=filters)
|
||||
|
||||
|
||||
@api.post("/agents")
|
||||
async def create_agent(agent: Agent):
|
||||
"""Create a new agent"""
|
||||
return create_entity(agent, Agent)
|
||||
|
||||
|
||||
@api.delete("/agents/delete")
|
||||
async def delete_agent(agent_id: int, user_id: str):
|
||||
"""Delete an agent"""
|
||||
filters = {"id": agent_id, "user_id": user_id}
|
||||
return delete_entity(Agent, filters=filters)
|
||||
|
||||
|
||||
@api.post("/agents/link/model/{agent_id}/{model_id}")
|
||||
async def link_agent_model(agent_id: int, model_id: int):
|
||||
"""Link a model to an agent"""
|
||||
return dbmanager.link(link_type="agent_model", primary_id=agent_id, secondary_id=model_id)
|
||||
|
||||
|
||||
@api.delete("/agents/link/model/{agent_id}/{model_id}")
|
||||
async def unlink_agent_model(agent_id: int, model_id: int):
|
||||
"""Unlink a model from an agent"""
|
||||
return dbmanager.unlink(link_type="agent_model", primary_id=agent_id, secondary_id=model_id)
|
||||
|
||||
|
||||
@api.get("/agents/link/model/{agent_id}")
|
||||
async def get_agent_models(agent_id: int):
|
||||
"""Get all models linked to an agent"""
|
||||
return dbmanager.get_linked_entities("agent_model", agent_id, return_json=True)
|
||||
|
||||
|
||||
@api.post("/agents/link/skill/{agent_id}/{skill_id}")
|
||||
async def link_agent_skill(agent_id: int, skill_id: int):
|
||||
"""Link an a skill to an agent"""
|
||||
return dbmanager.link(link_type="agent_skill", primary_id=agent_id, secondary_id=skill_id)
|
||||
|
||||
|
||||
@api.delete("/agents/link/skill/{agent_id}/{skill_id}")
|
||||
async def unlink_agent_skill(agent_id: int, skill_id: int):
|
||||
"""Unlink an a skill from an agent"""
|
||||
return dbmanager.unlink(link_type="agent_skill", primary_id=agent_id, secondary_id=skill_id)
|
||||
|
||||
|
||||
@api.get("/agents/link/skill/{agent_id}")
|
||||
async def get_agent_skills(agent_id: int):
|
||||
"""Get all skills linked to an agent"""
|
||||
return dbmanager.get_linked_entities("agent_skill", agent_id, return_json=True)
|
||||
|
||||
|
||||
@api.post("/agents/link/agent/{primary_agent_id}/{secondary_agent_id}")
|
||||
async def link_agent_agent(primary_agent_id: int, secondary_agent_id: int):
|
||||
"""Link an agent to another agent"""
|
||||
return dbmanager.link(
|
||||
link_type="agent_agent",
|
||||
primary_id=primary_agent_id,
|
||||
secondary_id=secondary_agent_id,
|
||||
)
|
||||
|
||||
|
||||
@api.delete("/agents/link/agent/{primary_agent_id}/{secondary_agent_id}")
|
||||
async def unlink_agent_agent(primary_agent_id: int, secondary_agent_id: int):
|
||||
"""Unlink an agent from another agent"""
|
||||
return dbmanager.unlink(
|
||||
link_type="agent_agent",
|
||||
primary_id=primary_agent_id,
|
||||
secondary_id=secondary_agent_id,
|
||||
)
|
||||
|
||||
|
||||
@api.get("/agents/link/agent/{agent_id}")
|
||||
async def get_linked_agents(agent_id: int):
|
||||
"""Get all agents linked to an agent"""
|
||||
return dbmanager.get_linked_entities("agent_agent", agent_id, return_json=True)
|
||||
|
||||
|
||||
@api.get("/workflows")
|
||||
async def list_workflows(user_id: str):
|
||||
"""List all workflows for a user"""
|
||||
filters = {"user_id": user_id}
|
||||
return list_entity(Workflow, filters=filters)
|
||||
|
||||
|
||||
@api.get("/workflows/{workflow_id}")
|
||||
async def get_workflow(workflow_id: int, user_id: str):
|
||||
"""Get a workflow"""
|
||||
filters = {"id": workflow_id, "user_id": user_id}
|
||||
return list_entity(Workflow, filters=filters)
|
||||
|
||||
|
||||
@api.get("/workflows/export/{workflow_id}")
|
||||
async def export_workflow(workflow_id: int, user_id: str):
|
||||
"""Export a user workflow"""
|
||||
response = Response(message="Workflow exported successfully", status=True, data=None)
|
||||
try:
|
||||
workflow_details = workflow_from_id(workflow_id, dbmanager=dbmanager)
|
||||
response.data = workflow_details
|
||||
except Exception as ex_error:
|
||||
response.message = "Error occurred while exporting workflow: " + str(ex_error)
|
||||
response.status = False
|
||||
return response.model_dump(mode="json")
|
||||
|
||||
|
||||
@api.post("/workflows")
|
||||
async def create_workflow(workflow: Workflow):
|
||||
"""Create a new workflow"""
|
||||
return create_entity(workflow, Workflow)
|
||||
|
||||
|
||||
@api.delete("/workflows/delete")
|
||||
async def delete_workflow(workflow_id: int, user_id: str):
|
||||
"""Delete a workflow"""
|
||||
filters = {"id": workflow_id, "user_id": user_id}
|
||||
return delete_entity(Workflow, filters=filters)
|
||||
|
||||
|
||||
@api.post("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}")
|
||||
async def link_workflow_agent(workflow_id: int, agent_id: int, agent_type: str):
|
||||
"""Link an agent to a workflow"""
|
||||
return dbmanager.link(
|
||||
link_type="workflow_agent",
|
||||
primary_id=workflow_id,
|
||||
secondary_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
)
|
||||
|
||||
|
||||
@api.post("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}/{sequence_id}")
|
||||
async def link_workflow_agent_sequence(workflow_id: int, agent_id: int, agent_type: str, sequence_id: int):
|
||||
"""Link an agent to a workflow"""
|
||||
print("Sequence ID: ", sequence_id)
|
||||
return dbmanager.link(
|
||||
link_type="workflow_agent",
|
||||
primary_id=workflow_id,
|
||||
secondary_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
|
||||
|
||||
@api.delete("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}")
|
||||
async def unlink_workflow_agent(workflow_id: int, agent_id: int, agent_type: str):
|
||||
"""Unlink an agent from a workflow"""
|
||||
return dbmanager.unlink(
|
||||
link_type="workflow_agent",
|
||||
primary_id=workflow_id,
|
||||
secondary_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
)
|
||||
|
||||
|
||||
@api.delete("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}/{sequence_id}")
|
||||
async def unlink_workflow_agent_sequence(workflow_id: int, agent_id: int, agent_type: str, sequence_id: int):
|
||||
"""Unlink an agent from a workflow sequence"""
|
||||
return dbmanager.unlink(
|
||||
link_type="workflow_agent",
|
||||
primary_id=workflow_id,
|
||||
secondary_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
sequence_id=sequence_id,
|
||||
)
|
||||
|
||||
|
||||
@api.get("/workflows/link/agent/{workflow_id}")
|
||||
async def get_linked_workflow_agents(workflow_id: int):
|
||||
"""Get all agents linked to a workflow"""
|
||||
return dbmanager.get_linked_entities(
|
||||
link_type="workflow_agent",
|
||||
primary_id=workflow_id,
|
||||
return_json=True,
|
||||
)
|
||||
|
||||
|
||||
@api.get("/profiler/{message_id}")
|
||||
async def profile_agent_task_run(message_id: int):
|
||||
"""Profile an agent task run"""
|
||||
try:
|
||||
agent_message = dbmanager.get(Message, filters={"id": message_id}).data[0]
|
||||
|
||||
profile = profiler.profile(agent_message)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Agent task run profiled successfully",
|
||||
"data": profile,
|
||||
}
|
||||
except Exception as ex_error:
|
||||
return {
|
||||
"status": False,
|
||||
"message": "Error occurred while profiling agent task run: " + str(ex_error),
|
||||
}
|
||||
|
||||
|
||||
@api.get("/sessions")
|
||||
async def list_sessions(user_id: str):
|
||||
"""List all sessions for a user"""
|
||||
filters = {"user_id": user_id}
|
||||
return list_entity(Session, filters=filters)
|
||||
|
||||
|
||||
@api.post("/sessions")
|
||||
async def create_session(session: Session):
|
||||
"""Create a new session"""
|
||||
return create_entity(session, Session)
|
||||
|
||||
|
||||
@api.delete("/sessions/delete")
|
||||
async def delete_session(session_id: int, user_id: str):
|
||||
"""Delete a session"""
|
||||
filters = {"id": session_id, "user_id": user_id}
|
||||
return delete_entity(Session, filters=filters)
|
||||
|
||||
|
||||
@api.get("/sessions/{session_id}/messages")
|
||||
async def list_messages(user_id: str, session_id: int):
|
||||
"""List all messages for a use session"""
|
||||
filters = {"user_id": user_id, "session_id": session_id}
|
||||
return list_entity(Message, filters=filters, order="asc", return_json=True)
|
||||
|
||||
|
||||
@api.post("/sessions/{session_id}/workflow/{workflow_id}/run")
|
||||
async def run_session_workflow(message: Message, session_id: int, workflow_id: int):
|
||||
"""Runs a workflow on provided message"""
|
||||
try:
|
||||
user_message_history = (
|
||||
dbmanager.get(
|
||||
Message,
|
||||
filters={"user_id": message.user_id, "session_id": message.session_id},
|
||||
return_json=True,
|
||||
).data
|
||||
if session_id is not None
|
||||
else []
|
||||
)
|
||||
# save incoming message
|
||||
dbmanager.upsert(message)
|
||||
user_dir = os.path.join(folders["files_static_root"], "user", sha256_hash(message.user_id))
|
||||
os.makedirs(user_dir, exist_ok=True)
|
||||
workflow = workflow_from_id(workflow_id, dbmanager=dbmanager)
|
||||
agent_response: Message = await managers["chat"].a_chat(
|
||||
message=message,
|
||||
history=user_message_history,
|
||||
user_dir=user_dir,
|
||||
workflow=workflow,
|
||||
connection_id=message.connection_id,
|
||||
)
|
||||
|
||||
response: Response = dbmanager.upsert(agent_response)
|
||||
return response.model_dump(mode="json")
|
||||
except Exception as ex_error:
|
||||
return {
|
||||
"status": False,
|
||||
"message": "Error occurred while processing message: " + str(ex_error),
|
||||
}
|
||||
# Version endpoint
|
||||
|
||||
|
||||
@api.get("/version")
|
||||
async def get_version():
|
||||
"""Get API version"""
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Version retrieved successfully",
|
||||
"data": {"version": VERSION},
|
||||
}
|
||||
|
||||
|
||||
# websockets
|
||||
# Health check endpoint
|
||||
|
||||
|
||||
async def process_socket_message(data: dict, websocket: WebSocket, client_id: str):
|
||||
print(f"Client says: {data['type']}")
|
||||
if data["type"] == "user_message":
|
||||
user_message = Message(**data["data"])
|
||||
session_id = data["data"].get("session_id", None)
|
||||
workflow_id = data["data"].get("workflow_id", None)
|
||||
response = await run_session_workflow(message=user_message, session_id=session_id, workflow_id=workflow_id)
|
||||
response_socket_message = {
|
||||
"type": "agent_response",
|
||||
"data": response,
|
||||
"connection_id": client_id,
|
||||
}
|
||||
await websocket_manager.send_message(response_socket_message, websocket)
|
||||
@api.get("/health")
|
||||
async def health_check():
|
||||
"""API health check endpoint"""
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Service is healthy",
|
||||
}
|
||||
|
||||
# Mount static file directories
|
||||
app.mount("/api", api)
|
||||
app.mount(
|
||||
"/files",
|
||||
StaticFiles(directory=initializer.static_root, html=True),
|
||||
name="files",
|
||||
)
|
||||
app.mount("/", StaticFiles(directory=initializer.ui_root, html=True), name="ui")
|
||||
|
||||
# Error handlers
|
||||
|
||||
|
||||
@api.websocket("/ws/{client_id}")
|
||||
async def websocket_endpoint(websocket: WebSocket, client_id: str):
|
||||
await websocket_manager.connect(websocket, client_id)
|
||||
try:
|
||||
while True:
|
||||
data = await websocket.receive_json()
|
||||
await process_socket_message(data, websocket, client_id)
|
||||
except WebSocketDisconnect:
|
||||
print(f"Client #{client_id} is disconnected")
|
||||
await websocket_manager.disconnect(websocket)
|
||||
@app.exception_handler(500)
|
||||
async def internal_error_handler(request, exc):
|
||||
logger.error(f"Internal error: {str(exc)}")
|
||||
return {
|
||||
"status": False,
|
||||
"message": "Internal server error",
|
||||
"detail": str(exc) if settings.API_DOCS else "Internal server error"
|
||||
}
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""
|
||||
Factory function to create and configure the FastAPI application.
|
||||
Useful for testing and different deployment scenarios.
|
||||
"""
|
||||
return app
|
||||
|
||||
18
python/packages/autogen-studio/autogenstudio/web/config.py
Normal file
18
python/packages/autogen-studio/autogenstudio/web/config.py
Normal file
@@ -0,0 +1,18 @@
|
||||
# api/config.py
|
||||
from pydantic_settings import BaseSettings
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
DATABASE_URI: str = "sqlite:///./autogen.db"
|
||||
API_DOCS: bool = False
|
||||
CLEANUP_INTERVAL: int = 300 # 5 minutes
|
||||
SESSION_TIMEOUT: int = 3600 # 1 hour
|
||||
CONFIG_DIR: str = "configs" # Default config directory relative to app_root
|
||||
DEFAULT_USER_ID: str = "guestuser@gmail.com"
|
||||
UPGRADE_DATABASE: bool = False
|
||||
|
||||
class Config:
|
||||
env_prefix = "AUTOGENSTUDIO_"
|
||||
|
||||
|
||||
settings = Settings()
|
||||
201
python/packages/autogen-studio/autogenstudio/web/deps.py
Normal file
201
python/packages/autogen-studio/autogenstudio/web/deps.py
Normal file
@@ -0,0 +1,201 @@
|
||||
# api/deps.py
|
||||
from typing import Optional
|
||||
from fastapi import Depends, HTTPException, status
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
|
||||
from ..database import DatabaseManager
|
||||
from .managers.connection import WebSocketManager
|
||||
from ..teammanager import TeamManager
|
||||
from .config import settings
|
||||
from ..database import ConfigurationManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global manager instances
|
||||
_db_manager: Optional[DatabaseManager] = None
|
||||
_websocket_manager: Optional[WebSocketManager] = None
|
||||
_team_manager: Optional[TeamManager] = None
|
||||
|
||||
# Context manager for database sessions
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_db_context():
|
||||
"""Provide a transactional scope around a series of operations."""
|
||||
if not _db_manager:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Database manager not initialized"
|
||||
)
|
||||
try:
|
||||
yield _db_manager
|
||||
except Exception as e:
|
||||
logger.error(f"Database operation failed: {str(e)}")
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Database operation failed"
|
||||
)
|
||||
|
||||
# Dependency providers
|
||||
|
||||
|
||||
async def get_db() -> DatabaseManager:
|
||||
"""Dependency provider for database manager"""
|
||||
if not _db_manager:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Database manager not initialized"
|
||||
)
|
||||
return _db_manager
|
||||
|
||||
|
||||
async def get_websocket_manager() -> WebSocketManager:
|
||||
"""Dependency provider for connection manager"""
|
||||
if not _websocket_manager:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Connection manager not initialized"
|
||||
)
|
||||
return _websocket_manager
|
||||
|
||||
|
||||
async def get_team_manager() -> TeamManager:
|
||||
"""Dependency provider for team manager"""
|
||||
if not _team_manager:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail="Team manager not initialized"
|
||||
)
|
||||
return _team_manager
|
||||
|
||||
# Authentication dependency
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
# Add your authentication logic here
|
||||
# For example: token: str = Depends(oauth2_scheme)
|
||||
) -> str:
|
||||
"""
|
||||
Dependency for getting the current authenticated user.
|
||||
Replace with your actual authentication logic.
|
||||
"""
|
||||
# Implement your user authentication here
|
||||
return "user_id" # Replace with actual user identification
|
||||
|
||||
# Manager initialization and cleanup
|
||||
|
||||
|
||||
async def init_managers(database_uri: str, config_dir: str) -> None:
|
||||
"""Initialize all manager instances"""
|
||||
global _db_manager, _websocket_manager, _team_manager
|
||||
|
||||
logger.info("Initializing managers...")
|
||||
|
||||
try:
|
||||
# Initialize database manager
|
||||
_db_manager = DatabaseManager(
|
||||
engine_uri=database_uri, auto_upgrade=settings.UPGRADE_DATABASE)
|
||||
_db_manager.create_db_and_tables()
|
||||
|
||||
# init default team config
|
||||
|
||||
_team_config_manager = ConfigurationManager(db_manager=_db_manager)
|
||||
import_result = await _team_config_manager.import_directory(
|
||||
config_dir, settings.DEFAULT_USER_ID, check_exists=True)
|
||||
|
||||
# Initialize connection manager
|
||||
_websocket_manager = WebSocketManager(
|
||||
db_manager=_db_manager
|
||||
)
|
||||
logger.info("Connection manager initialized")
|
||||
|
||||
# Initialize team manager
|
||||
_team_manager = TeamManager()
|
||||
logger.info("Team manager initialized")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize managers: {str(e)}")
|
||||
await cleanup_managers() # Cleanup any partially initialized managers
|
||||
raise
|
||||
|
||||
|
||||
async def cleanup_managers() -> None:
|
||||
"""Cleanup and shutdown all manager instances"""
|
||||
global _db_manager, _websocket_manager, _team_manager
|
||||
|
||||
logger.info("Cleaning up managers...")
|
||||
|
||||
# Cleanup connection manager first to ensure all active connections are closed
|
||||
if _websocket_manager:
|
||||
try:
|
||||
await _websocket_manager.cleanup()
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up connection manager: {str(e)}")
|
||||
finally:
|
||||
_websocket_manager = None
|
||||
|
||||
# TeamManager doesn't need explicit cleanup since WebSocketManager handles it
|
||||
_team_manager = None
|
||||
|
||||
# Cleanup database manager last
|
||||
if _db_manager:
|
||||
try:
|
||||
await _db_manager.close()
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up database manager: {str(e)}")
|
||||
finally:
|
||||
_db_manager = None
|
||||
|
||||
logger.info("All managers cleaned up")
|
||||
|
||||
# Utility functions for dependency management
|
||||
|
||||
|
||||
def get_manager_status() -> dict:
|
||||
"""Get the initialization status of all managers"""
|
||||
return {
|
||||
"database_manager": _db_manager is not None,
|
||||
"websocket_manager": _websocket_manager is not None,
|
||||
"team_manager": _team_manager is not None
|
||||
}
|
||||
|
||||
# Combined dependencies
|
||||
|
||||
|
||||
async def get_managers():
|
||||
"""Get all managers in one dependency"""
|
||||
return {
|
||||
"db": await get_db(),
|
||||
"connection": await get_websocket_manager(),
|
||||
"team": await get_team_manager()
|
||||
}
|
||||
|
||||
# Error handling for manager operations
|
||||
|
||||
|
||||
class ManagerOperationError(Exception):
|
||||
"""Custom exception for manager operation errors"""
|
||||
|
||||
def __init__(self, manager_name: str, operation: str, detail: str):
|
||||
self.manager_name = manager_name
|
||||
self.operation = operation
|
||||
self.detail = detail
|
||||
super().__init__(f"{manager_name} failed during {operation}: {detail}")
|
||||
|
||||
# Dependency for requiring specific managers
|
||||
|
||||
|
||||
def require_managers(*manager_names: str):
|
||||
"""Decorator to require specific managers for a route"""
|
||||
async def dependency():
|
||||
status = get_manager_status()
|
||||
missing = [name for name in manager_names if not status.get(
|
||||
f"{name}_manager")]
|
||||
if missing:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail=f"Required managers not available: {', '.join(missing)}"
|
||||
)
|
||||
return True
|
||||
return Depends(dependency)
|
||||
@@ -0,0 +1,110 @@
|
||||
# api/initialization.py
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
from pydantic import BaseModel
|
||||
from loguru import logger
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from .config import Settings
|
||||
|
||||
|
||||
class _AppPaths(BaseModel):
|
||||
"""Internal model representing all application paths"""
|
||||
app_root: Path
|
||||
static_root: Path
|
||||
user_files: Path
|
||||
ui_root: Path
|
||||
config_dir: Path
|
||||
database_uri: str
|
||||
|
||||
|
||||
class AppInitializer:
|
||||
"""Handles application initialization including paths and environment setup"""
|
||||
|
||||
def __init__(self, settings: Settings, app_path: str):
|
||||
"""
|
||||
Initialize the application structure.
|
||||
|
||||
Args:
|
||||
settings: Application settings
|
||||
app_path: Path to the application code directory
|
||||
"""
|
||||
self.settings = settings
|
||||
self._app_path = Path(app_path)
|
||||
self._paths = self._init_paths()
|
||||
self._create_directories()
|
||||
self._load_environment()
|
||||
logger.info(f"Initialized application data folder: {self.app_root}")
|
||||
|
||||
def _get_app_root(self) -> Path:
|
||||
"""Determine application root directory"""
|
||||
if app_dir := os.getenv("AUTOGENSTUDIO_APPDIR"):
|
||||
return Path(app_dir)
|
||||
return Path.home() / ".autogenstudio"
|
||||
|
||||
def _get_database_uri(self, app_root: Path) -> str:
|
||||
"""Generate database URI based on settings or environment"""
|
||||
if db_uri := os.getenv("AUTOGENSTUDIO_DATABASE_URI"):
|
||||
return db_uri
|
||||
return self.settings.DATABASE_URI.replace(
|
||||
"./", str(app_root) + "/"
|
||||
)
|
||||
|
||||
def _init_paths(self) -> _AppPaths:
|
||||
"""Initialize and return AppPaths instance"""
|
||||
app_root = self._get_app_root()
|
||||
return _AppPaths(
|
||||
app_root=app_root,
|
||||
static_root=app_root / "files",
|
||||
user_files=app_root / "files" / "user",
|
||||
ui_root=self._app_path / "ui",
|
||||
config_dir=app_root / self.settings.CONFIG_DIR,
|
||||
database_uri=self._get_database_uri(app_root)
|
||||
)
|
||||
|
||||
def _create_directories(self) -> None:
|
||||
"""Create all required directories"""
|
||||
self.app_root.mkdir(parents=True, exist_ok=True)
|
||||
dirs = [self.static_root, self.user_files,
|
||||
self.ui_root, self.config_dir]
|
||||
for path in dirs:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _load_environment(self) -> None:
|
||||
"""Load environment variables from .env file if it exists"""
|
||||
env_file = self.app_root / ".env"
|
||||
if env_file.exists():
|
||||
logger.info(f"Loading environment variables from {env_file}")
|
||||
load_dotenv(str(env_file))
|
||||
|
||||
# Properties for accessing paths
|
||||
@property
|
||||
def app_root(self) -> Path:
|
||||
"""Root directory for the application"""
|
||||
return self._paths.app_root
|
||||
|
||||
@property
|
||||
def static_root(self) -> Path:
|
||||
"""Directory for static files"""
|
||||
return self._paths.static_root
|
||||
|
||||
@property
|
||||
def user_files(self) -> Path:
|
||||
"""Directory for user files"""
|
||||
return self._paths.user_files
|
||||
|
||||
@property
|
||||
def ui_root(self) -> Path:
|
||||
"""Directory for UI files"""
|
||||
return self._paths.ui_root
|
||||
|
||||
@property
|
||||
def config_dir(self) -> Path:
|
||||
"""Directory for configuration files"""
|
||||
return self._paths.config_dir
|
||||
|
||||
@property
|
||||
def database_uri(self) -> str:
|
||||
"""Database connection URI"""
|
||||
return self._paths.database_uri
|
||||
@@ -0,0 +1,247 @@
|
||||
# managers/connection.py
|
||||
from fastapi import WebSocket, WebSocketDisconnect
|
||||
from typing import Dict, Optional, Any
|
||||
from uuid import UUID
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from ...datamodel import Run, RunStatus, TeamResult
|
||||
from ...database import DatabaseManager
|
||||
from autogen_agentchat.messages import InnerMessage, ChatMessage
|
||||
from autogen_core.base import CancellationToken
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WebSocketManager:
|
||||
"""Manages WebSocket connections and message streaming for team task execution"""
|
||||
|
||||
def __init__(self, db_manager: DatabaseManager):
|
||||
self.db_manager = db_manager
|
||||
self._connections: Dict[UUID, WebSocket] = {}
|
||||
self._cancellation_tokens: Dict[UUID, CancellationToken] = {}
|
||||
|
||||
async def connect(self, websocket: WebSocket, run_id: UUID) -> bool:
|
||||
"""Initialize WebSocket connection for a run
|
||||
|
||||
Args:
|
||||
websocket: The WebSocket connection to initialize
|
||||
run_id: UUID of the run to associate with this connection
|
||||
|
||||
Returns:
|
||||
bool: True if connection was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
await websocket.accept()
|
||||
self._connections[run_id] = websocket
|
||||
|
||||
run = await self._get_run(run_id)
|
||||
if run:
|
||||
run.status = RunStatus.ACTIVE
|
||||
self.db_manager.upsert(run)
|
||||
|
||||
await self._send_message(run_id, {
|
||||
"type": "system",
|
||||
"status": "connected",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
})
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Connection error for run {run_id}: {e}")
|
||||
return False
|
||||
|
||||
async def start_stream(
|
||||
self,
|
||||
run_id: UUID,
|
||||
team_manager: Any,
|
||||
task: str,
|
||||
team_config: dict
|
||||
) -> None:
|
||||
"""Start streaming task execution
|
||||
|
||||
Args:
|
||||
run_id: UUID of the run
|
||||
team_manager: Instance of the team manager
|
||||
task: Task string to execute
|
||||
team_config: Team configuration dictionary
|
||||
"""
|
||||
if run_id not in self._connections:
|
||||
raise ValueError(f"No active connection for run {run_id}")
|
||||
|
||||
cancellation_token = CancellationToken()
|
||||
self._cancellation_tokens[run_id] = cancellation_token
|
||||
|
||||
try:
|
||||
async for message in team_manager.run_stream(
|
||||
task=task,
|
||||
team_config=team_config,
|
||||
cancellation_token=cancellation_token
|
||||
):
|
||||
|
||||
if cancellation_token.is_cancelled():
|
||||
logger.info(f"Stream cancelled for run {run_id}")
|
||||
break
|
||||
|
||||
formatted_message = self._format_message(message)
|
||||
if formatted_message:
|
||||
await self._send_message(run_id, formatted_message)
|
||||
|
||||
# Only send completion if not cancelled
|
||||
if not cancellation_token.is_cancelled():
|
||||
# await self._send_message(run_id, {
|
||||
# "type": "completion",
|
||||
# "status": "complete",
|
||||
# "timestamp": datetime.now(timezone.utc).isoformat()
|
||||
# })
|
||||
await self._update_run_status(run_id, RunStatus.COMPLETE)
|
||||
else:
|
||||
await self._send_message(run_id, {
|
||||
"type": "completion",
|
||||
"status": "cancelled",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
})
|
||||
await self._update_run_status(run_id, RunStatus.STOPPED)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Stream error for run {run_id}: {e}")
|
||||
await self._handle_stream_error(run_id, e)
|
||||
|
||||
finally:
|
||||
self._cancellation_tokens.pop(run_id, None)
|
||||
|
||||
async def stop_run(self, run_id: UUID) -> None:
|
||||
"""Stop a running task"""
|
||||
if run_id in self._cancellation_tokens:
|
||||
logger.info(f"Stopping run {run_id}")
|
||||
self._cancellation_tokens[run_id].cancel()
|
||||
|
||||
# Send final message if connection still exists
|
||||
if run_id in self._connections:
|
||||
try:
|
||||
await self._send_message(run_id, {
|
||||
"type": "completion",
|
||||
"status": "cancelled",
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def disconnect(self, run_id: UUID) -> None:
|
||||
"""Clean up connection and associated resources"""
|
||||
logger.info(f"Disconnecting run {run_id}")
|
||||
|
||||
# First cancel any running tasks
|
||||
await self.stop_run(run_id)
|
||||
|
||||
# Then clean up resources without trying to close the socket again
|
||||
if run_id in self._connections:
|
||||
self._connections.pop(run_id, None)
|
||||
self._cancellation_tokens.pop(run_id, None)
|
||||
|
||||
async def _send_message(self, run_id: UUID, message: dict) -> None:
|
||||
"""Send a message through the WebSocket
|
||||
|
||||
Args:
|
||||
run_id: UUID of the run
|
||||
message: Message dictionary to send
|
||||
"""
|
||||
try:
|
||||
if run_id in self._connections:
|
||||
await self._connections[run_id].send_json(message)
|
||||
except WebSocketDisconnect:
|
||||
logger.warning(
|
||||
f"WebSocket disconnected while sending message for run {run_id}")
|
||||
await self.disconnect(run_id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error sending message for run {run_id}: {e}")
|
||||
await self._handle_stream_error(run_id, e)
|
||||
|
||||
async def _handle_stream_error(self, run_id: UUID, error: Exception) -> None:
|
||||
"""Handle stream errors consistently
|
||||
|
||||
Args:
|
||||
run_id: UUID of the run
|
||||
error: Exception that occurred
|
||||
"""
|
||||
try:
|
||||
await self._send_message(run_id, {
|
||||
"type": "completion",
|
||||
"status": "error",
|
||||
"error": str(error),
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
})
|
||||
except Exception as send_error:
|
||||
logger.error(
|
||||
f"Failed to send error message for run {run_id}: {send_error}")
|
||||
|
||||
await self._update_run_status(run_id, RunStatus.ERROR, str(error))
|
||||
|
||||
def _format_message(self, message: Any) -> Optional[dict]:
|
||||
"""Format message for WebSocket transmission
|
||||
|
||||
Args:
|
||||
message: Message to format
|
||||
|
||||
Returns:
|
||||
Optional[dict]: Formatted message or None if formatting fails
|
||||
"""
|
||||
try:
|
||||
if isinstance(message, (InnerMessage, ChatMessage)):
|
||||
return {
|
||||
"type": "message",
|
||||
"data": message.model_dump()
|
||||
}
|
||||
elif isinstance(message, TeamResult):
|
||||
return {
|
||||
"type": "result",
|
||||
"data": message.model_dump(),
|
||||
"status": "complete",
|
||||
}
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Message formatting error: {e}")
|
||||
return None
|
||||
|
||||
async def _get_run(self, run_id: UUID) -> Optional[Run]:
|
||||
"""Get run from database
|
||||
|
||||
Args:
|
||||
run_id: UUID of the run to retrieve
|
||||
|
||||
Returns:
|
||||
Optional[Run]: Run object if found, None otherwise
|
||||
"""
|
||||
response = self.db_manager.get(
|
||||
Run, filters={"id": run_id}, return_json=False)
|
||||
return response.data[0] if response.status and response.data else None
|
||||
|
||||
async def _update_run_status(
|
||||
self,
|
||||
run_id: UUID,
|
||||
status: RunStatus,
|
||||
error: Optional[str] = None
|
||||
) -> None:
|
||||
"""Update run status in database
|
||||
|
||||
Args:
|
||||
run_id: UUID of the run to update
|
||||
status: New status to set
|
||||
error: Optional error message
|
||||
"""
|
||||
run = await self._get_run(run_id)
|
||||
if run:
|
||||
run.status = status
|
||||
run.error_message = error
|
||||
self.db_manager.upsert(run)
|
||||
|
||||
@property
|
||||
def active_connections(self) -> set[UUID]:
|
||||
"""Get set of active run IDs"""
|
||||
return set(self._connections.keys())
|
||||
|
||||
@property
|
||||
def active_runs(self) -> set[UUID]:
|
||||
"""Get set of runs with active cancellation tokens"""
|
||||
return set(self._cancellation_tokens.keys())
|
||||
@@ -0,0 +1,181 @@
|
||||
# api/routes/agents.py
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from ..deps import get_db
|
||||
from ...datamodel import Agent, Model, Tool
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def list_agents(
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all agents for a user"""
|
||||
response = db.get(Agent, filters={"user_id": user_id})
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{agent_id}")
|
||||
async def get_agent(
|
||||
agent_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get a specific agent"""
|
||||
response = db.get(
|
||||
Agent,
|
||||
filters={"id": agent_id, "user_id": user_id}
|
||||
)
|
||||
if not response.status or not response.data:
|
||||
raise HTTPException(status_code=404, detail="Agent not found")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data[0]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_agent(
|
||||
agent: Agent,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Create a new agent"""
|
||||
response = db.upsert(agent)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{agent_id}")
|
||||
async def delete_agent(
|
||||
agent_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Delete an agent"""
|
||||
response = db.delete(
|
||||
filters={"id": agent_id, "user_id": user_id},
|
||||
model_class=Agent
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Agent deleted successfully"
|
||||
}
|
||||
|
||||
# Agent-Model link endpoints
|
||||
|
||||
|
||||
@router.post("/{agent_id}/models/{model_id}")
|
||||
async def link_agent_model(
|
||||
agent_id: int,
|
||||
model_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Link a model to an agent"""
|
||||
response = db.link(
|
||||
link_type="agent_model",
|
||||
primary_id=agent_id,
|
||||
secondary_id=model_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Model linked to agent successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{agent_id}/models/{model_id}")
|
||||
async def unlink_agent_model(
|
||||
agent_id: int,
|
||||
model_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Unlink a model from an agent"""
|
||||
response = db.unlink(
|
||||
link_type="agent_model",
|
||||
primary_id=agent_id,
|
||||
secondary_id=model_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Model unlinked from agent successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{agent_id}/models")
|
||||
async def get_agent_models(
|
||||
agent_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get all models linked to an agent"""
|
||||
response = db.get_linked_entities(
|
||||
link_type="agent_model",
|
||||
primary_id=agent_id,
|
||||
return_json=True
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
# Agent-Tool link endpoints
|
||||
|
||||
|
||||
@router.post("/{agent_id}/tools/{tool_id}")
|
||||
async def link_agent_tool(
|
||||
agent_id: int,
|
||||
tool_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Link a tool to an agent"""
|
||||
response = db.link(
|
||||
link_type="agent_tool",
|
||||
primary_id=agent_id,
|
||||
secondary_id=tool_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Tool linked to agent successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{agent_id}/tools/{tool_id}")
|
||||
async def unlink_agent_tool(
|
||||
agent_id: int,
|
||||
tool_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Unlink a tool from an agent"""
|
||||
response = db.unlink(
|
||||
link_type="agent_tool",
|
||||
primary_id=agent_id,
|
||||
secondary_id=tool_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Tool unlinked from agent successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{agent_id}/tools")
|
||||
async def get_agent_tools(
|
||||
agent_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get all tools linked to an agent"""
|
||||
response = db.get_linked_entities(
|
||||
link_type="agent_tool",
|
||||
primary_id=agent_id,
|
||||
return_json=True
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
# api/routes/models.py
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from openai import OpenAIError
|
||||
from ..deps import get_db
|
||||
from ...datamodel import Model
|
||||
from ...utils import test_model
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def list_models(
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all models for a user"""
|
||||
response = db.get(Model, filters={"user_id": user_id})
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{model_id}")
|
||||
async def get_model(
|
||||
model_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get a specific model"""
|
||||
response = db.get(
|
||||
Model,
|
||||
filters={"id": model_id, "user_id": user_id}
|
||||
)
|
||||
if not response.status or not response.data:
|
||||
raise HTTPException(status_code=404, detail="Model not found")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data[0]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_model(
|
||||
model: Model,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Create a new model"""
|
||||
response = db.upsert(model)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{model_id}")
|
||||
async def delete_model(
|
||||
model_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Delete a model"""
|
||||
response = db.delete(
|
||||
filters={"id": model_id, "user_id": user_id},
|
||||
model_class=Model
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Model deleted successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.post("/test")
|
||||
async def test_model_endpoint(model: Model) -> Dict:
|
||||
"""Test a model configuration"""
|
||||
try:
|
||||
response = test_model(model)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Model tested successfully",
|
||||
"data": response
|
||||
}
|
||||
except OpenAIError as e:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"OpenAI API error: {str(e)}"
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error testing model: {str(e)}"
|
||||
)
|
||||
@@ -0,0 +1,76 @@
|
||||
# /api/runs routes
|
||||
from fastapi import APIRouter, Body, Depends, HTTPException
|
||||
from uuid import UUID
|
||||
from typing import Dict
|
||||
|
||||
from pydantic import BaseModel
|
||||
from ..deps import get_db, get_websocket_manager, get_team_manager
|
||||
from ...datamodel import Run, Session, Message, Team, RunStatus, MessageConfig
|
||||
|
||||
from ...teammanager import TeamManager
|
||||
from autogen_core.base import CancellationToken
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class CreateRunRequest(BaseModel):
|
||||
session_id: int
|
||||
user_id: str
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_run(
|
||||
request: CreateRunRequest,
|
||||
db=Depends(get_db),
|
||||
) -> Dict:
|
||||
"""Create a new run"""
|
||||
session_response = db.get(
|
||||
Session,
|
||||
filters={"id": request.session_id, "user_id": request.user_id},
|
||||
return_json=False
|
||||
)
|
||||
if not session_response.status or not session_response.data:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
|
||||
try:
|
||||
|
||||
run = db.upsert(Run(session_id=request.session_id), return_json=False)
|
||||
return {
|
||||
"status": run.status,
|
||||
"data": {"run_id": str(run.data.id)}
|
||||
}
|
||||
|
||||
# }
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post("/{run_id}/start")
|
||||
async def start_run(
|
||||
run_id: UUID,
|
||||
message: Message = Body(...),
|
||||
ws_manager=Depends(get_websocket_manager),
|
||||
team_manager=Depends(get_team_manager),
|
||||
db=Depends(get_db),
|
||||
) -> Dict:
|
||||
"""Start streaming task execution"""
|
||||
|
||||
if isinstance(message.config, dict):
|
||||
message.config = MessageConfig(**message.config)
|
||||
|
||||
session = db.get(Session, filters={
|
||||
"id": message.session_id}, return_json=False)
|
||||
|
||||
team = db.get(
|
||||
Team, filters={"id": session.data[0].team_id}, return_json=False)
|
||||
|
||||
try:
|
||||
await ws_manager.start_stream(run_id, team_manager, message.config.content, team.data[0].config)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Stream started successfully",
|
||||
"data": {"run_id": str(run_id)}
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,114 @@
|
||||
# api/routes/sessions.py
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from ..deps import get_db
|
||||
from ...datamodel import Session, Message
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def list_sessions(
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all sessions for a user"""
|
||||
response = db.get(Session, filters={"user_id": user_id})
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{session_id}")
|
||||
async def get_session(
|
||||
session_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get a specific session"""
|
||||
response = db.get(
|
||||
Session,
|
||||
filters={"id": session_id, "user_id": user_id}
|
||||
)
|
||||
if not response.status or not response.data:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data[0]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_session(
|
||||
session: Session,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Create a new session"""
|
||||
response = db.upsert(session)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.put("/{session_id}")
|
||||
async def update_session(
|
||||
session_id: int,
|
||||
user_id: str,
|
||||
session: Session,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Update an existing session"""
|
||||
# First verify the session belongs to user
|
||||
existing = db.get(
|
||||
Session,
|
||||
filters={"id": session_id, "user_id": user_id}
|
||||
)
|
||||
if not existing.status or not existing.data:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
|
||||
# Update the session
|
||||
response = db.upsert(session)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data,
|
||||
"message": "Session updated successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{session_id}")
|
||||
async def delete_session(
|
||||
session_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Delete a session"""
|
||||
response = db.delete(
|
||||
filters={"id": session_id, "user_id": user_id},
|
||||
model_class=Session
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Session deleted successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{session_id}/messages")
|
||||
async def list_messages(
|
||||
session_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all messages for a session"""
|
||||
filters = {"session_id": session_id, "user_id": user_id}
|
||||
response = db.get(Message, filters=filters, order="asc")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
146
python/packages/autogen-studio/autogenstudio/web/routes/teams.py
Normal file
146
python/packages/autogen-studio/autogenstudio/web/routes/teams.py
Normal file
@@ -0,0 +1,146 @@
|
||||
# api/routes/teams.py
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from ..deps import get_db
|
||||
from ...datamodel import Team
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def list_teams(
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all teams for a user"""
|
||||
response = db.get(Team, filters={"user_id": user_id})
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{team_id}")
|
||||
async def get_team(
|
||||
team_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get a specific team"""
|
||||
response = db.get(
|
||||
Team,
|
||||
filters={"id": team_id, "user_id": user_id}
|
||||
)
|
||||
if not response.status or not response.data:
|
||||
raise HTTPException(status_code=404, detail="Team not found")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data[0]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_team(
|
||||
team: Team,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Create a new team"""
|
||||
response = db.upsert(team)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{team_id}")
|
||||
async def delete_team(
|
||||
team_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Delete a team"""
|
||||
response = db.delete(
|
||||
filters={"id": team_id, "user_id": user_id},
|
||||
model_class=Team
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Team deleted successfully"
|
||||
}
|
||||
|
||||
# Team-Agent link endpoints
|
||||
|
||||
|
||||
@router.post("/{team_id}/agents/{agent_id}")
|
||||
async def link_team_agent(
|
||||
team_id: int,
|
||||
agent_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Link an agent to a team"""
|
||||
response = db.link(
|
||||
link_type="team_agent",
|
||||
primary_id=team_id,
|
||||
secondary_id=agent_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Agent linked to team successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.post("/{team_id}/agents/{agent_id}/{sequence_id}")
|
||||
async def link_team_agent_sequence(
|
||||
team_id: int,
|
||||
agent_id: int,
|
||||
sequence_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Link an agent to a team with sequence"""
|
||||
response = db.link(
|
||||
link_type="team_agent",
|
||||
primary_id=team_id,
|
||||
secondary_id=agent_id,
|
||||
sequence_id=sequence_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Agent linked to team with sequence successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{team_id}/agents/{agent_id}")
|
||||
async def unlink_team_agent(
|
||||
team_id: int,
|
||||
agent_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Unlink an agent from a team"""
|
||||
response = db.unlink(
|
||||
link_type="team_agent",
|
||||
primary_id=team_id,
|
||||
secondary_id=agent_id
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Agent unlinked from team successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{team_id}/agents")
|
||||
async def get_team_agents(
|
||||
team_id: int,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get all agents linked to a team"""
|
||||
response = db.get_linked_entities(
|
||||
link_type="team_agent",
|
||||
primary_id=team_id,
|
||||
return_json=True
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
103
python/packages/autogen-studio/autogenstudio/web/routes/tools.py
Normal file
103
python/packages/autogen-studio/autogenstudio/web/routes/tools.py
Normal file
@@ -0,0 +1,103 @@
|
||||
# api/routes/tools.py
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from ..deps import get_db
|
||||
from ...datamodel import Tool
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/")
|
||||
async def list_tools(
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""List all tools for a user"""
|
||||
response = db.get(Tool, filters={"user_id": user_id})
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{tool_id}")
|
||||
async def get_tool(
|
||||
tool_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Get a specific tool"""
|
||||
response = db.get(
|
||||
Tool,
|
||||
filters={"id": tool_id, "user_id": user_id}
|
||||
)
|
||||
if not response.status or not response.data:
|
||||
raise HTTPException(status_code=404, detail="Tool not found")
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data[0]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/")
|
||||
async def create_tool(
|
||||
tool: Tool,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Create a new tool"""
|
||||
response = db.upsert(tool)
|
||||
if not response.status:
|
||||
raise HTTPException(status_code=400, detail=response.message)
|
||||
return {
|
||||
"status": True,
|
||||
"data": response.data
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{tool_id}")
|
||||
async def delete_tool(
|
||||
tool_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Delete a tool"""
|
||||
response = db.delete(
|
||||
filters={"id": tool_id, "user_id": user_id},
|
||||
model_class=Tool
|
||||
)
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Tool deleted successfully"
|
||||
}
|
||||
|
||||
|
||||
@router.post("/{tool_id}/test")
|
||||
async def test_tool(
|
||||
tool_id: int,
|
||||
user_id: str,
|
||||
db=Depends(get_db)
|
||||
) -> Dict:
|
||||
"""Test a tool configuration"""
|
||||
# Get tool
|
||||
tool_response = db.get(
|
||||
Tool,
|
||||
filters={"id": tool_id, "user_id": user_id}
|
||||
)
|
||||
if not tool_response.status or not tool_response.data:
|
||||
raise HTTPException(status_code=404, detail="Tool not found")
|
||||
|
||||
tool = tool_response.data[0]
|
||||
|
||||
try:
|
||||
# Implement tool testing logic here
|
||||
# This would depend on the tool type and configuration
|
||||
return {
|
||||
"status": True,
|
||||
"message": "Tool tested successfully",
|
||||
"data": {"tool_id": tool_id}
|
||||
}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error testing tool: {str(e)}"
|
||||
)
|
||||
@@ -0,0 +1,74 @@
|
||||
# api/ws.py
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException
|
||||
from typing import Dict
|
||||
from uuid import UUID
|
||||
import logging
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
from ..deps import get_websocket_manager, get_db, get_team_manager
|
||||
from ...datamodel import Run, RunStatus
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@router.websocket("/runs/{run_id}")
|
||||
async def run_websocket(
|
||||
websocket: WebSocket,
|
||||
run_id: UUID,
|
||||
ws_manager=Depends(get_websocket_manager),
|
||||
db=Depends(get_db),
|
||||
team_manager=Depends(get_team_manager)
|
||||
):
|
||||
"""WebSocket endpoint for run communication"""
|
||||
# Verify run exists and is in valid state
|
||||
run_response = db.get(Run, filters={"id": run_id}, return_json=False)
|
||||
if not run_response.status or not run_response.data:
|
||||
await websocket.close(code=4004, reason="Run not found")
|
||||
return
|
||||
|
||||
run = run_response.data[0]
|
||||
if run.status not in [RunStatus.CREATED, RunStatus.ACTIVE]:
|
||||
await websocket.close(code=4003, reason="Run not in valid state")
|
||||
return
|
||||
|
||||
# Connect websocket
|
||||
connected = await ws_manager.connect(websocket, run_id)
|
||||
if not connected:
|
||||
await websocket.close(code=4002, reason="Failed to establish connection")
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info(f"WebSocket connection established for run {run_id}")
|
||||
|
||||
while True:
|
||||
try:
|
||||
raw_message = await websocket.receive_text()
|
||||
message = json.loads(raw_message)
|
||||
|
||||
if message.get("type") == "stop":
|
||||
logger.info(f"Received stop request for run {run_id}")
|
||||
await ws_manager.stop_run(run_id)
|
||||
break
|
||||
|
||||
elif message.get("type") == "ping":
|
||||
await websocket.send_json({
|
||||
"type": "pong",
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Invalid JSON received: {raw_message}")
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
"error": "Invalid message format",
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info(f"WebSocket disconnected for run {run_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket error: {str(e)}")
|
||||
finally:
|
||||
await ws_manager.disconnect(run_id)
|
||||
Reference in New Issue
Block a user