agno 2.0.0rc1__py3-none-any.whl → 2.0.1__py3-none-any.whl

agno/os/mcp.py

@@ -25,7 +25,6 @@ from agno.os.schema import (
 from agno.os.utils import (
     get_agent_by_id,
     get_db,
-    get_knowledge_instance_by_db_id,
     get_team_by_id,
     get_workflow_by_id,
 )
@@ -50,7 +49,7 @@ def get_mcp_server(
     @mcp.tool(
         name="get_agentos_config",
         description="Get the configuration of the AgentOS",
-        tags=["core"],
+        tags={"core"},
         output_schema=ConfigResponse.model_json_schema(),
     )  # type: ignore
     async def config() -> ConfigResponse:
@@ -74,21 +73,21 @@ def get_mcp_server(
             ],
         )
 
-    @mcp.tool(name="run_agent", description="Run an agent", tags=["core"])  # type: ignore
+    @mcp.tool(name="run_agent", description="Run an agent", tags={"core"})  # type: ignore
     async def run_agent(agent_id: str, message: str) -> RunOutput:
         agent = get_agent_by_id(agent_id, os.agents)
         if agent is None:
             raise Exception(f"Agent {agent_id} not found")
         return agent.run(message)
 
-    @mcp.tool(name="run_team", description="Run a team", tags=["core"])  # type: ignore
+    @mcp.tool(name="run_team", description="Run a team", tags={"core"})  # type: ignore
     async def run_team(team_id: str, message: str) -> TeamRunOutput:
         team = get_team_by_id(team_id, os.teams)
         if team is None:
             raise Exception(f"Team {team_id} not found")
         return team.run(message)
 
-    @mcp.tool(name="run_workflow", description="Run a workflow", tags=["core"])  # type: ignore
+    @mcp.tool(name="run_workflow", description="Run a workflow", tags={"core"})  # type: ignore
     async def run_workflow(workflow_id: str, message: str) -> WorkflowRunOutput:
         workflow = get_workflow_by_id(workflow_id, os.workflows)
         if workflow is None:
@@ -96,7 +95,7 @@ def get_mcp_server(
         return workflow.run(message)
 
     # Session Management Tools
-    @mcp.tool(name="get_sessions_for_agent", description="Get list of sessions for an agent", tags=["session"])  # type: ignore
+    @mcp.tool(name="get_sessions_for_agent", description="Get list of sessions for an agent", tags={"session"})  # type: ignore
     async def get_sessions_for_agent(
         agent_id: str,
         db_id: str,
@@ -105,7 +104,7 @@ def get_mcp_server(
         sort_order: str = "desc",
     ):
         db = get_db(os.dbs, db_id)
-        sessions = db.get_sessions(
+        sessions, _ = db.get_sessions(
             session_type=SessionType.AGENT,
             component_id=agent_id,
             user_id=user_id,
@@ -118,7 +117,7 @@ def get_mcp_server(
             "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
         }
 
-    @mcp.tool(name="get_sessions_for_team", description="Get list of sessions for a team", tags=["session"])  # type: ignore
+    @mcp.tool(name="get_sessions_for_team", description="Get list of sessions for a team", tags={"session"})  # type: ignore
     async def get_sessions_for_team(
         team_id: str,
         db_id: str,
@@ -127,7 +126,7 @@ def get_mcp_server(
         sort_order: str = "desc",
     ):
         db = get_db(os.dbs, db_id)
-        sessions = db.get_sessions(
+        sessions, _ = db.get_sessions(
             session_type=SessionType.TEAM,
             component_id=team_id,
             user_id=user_id,
@@ -137,11 +136,33 @@ def get_mcp_server(
         )
 
         return {
-            "data": sessions,  # type: ignore
+            "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
+        }
+
+    @mcp.tool(name="get_sessions_for_workflow", description="Get list of sessions for a workflow", tags={"session"})  # type: ignore
+    async def get_sessions_for_workflow(
+        workflow_id: str,
+        db_id: str,
+        user_id: Optional[str] = None,
+        sort_by: str = "created_at",
+        sort_order: str = "desc",
+    ):
+        db = get_db(os.dbs, db_id)
+        sessions, _ = db.get_sessions(
+            session_type=SessionType.WORKFLOW,
+            component_id=workflow_id,
+            user_id=user_id,
+            sort_by=sort_by,
+            sort_order=sort_order,
+            deserialize=False,
+        )
+
+        return {
+            "data": [SessionSchema.from_dict(session) for session in sessions],  # type: ignore
         }
 
     # Memory Management Tools
-    @mcp.tool(name="create_memory", description="Create a new user memory", tags=["memory"])  # type: ignore
+    @mcp.tool(name="create_memory", description="Create a new user memory", tags={"memory"})  # type: ignore
     async def create_memory(
         db_id: str,
         memory: str,
@@ -161,27 +182,27 @@ def get_mcp_server(
         if not user_memory:
             raise Exception("Failed to create memory")
 
-        return UserMemorySchema.from_dict(user_memory.to_dict())  # type: ignore
+        return UserMemorySchema.from_dict(user_memory)  # type: ignore
 
-    @mcp.tool(name="get_memories_for_user", description="Get a list of memories for a user", tags=["memory"])  # type: ignore
-    async def get_memories(
+    @mcp.tool(name="get_memories_for_user", description="Get a list of memories for a user", tags={"memory"})  # type: ignore
+    async def get_memories_for_user(
         user_id: str,
         sort_by: str = "updated_at",
         sort_order: str = "desc",
         db_id: Optional[str] = None,
     ):
         db = get_db(os.dbs, db_id)
-        user_memories = db.get_user_memories(
+        user_memories, _ = db.get_user_memories(
             user_id=user_id,
             sort_by=sort_by,
             sort_order=sort_order,
             deserialize=False,
         )
         return {
-            "data": user_memories,  # type: ignore
+            "data": [UserMemorySchema.from_dict(user_memory) for user_memory in user_memories],  # type: ignore
         }
 
-    @mcp.tool(name="update_memory", description="Update a memory", tags=["memory"])  # type: ignore
+    @mcp.tool(name="update_memory", description="Update a memory", tags={"memory"})  # type: ignore
     async def update_memory(
         db_id: str,
         memory_id: str,
@@ -202,7 +223,7 @@ def get_mcp_server(
 
         return UserMemorySchema.from_dict(user_memory)  # type: ignore
 
-    @mcp.tool(name="delete_memory", description="Delete a memory by ID", tags=["memory"])  # type: ignore
+    @mcp.tool(name="delete_memory", description="Delete a memory by ID", tags={"memory"})  # type: ignore
     async def delete_memory(
         db_id: str,
         memory_id: str,
@@ -210,46 +231,5 @@ def get_mcp_server(
         db = get_db(os.dbs, db_id)
         db.delete_user_memory(memory_id=memory_id)
 
-    # Knowledge Management Tools
-    @mcp.tool(name="get_content", description="Get paginated list of knowledge content", tags=["knowledge"])  # type: ignore
-    async def get_content(
-        sort_by: str = "created_at",
-        sort_order: str = "desc",
-        db_id: Optional[str] = None,
-    ):
-        knowledge = get_knowledge_instance_by_db_id(
-            os.knowledge_instances if hasattr(os, "knowledge_instances") else [], db_id
-        )
-        contents, count = knowledge.get_content(sort_by=sort_by, sort_order=sort_order)
-
-        return {
-            "data": [
-                {
-                    "id": content.id,
-                    "name": content.name,
-                    "description": content.description,
-                    "file_type": content.file_type,
-                    "size": content.size,
-                    "metadata": content.metadata,
-                    "status": content.status,
-                    "status_message": content.status_message,
-                    "created_at": content.created_at,
-                    "updated_at": content.updated_at,
-                }
-                for content in contents
-            ]
-        }
-
-    @mcp.tool(name="delete_content_by_id", description="Delete knowledge content by ID", tags=["knowledge"])  # type: ignore
-    async def delete_content_by_id(
-        content_id: str,
-        db_id: str,
-    ) -> str:
-        knowledge = get_knowledge_instance_by_db_id(
-            os.knowledge_instances if hasattr(os, "knowledge_instances") else [], db_id
-        )
-        knowledge.remove_content_by_id(content_id=content_id)
-        return "Successfully deleted content"
-
     mcp_app = mcp.http_app(path="/mcp")
     return mcp_app