agno 2.0.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/os/router.py

@@ -1,5 +1,5 @@
 import json
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Dict, List, Optional, Union, cast
 from uuid import uuid4
 
 from fastapi import (
@@ -8,22 +8,26 @@ from fastapi import (
     File,
     Form,
     HTTPException,
+    Request,
     UploadFile,
     WebSocket,
 )
 from fastapi.responses import JSONResponse, StreamingResponse
+from packaging import version
 from pydantic import BaseModel
 
 from agno.agent.agent import Agent
+from agno.db.base import AsyncBaseDb
+from agno.db.migrations.manager import MigrationManager
+from agno.exceptions import InputCheckError, OutputCheckError
 from agno.media import Audio, Image, Video
 from agno.media import File as FileMedia
-from agno.os.auth import get_authentication_dependency
+from agno.os.auth import get_authentication_dependency, validate_websocket_token
 from agno.os.schema import (
     AgentResponse,
     AgentSummaryResponse,
     BadRequestResponse,
     ConfigResponse,
-    HealthResponse,
     InterfaceResponse,
     InternalServerErrorResponse,
     Model,
@@ -38,6 +42,7 @@ from agno.os.schema import (
 from agno.os.settings import AgnoAPISettings
 from agno.os.utils import (
     get_agent_by_id,
+    get_db,
     get_team_by_id,
     get_workflow_by_id,
     process_audio,
@@ -45,9 +50,10 @@ from agno.os.utils import (
     process_image,
     process_video,
 )
-from agno.run.agent import RunErrorEvent, RunOutput
+from agno.run.agent import RunErrorEvent, RunOutput, RunOutputEvent
 from agno.run.team import RunErrorEvent as TeamRunErrorEvent
-from agno.run.workflow import WorkflowErrorEvent
+from agno.run.team import TeamRunOutputEvent
+from agno.run.workflow import WorkflowErrorEvent, WorkflowRunOutput, WorkflowRunOutputEvent
 from agno.team.team import Team
 from agno.utils.log import log_debug, log_error, log_warning, logger
 from agno.workflow.workflow import Workflow
@@ -56,11 +62,98 @@ if TYPE_CHECKING:
     from agno.os.app import AgentOS
 
 
-def format_sse_event(json_data: str) -> str:
+async def _get_request_kwargs(request: Request, endpoint_func: Callable) -> Dict[str, Any]:
+    """Given a Request and an endpoint function, return a dictionary with all extra form data fields.
+    Args:
+        request: The FastAPI Request object
+        endpoint_func: The function exposing the endpoint that received the request
+
+    Returns:
+        A dictionary of kwargs
+    """
+    import inspect
+
+    form_data = await request.form()
+    sig = inspect.signature(endpoint_func)
+    known_fields = set(sig.parameters.keys())
+    kwargs: Dict[str, Any] = {key: value for key, value in form_data.items() if key not in known_fields}
+
+    # Handle JSON parameters. They are passed as strings and need to be deserialized.
+    if session_state := kwargs.get("session_state"):
+        try:
+            if isinstance(session_state, str):
+                session_state_dict = json.loads(session_state)  # type: ignore
+                kwargs["session_state"] = session_state_dict
+        except json.JSONDecodeError:
+            kwargs.pop("session_state")
+            log_warning(f"Invalid session_state parameter couldn't be loaded: {session_state}")
+
+    if dependencies := kwargs.get("dependencies"):
+        try:
+            if isinstance(dependencies, str):
+                dependencies_dict = json.loads(dependencies)  # type: ignore
+                kwargs["dependencies"] = dependencies_dict
+        except json.JSONDecodeError:
+            kwargs.pop("dependencies")
+            log_warning(f"Invalid dependencies parameter couldn't be loaded: {dependencies}")
+
+    if metadata := kwargs.get("metadata"):
+        try:
+            if isinstance(metadata, str):
+                metadata_dict = json.loads(metadata)  # type: ignore
+                kwargs["metadata"] = metadata_dict
+        except json.JSONDecodeError:
+            kwargs.pop("metadata")
+            log_warning(f"Invalid metadata parameter couldn't be loaded: {metadata}")
+
+    if knowledge_filters := kwargs.get("knowledge_filters"):
+        try:
+            if isinstance(knowledge_filters, str):
+                knowledge_filters_dict = json.loads(knowledge_filters)  # type: ignore
+
+                # Try to deserialize FilterExpr objects
+                from agno.filters import from_dict
+
+                # Check if it's a single FilterExpr dict or a list of FilterExpr dicts
+                if isinstance(knowledge_filters_dict, dict) and "op" in knowledge_filters_dict:
+                    # Single FilterExpr - convert to list format
+                    kwargs["knowledge_filters"] = [from_dict(knowledge_filters_dict)]
+                elif isinstance(knowledge_filters_dict, list):
+                    # List of FilterExprs or mixed content
+                    deserialized = []
+                    for item in knowledge_filters_dict:
+                        if isinstance(item, dict) and "op" in item:
+                            deserialized.append(from_dict(item))
+                        else:
+                            # Keep non-FilterExpr items as-is
+                            deserialized.append(item)
+                    kwargs["knowledge_filters"] = deserialized
+                else:
+                    # Regular dict filter
+                    kwargs["knowledge_filters"] = knowledge_filters_dict
+        except json.JSONDecodeError:
+            kwargs.pop("knowledge_filters")
+            log_warning(f"Invalid knowledge_filters parameter couldn't be loaded: {knowledge_filters}")
+        except ValueError as e:
+            # Filter deserialization failed
+            kwargs.pop("knowledge_filters")
+            log_warning(f"Invalid FilterExpr in knowledge_filters: {e}")
+
+    # Parse boolean and null values
+    for key, value in kwargs.items():
+        if isinstance(value, str) and value.lower() in ["true", "false"]:
+            kwargs[key] = value.lower() == "true"
+        elif isinstance(value, str) and value.lower() in ["null", "none"]:
+            kwargs[key] = None
+
+    return kwargs
+
+
+def format_sse_event(event: Union[RunOutputEvent, TeamRunOutputEvent, WorkflowRunOutputEvent]) -> str:
     """Parse JSON data into SSE-compliant format.
 
     Args:
-        json_data: JSON string containing the event data
+        event_dict: Dictionary containing the event data
 
     Returns:
         SSE-formatted response:
@@ -75,20 +168,22 @@ def format_sse_event(json_data: str) -> str:
     """
     try:
         # Parse the JSON to extract the event type
-        data = json.loads(json_data)
-        event_type = data.get("event", "message")
+        event_type = event.event or "message"
+
+        # Serialize to valid JSON with double quotes and no newlines
+        clean_json = event.to_json(separators=(",", ":"), indent=None)
 
-        # Format as SSE: event: <event_type>\ndata: <json_data>\n\n
-        return f"event: {event_type}\ndata: {json_data}\n\n"
-    except (json.JSONDecodeError, KeyError):
-        # Fallback to generic message event if parsing fails
-        return f"event: message\ndata: {json_data}\n\n"
+        return f"event: {event_type}\ndata: {clean_json}\n\n"
+    except json.JSONDecodeError:
+        clean_json = event.to_json(separators=(",", ":"), indent=None)
+        return f"event: message\ndata: {clean_json}\n\n"
 
 
 class WebSocketManager:
     """Manages WebSocket connections for workflow runs"""
 
     active_connections: Dict[str, WebSocket]  # {run_id: websocket}
+    authenticated_connections: Dict[WebSocket, bool]  # {websocket: is_authenticated}
 
     def __init__(
         self,
@@ -96,22 +191,51 @@ class WebSocketManager:
     ):
         # Store active connections: {run_id: websocket}
         self.active_connections = active_connections or {}
+        # Track authentication state for each websocket
+        self.authenticated_connections = {}
 
-    async def connect(self, websocket: WebSocket):
+    async def connect(self, websocket: WebSocket, requires_auth: bool = True):
         """Accept WebSocket connection"""
         await websocket.accept()
         logger.debug("WebSocket connected")
 
-        # Send connection confirmation
+        # If auth is not required, mark as authenticated immediately
+        self.authenticated_connections[websocket] = not requires_auth
+
+        # Send connection confirmation with auth requirement info
         await websocket.send_text(
             json.dumps(
                 {
                     "event": "connected",
-                    "message": "Connected to workflow events",
+                    "message": (
+                        "Connected to workflow events. Please authenticate to continue."
+                        if requires_auth
+                        else "Connected to workflow events. Authentication not required."
+                    ),
+                    "requires_auth": requires_auth,
                 }
             )
         )
 
+    async def authenticate_websocket(self, websocket: WebSocket):
+        """Mark a WebSocket connection as authenticated"""
+        self.authenticated_connections[websocket] = True
+        logger.debug("WebSocket authenticated")
+
+        # Send authentication confirmation
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "authenticated",
+                    "message": "Authentication successful. You can now send commands.",
+                }
+            )
+        )
+
+    def is_authenticated(self, websocket: WebSocket) -> bool:
+        """Check if a WebSocket connection is authenticated"""
+        return self.authenticated_connections.get(websocket, False)
+
     async def register_workflow_websocket(self, run_id: str, websocket: WebSocket):
         """Register a workflow run with its WebSocket connection"""
         self.active_connections[run_id] = websocket
@@ -120,9 +244,26 @@ class WebSocketManager:
     async def disconnect_by_run_id(self, run_id: str):
         """Remove WebSocket connection by run_id"""
         if run_id in self.active_connections:
+            websocket = self.active_connections[run_id]
             del self.active_connections[run_id]
+            # Clean up authentication state
+            if websocket in self.authenticated_connections:
+                del self.authenticated_connections[websocket]
             logger.debug(f"WebSocket disconnected for run_id: {run_id}")
 
+    async def disconnect_websocket(self, websocket: WebSocket):
+        """Remove WebSocket connection and clean up all associated state"""
+        # Remove from authenticated connections
+        if websocket in self.authenticated_connections:
+            del self.authenticated_connections[websocket]
+
+        # Remove from active connections
+        runs_to_remove = [run_id for run_id, ws in self.active_connections.items() if ws == websocket]
+        for run_id in runs_to_remove:
+            del self.active_connections[run_id]
+
+        logger.debug("WebSocket disconnected and cleaned up")
+
     async def get_websocket_for_run(self, run_id: str) -> Optional[WebSocket]:
         """Get WebSocket connection for a workflow run"""
         return self.active_connections.get(run_id)
@@ -143,6 +284,7 @@ async def agent_response_streamer(
     audio: Optional[List[Audio]] = None,
     videos: Optional[List[Video]] = None,
     files: Optional[List[FileMedia]] = None,
+    **kwargs: Any,
 ) -> AsyncGenerator:
     try:
         run_response = agent.arun(
@@ -154,11 +296,19 @@ async def agent_response_streamer(
             videos=videos,
             files=files,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
+            **kwargs,
         )
         async for run_response_chunk in run_response:
-            yield format_sse_event(run_response_chunk.to_json())
-
+            yield format_sse_event(run_response_chunk)  # type: ignore
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = RunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
     except Exception as e:
         import traceback
 
@@ -166,7 +316,7 @@ async def agent_response_streamer(
         error_response = RunErrorEvent(
             content=str(e),
         )
-        yield format_sse_event(error_response.to_json())
+        yield format_sse_event(error_response)
 
 
 async def agent_continue_response_streamer(
@@ -183,10 +333,18 @@ async def agent_continue_response_streamer(
             session_id=session_id,
             user_id=user_id,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
         )
         async for run_response_chunk in continue_response:
-            yield format_sse_event(run_response_chunk.to_json())
+            yield format_sse_event(run_response_chunk)  # type: ignore
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = RunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
 
     except Exception as e:
         import traceback
@@ -194,8 +352,10 @@ async def agent_continue_response_streamer(
         traceback.print_exc(limit=3)
         error_response = RunErrorEvent(
             content=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
-        yield format_sse_event(error_response.to_json())
+        yield format_sse_event(error_response)
         return
 
 
@@ -208,6 +368,7 @@ async def team_response_streamer(
     audio: Optional[List[Audio]] = None,
     videos: Optional[List[Video]] = None,
     files: Optional[List[FileMedia]] = None,
+    **kwargs: Any,
 ) -> AsyncGenerator:
     """Run the given team asynchronously and yield its response"""
     try:
@@ -220,10 +381,19 @@ async def team_response_streamer(
             videos=videos,
             files=files,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
+            **kwargs,
         )
         async for run_response_chunk in run_response:
-            yield format_sse_event(run_response_chunk.to_json())
+            yield format_sse_event(run_response_chunk)  # type: ignore
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = TeamRunErrorEvent(
+            content=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
 
     except Exception as e:
         import traceback
@@ -231,8 +401,10 @@ async def team_response_streamer(
         traceback.print_exc()
         error_response = TeamRunErrorEvent(
             content=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
-        yield format_sse_event(error_response.to_json())
+        yield format_sse_event(error_response)
         return
 
 
@@ -263,19 +435,42 @@ async def handle_workflow_via_websocket(websocket: WebSocket, message: dict, os:
                 session_id = str(uuid4())
 
         # Execute workflow in background with streaming
-        await workflow.arun(
+        workflow_result = await workflow.arun(  # type: ignore
             input=user_message,
             session_id=session_id,
             user_id=user_id,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
             background=True,
             websocket=websocket,
         )
 
+        workflow_run_output = cast(WorkflowRunOutput, workflow_result)
+
+        await websocket_manager.register_workflow_websocket(workflow_run_output.run_id, websocket)  # type: ignore
+
+    except (InputCheckError, OutputCheckError) as e:
+        await websocket.send_text(
+            json.dumps(
+                {
+                    "event": "error",
+                    "error": str(e),
+                    "error_type": e.type,
+                    "error_id": e.error_id,
+                    "additional_data": e.additional_data,
+                }
+            )
+        )
     except Exception as e:
         logger.error(f"Error executing workflow via WebSocket: {e}")
-        await websocket.send_text(json.dumps({"event": "error", "error": str(e)}))
+        error_payload = {
+            "event": "error",
+            "error": str(e),
+            "error_type": e.type if hasattr(e, "type") else None,
+            "error_id": e.error_id if hasattr(e, "error_id") else None,
+        }
+        error_payload = {k: v for k, v in error_payload.items() if v is not None}
+        await websocket.send_text(json.dumps(error_payload))
 
 
 async def workflow_response_streamer(
@@ -286,17 +481,26 @@ async def workflow_response_streamer(
     **kwargs: Any,
 ) -> AsyncGenerator:
     try:
-        run_response = await workflow.arun(
+        run_response = workflow.arun(
             input=input,
             session_id=session_id,
             user_id=user_id,
             stream=True,
-            stream_intermediate_steps=True,
+            stream_events=True,
             **kwargs,
         )
 
         async for run_response_chunk in run_response:
-            yield format_sse_event(run_response_chunk.to_json())
+            yield format_sse_event(run_response_chunk)  # type: ignore
+
+    except (InputCheckError, OutputCheckError) as e:
+        error_response = WorkflowErrorEvent(
+            error=str(e),
+            error_type=e.type,
+            error_id=e.error_id,
+            additional_data=e.additional_data,
+        )
+        yield format_sse_event(error_response)
 
     except Exception as e:
         import traceback
@@ -304,11 +508,84 @@ async def workflow_response_streamer(
         traceback.print_exc()
         error_response = WorkflowErrorEvent(
             error=str(e),
+            error_type=e.type if hasattr(e, "type") else None,
+            error_id=e.error_id if hasattr(e, "error_id") else None,
         )
-        yield format_sse_event(error_response.to_json())
+        yield format_sse_event(error_response)
         return
 
 
+def get_websocket_router(
+    os: "AgentOS",
+    settings: AgnoAPISettings = AgnoAPISettings(),
+) -> APIRouter:
+    """
+    Create WebSocket router without HTTP authentication dependencies.
+    WebSocket endpoints handle authentication internally via message-based auth.
+    """
+    ws_router = APIRouter()
+
+    @ws_router.websocket(
+        "/workflows/ws",
+        name="workflow_websocket",
+    )
+    async def workflow_websocket_endpoint(websocket: WebSocket):
+        """WebSocket endpoint for receiving real-time workflow events"""
+        requires_auth = bool(settings.os_security_key)
+        await websocket_manager.connect(websocket, requires_auth=requires_auth)
+
+        try:
+            while True:
+                data = await websocket.receive_text()
+                message = json.loads(data)
+                action = message.get("action")
+
+                # Handle authentication first
+                if action == "authenticate":
+                    token = message.get("token")
+                    if not token:
+                        await websocket.send_text(json.dumps({"event": "auth_error", "error": "Token is required"}))
+                        continue
+
+                    if validate_websocket_token(token, settings):
+                        await websocket_manager.authenticate_websocket(websocket)
+                    else:
+                        await websocket.send_text(json.dumps({"event": "auth_error", "error": "Invalid token"}))
+                        continue
+
+                # Check authentication for all other actions (only when required)
+                elif requires_auth and not websocket_manager.is_authenticated(websocket):
+                    await websocket.send_text(
+                        json.dumps(
+                            {
+                                "event": "auth_required",
+                                "error": "Authentication required. Send authenticate action with valid token.",
+                            }
+                        )
+                    )
+                    continue
+
+                # Handle authenticated actions
+                elif action == "ping":
+                    await websocket.send_text(json.dumps({"event": "pong"}))
+
+                elif action == "start-workflow":
+                    # Handle workflow execution directly via WebSocket
+                    await handle_workflow_via_websocket(websocket, message, os)
+
+                else:
+                    await websocket.send_text(json.dumps({"event": "error", "error": f"Unknown action: {action}"}))
+
+        except Exception as e:
+            if "1012" not in str(e) and "1001" not in str(e):
+                logger.error(f"WebSocket error: {e}")
+        finally:
+            # Clean up the websocket connection
+            await websocket_manager.disconnect_websocket(websocket)
+
+    return ws_router
+
+
 def get_base_router(
     os: "AgentOS",
     settings: AgnoAPISettings = AgnoAPISettings(),
@@ -321,7 +598,6 @@ def get_base_router(
     - Agent management and execution
     - Team collaboration and coordination
     - Workflow automation and orchestration
-    - Real-time WebSocket communications
 
     All endpoints include detailed documentation, examples, and proper error handling.
     """
@@ -337,24 +613,6 @@ def get_base_router(
     )
 
     # -- Main Routes ---
-
-    @router.get(
-        "/health",
-        tags=["Core"],
-        operation_id="health_check",
-        summary="Health Check",
-        description="Check the health status of the AgentOS API. Returns a simple status indicator.",
-        response_model=HealthResponse,
-        responses={
-            200: {
-                "description": "API is healthy and operational",
-                "content": {"application/json": {"example": {"status": "ok"}}},
-            }
-        },
-    )
-    async def health_check() -> HealthResponse:
-        return HealthResponse(status="ok")
-
     @router.get(
         "/config",
         response_model=ConfigResponse,
@@ -375,7 +633,7 @@ def get_base_router(
                 "content": {
                     "application/json": {
                         "example": {
-                            "os_id": "demo",
+                            "id": "demo",
                             "description": "Example AgentOS configuration",
                             "available_models": [],
                             "databases": ["9c884dc4-9066-448c-9074-ef49ec7eb73c"],
@@ -437,10 +695,10 @@ def get_base_router(
     )
     async def config() -> ConfigResponse:
         return ConfigResponse(
-            os_id=os.os_id or "Unnamed OS",
+            os_id=os.id or "Unnamed OS",
             description=os.description,
             available_models=os.config.available_models if os.config else [],
-            databases=[db.id for db in os.dbs.values()],
+            databases=list({db.id for db_id, dbs in os.dbs.items() for db in dbs}),
             chat=os.config.chat if os.config else None,
             session=os._get_session_config(),
             memory=os._get_memory_config(),
@@ -451,7 +709,7 @@ def get_base_router(
             teams=[TeamSummaryResponse.from_team(team) for team in os.teams] if os.teams else [],
             workflows=[WorkflowSummaryResponse.from_workflow(w) for w in os.workflows] if os.workflows else [],
             interfaces=[
-                InterfaceResponse(type=interface.type, version=interface.version, route=interface.router_prefix)
+                InterfaceResponse(type=interface.type, version=interface.version, route=interface.prefix)
                 for interface in os.interfaces
             ],
         )
@@ -524,7 +782,7 @@ def get_base_router(
                 "content": {
                     "text/event-stream": {
                         "examples": {
-                            "event_strea": {
+                            "event_stream": {
                                 "summary": "Example event stream response",
                                 "value": 'event: RunStarted\ndata: {"content": "Hello!", "run_id": "123..."}\n\n',
                             }
@@ -538,12 +796,39 @@ def get_base_router(
     )
     async def create_agent_run(
         agent_id: str,
+        request: Request,
         message: str = Form(...),
         stream: bool = Form(False),
         session_id: Optional[str] = Form(None),
         user_id: Optional[str] = Form(None),
         files: Optional[List[UploadFile]] = File(None),
     ):
+        kwargs = await _get_request_kwargs(request, create_agent_run)
+
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+        if hasattr(request.state, "metadata"):
+            metadata = request.state.metadata
+            if "metadata" in kwargs:
+                log_warning("Metadata parameter passed in both request state and kwargs, using request state")
+            kwargs["metadata"] = metadata
+
         agent = get_agent_by_id(agent_id, os.agents)
         if agent is None:
             raise HTTPException(status_code=404, detail="Agent not found")
@@ -559,19 +844,39 @@ def get_base_router(
 
         if files:
             for file in files:
-                if file.content_type in ["image/png", "image/jpeg", "image/jpg", "image/webp"]:
+                if file.content_type in [
+                    "image/png",
+                    "image/jpeg",
+                    "image/jpg",
+                    "image/gif",
+                    "image/webp",
+                    "image/bmp",
+                    "image/tiff",
+                    "image/tif",
+                    "image/avif",
+                ]:
                     try:
                         base64_image = process_image(file)
                         base64_images.append(base64_image)
                     except Exception as e:
                         log_error(f"Error processing image {file.filename}: {e}")
                         continue
-                elif file.content_type in ["audio/wav", "audio/mp3", "audio/mpeg"]:
+                elif file.content_type in [
+                    "audio/wav",
+                    "audio/wave",
+                    "audio/mp3",
+                    "audio/mpeg",
+                    "audio/ogg",
+                    "audio/mp4",
+                    "audio/m4a",
+                    "audio/aac",
+                    "audio/flac",
+                ]:
                     try:
-                        base64_audio = process_audio(file)
-                        base64_audios.append(base64_audio)
+                        audio = process_audio(file)
+                        base64_audios.append(audio)
                     except Exception as e:
-                        log_error(f"Error processing audio {file.filename}: {e}")
+                        log_error(f"Error processing audio {file.filename} with content type {file.content_type}: {e}")
                         continue
                 elif file.content_type in [
                     "video/x-flv",
@@ -594,15 +899,25 @@ def get_base_router(
                         continue
                 elif file.content_type in [
                     "application/pdf",
-                    "text/csv",
+                    "application/json",
+                    "application/x-javascript",
                     "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+                    "text/javascript",
+                    "application/x-python",
+                    "text/x-python",
                     "text/plain",
-                    "application/json",
+                    "text/html",
+                    "text/css",
+                    "text/md",
+                    "text/csv",
+                    "text/xml",
+                    "text/rtf",
                 ]:
                     # Process document files
                     try:
-                        file_content = await file.read()
-                        input_files.append(FileMedia(content=file_content))
+                        input_file = process_document(file)
+                        if input_file is not None:
+                            input_files.append(input_file)
                     except Exception as e:
                         log_error(f"Error processing file {file.filename}: {e}")
                         continue
@@ -620,24 +935,30 @@ def get_base_router(
                     audio=base64_audios if base64_audios else None,
                     videos=base64_videos if base64_videos else None,
                     files=input_files if input_files else None,
+                    **kwargs,
                 ),
                 media_type="text/event-stream",
             )
         else:
-            run_response = cast(
-                RunOutput,
-                await agent.arun(
-                    input=message,
-                    session_id=session_id,
-                    user_id=user_id,
-                    images=base64_images if base64_images else None,
-                    audio=base64_audios if base64_audios else None,
-                    videos=base64_videos if base64_videos else None,
-                    files=input_files if input_files else None,
-                    stream=False,
-                ),
-            )
-            return run_response.to_dict()
+            try:
+                run_response = cast(
+                    RunOutput,
+                    await agent.arun(
+                        input=message,
+                        session_id=session_id,
+                        user_id=user_id,
+                        images=base64_images if base64_images else None,
+                        audio=base64_audios if base64_audios else None,
+                        videos=base64_videos if base64_videos else None,
+                        files=input_files if input_files else None,
+                        stream=False,
+                        **kwargs,
+                    ),
+                )
+                return run_response.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.post(
         "/agents/{agent_id}/runs/{run_id}/cancel",
@@ -698,11 +1019,17 @@ def get_base_router(
     async def continue_agent_run(
         agent_id: str,
         run_id: str,
+        request: Request,
         tools: str = Form(...),  # JSON string of tools
         session_id: Optional[str] = Form(None),
         user_id: Optional[str] = Form(None),
         stream: bool = Form(True),
     ):
+        if hasattr(request.state, "user_id"):
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            session_id = request.state.session_id
+
         # Parse the JSON string manually
         try:
             tools_data = json.loads(tools) if tools else None
@@ -740,17 +1067,21 @@ def get_base_router(
                 media_type="text/event-stream",
             )
         else:
-            run_response_obj = cast(
-                RunOutput,
-                await agent.acontinue_run(
-                    run_id=run_id,  # run_id from path
-                    updated_tools=updated_tools,
-                    session_id=session_id,
-                    user_id=user_id,
-                    stream=False,
-                ),
-            )
-            return run_response_obj.to_dict()
+            try:
+                run_response_obj = cast(
+                    RunOutput,
+                    await agent.acontinue_run(
+                        run_id=run_id,  # run_id from path
+                        updated_tools=updated_tools,
+                        session_id=session_id,
+                        user_id=user_id,
+                        stream=False,
+                    ),
+                )
+                return run_response_obj.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.get(
         "/agents",
@@ -797,7 +1128,8 @@ def get_base_router(
 
         agents = []
         for agent in os.agents:
-            agents.append(AgentResponse.from_agent(agent=agent))
+            agent_response = await AgentResponse.from_agent(agent=agent)
+            agents.append(agent_response)
 
         return agents
 
@@ -844,7 +1176,7 @@ def get_base_router(
         if agent is None:
             raise HTTPException(status_code=404, detail="Agent not found")
 
-        return AgentResponse.from_agent(agent)
+        return await AgentResponse.from_agent(agent)
 
     # -- Team routes ---
 
@@ -880,6 +1212,7 @@ def get_base_router(
     )
     async def create_team_run(
         team_id: str,
+        request: Request,
         message: str = Form(...),
         stream: bool = Form(True),
         monitor: bool = Form(True),
@@ -887,7 +1220,34 @@ def get_base_router(
         user_id: Optional[str] = Form(None),
         files: Optional[List[UploadFile]] = File(None),
     ):
-        logger.debug(f"Creating team run: {message} {session_id} {monitor} {user_id} {team_id} {files}")
+        kwargs = await _get_request_kwargs(request, create_team_run)
+
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+        if hasattr(request.state, "metadata"):
+            metadata = request.state.metadata
+            if "metadata" in kwargs:
+                log_warning("Metadata parameter passed in both request state and kwargs, using request state")
+            kwargs["metadata"] = metadata
+
+        logger.debug(f"Creating team run: {message=} {session_id=} {monitor=} {user_id=} {team_id=} {files=} {kwargs=}")
+
         team = get_team_by_id(team_id, os.teams)
         if team is None:
             raise HTTPException(status_code=404, detail="Team not found")
@@ -962,21 +1322,27 @@ def get_base_router(
                     audio=base64_audios if base64_audios else None,
                     videos=base64_videos if base64_videos else None,
                     files=document_files if document_files else None,
+                    **kwargs,
                 ),
                 media_type="text/event-stream",
             )
         else:
-            run_response = await team.arun(
-                input=message,
-                session_id=session_id,
-                user_id=user_id,
-                images=base64_images if base64_images else None,
-                audio=base64_audios if base64_audios else None,
-                videos=base64_videos if base64_videos else None,
-                files=document_files if document_files else None,
-                stream=False,
-            )
-            return run_response.to_dict()
+            try:
+                run_response = await team.arun(
+                    input=message,
+                    session_id=session_id,
+                    user_id=user_id,
+                    images=base64_images if base64_images else None,
+                    audio=base64_audios if base64_audios else None,
+                    videos=base64_videos if base64_videos else None,
+                    files=document_files if document_files else None,
+                    stream=False,
+                    **kwargs,
+                )
+                return run_response.to_dict()
+
+            except InputCheckError as e:
+                raise HTTPException(status_code=400, detail=str(e))
 
     @router.post(
         "/teams/{team_id}/runs/{run_id}/cancel",
@@ -1095,7 +1461,8 @@ def get_base_router(
 
         teams = []
         for team in os.teams:
-            teams.append(TeamResponse.from_team(team=team))
+            team_response = await TeamResponse.from_team(team=team)
+            teams.append(team_response)
 
         return teams
 
@@ -1188,39 +1555,10 @@ def get_base_router(
         if team is None:
             raise HTTPException(status_code=404, detail="Team not found")
 
-        return TeamResponse.from_team(team)
+        return await TeamResponse.from_team(team)
 
     # -- Workflow routes ---
 
-    @router.websocket(
-        "/workflows/ws",
-        name="workflow_websocket",
-    )
-    async def workflow_websocket_endpoint(websocket: WebSocket):
-        """WebSocket endpoint for receiving real-time workflow events"""
-        await websocket_manager.connect(websocket)
-
-        try:
-            while True:
-                data = await websocket.receive_text()
-                message = json.loads(data)
-                action = message.get("action")
-
-                if action == "ping":
-                    await websocket.send_text(json.dumps({"event": "pong"}))
-
-                elif action == "start-workflow":
-                    # Handle workflow execution directly via WebSocket
-                    await handle_workflow_via_websocket(websocket, message, os)
-        except Exception as e:
-            if "1012" not in str(e):
-                logger.error(f"WebSocket error: {e}")
-        finally:
-            # Clean up any run_ids associated with this websocket
-            runs_to_remove = [run_id for run_id, ws in websocket_manager.active_connections.items() if ws == websocket]
-            for run_id in runs_to_remove:
-                await websocket_manager.disconnect_by_run_id(run_id)
-
     @router.get(
         "/workflows",
         response_model=List[WorkflowSummaryResponse],
@@ -1290,7 +1628,7 @@ def get_base_router(
         if workflow is None:
             raise HTTPException(status_code=404, detail="Workflow not found")
 
-        return WorkflowResponse.from_workflow(workflow)
+        return await WorkflowResponse.from_workflow(workflow)
 
     @router.post(
         "/workflows/{workflow_id}/runs",
@@ -1328,12 +1666,38 @@ def get_base_router(
     )
     async def create_workflow_run(
         workflow_id: str,
+        request: Request,
         message: str = Form(...),
         stream: bool = Form(True),
         session_id: Optional[str] = Form(None),
         user_id: Optional[str] = Form(None),
-        **kwargs: Any,
     ):
+        kwargs = await _get_request_kwargs(request, create_workflow_run)
+
+        if hasattr(request.state, "user_id"):
+            if user_id:
+                log_warning("User ID parameter passed in both request state and kwargs, using request state")
+            user_id = request.state.user_id
+        if hasattr(request.state, "session_id"):
+            if session_id:
+                log_warning("Session ID parameter passed in both request state and kwargs, using request state")
+            session_id = request.state.session_id
+        if hasattr(request.state, "session_state"):
+            session_state = request.state.session_state
+            if "session_state" in kwargs:
+                log_warning("Session state parameter passed in both request state and kwargs, using request state")
+            kwargs["session_state"] = session_state
+        if hasattr(request.state, "dependencies"):
+            dependencies = request.state.dependencies
+            if "dependencies" in kwargs:
+                log_warning("Dependencies parameter passed in both request state and kwargs, using request state")
+            kwargs["dependencies"] = dependencies
+        if hasattr(request.state, "metadata"):
+            metadata = request.state.metadata
+            if "metadata" in kwargs:
+                log_warning("Metadata parameter passed in both request state and kwargs, using request state")
+            kwargs["metadata"] = metadata
+
         # Retrieve the workflow by ID
         workflow = get_workflow_by_id(workflow_id, os.workflows)
         if workflow is None:
@@ -1367,6 +1731,9 @@ def get_base_router(
                     **kwargs,
                 )
                 return run_response.to_dict()
+
+        except InputCheckError as e:
+            raise HTTPException(status_code=400, detail=str(e))
         except Exception as e:
             # Handle unexpected runtime errors
             raise HTTPException(status_code=500, detail=f"Error running workflow: {str(e)}")
@@ -1397,4 +1764,54 @@ def get_base_router(
 
         return JSONResponse(content={}, status_code=200)
 
+    # -- Database Migration routes ---
+
+    @router.post(
+        "/databases/{db_id}/migrate",
+        tags=["Database"],
+        operation_id="migrate_database",
+        summary="Migrate Database",
+        description=(
+            "Migrate the given database schema to the given target version. "
+            "If a target version is not provided, the database will be migrated to the latest version. "
+        ),
+        responses={
+            200: {
+                "description": "Database migrated successfully",
+                "content": {
+                    "application/json": {
+                        "example": {"message": "Database migrated successfully to version 3.0.0"},
+                    }
+                },
+            },
+            404: {"description": "Database not found", "model": NotFoundResponse},
+            500: {"description": "Failed to migrate database", "model": InternalServerErrorResponse},
+        },
+    )
+    async def migrate_database(db_id: str, target_version: Optional[str] = None):
+        db = await get_db(os.dbs, db_id)
+        if not db:
+            raise HTTPException(status_code=404, detail="Database not found")
+
+        if target_version:
+            
+            # Use the session table as proxy for the database schema version
+            if isinstance(db, AsyncBaseDb):
+                current_version = await db.get_latest_schema_version(db.session_table_name)
+            else:
+                current_version = db.get_latest_schema_version(db.session_table_name)
+
+            if version.parse(target_version) > version.parse(current_version):  # type: ignore
+                MigrationManager(db).up(target_version)  # type: ignore
+            else:
+                MigrationManager(db).down(target_version)  # type: ignore
+
+        # If the target version is not provided, migrate to the latest version
+        else:
+            MigrationManager(db).up()  # type: ignore
+
+        return JSONResponse(
+            content={"message": f"Database migrated successfully to version {target_version}"}, status_code=200
+        )
+
     return router