letta-nightly 0.7.30.dev20250603104343__py3-none-any.whl → 0.8.0.dev20250604104349__py3-none-any.whl

letta/services/identity_manager.py

@@ -3,7 +3,6 @@ from typing import List, Optional
 from fastapi import HTTPException
 from sqlalchemy import select
 from sqlalchemy.exc import NoResultFound
-from sqlalchemy.orm import Session
 
 from letta.orm.agent import Agent as AgentModel
 from letta.orm.block import Block as BlockModel
@@ -60,26 +59,29 @@ class IdentityManager:
     @trace_method
     async def create_identity_async(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
         async with db_registry.async_session() as session:
-            new_identity = IdentityModel(**identity.model_dump(exclude={"agent_ids", "block_ids"}, exclude_unset=True))
-            new_identity.organization_id = actor.organization_id
-            await self._process_relationship_async(
-                session=session,
-                identity=new_identity,
-                relationship_name="agents",
-                model_class=AgentModel,
-                item_ids=identity.agent_ids,
-                allow_partial=False,
-            )
-            await self._process_relationship_async(
-                session=session,
-                identity=new_identity,
-                relationship_name="blocks",
-                model_class=BlockModel,
-                item_ids=identity.block_ids,
-                allow_partial=False,
-            )
-            await new_identity.create_async(session, actor=actor)
-            return new_identity.to_pydantic()
+            return await self._create_identity_async(db_session=session, identity=identity, actor=actor)
+
+    async def _create_identity_async(self, db_session, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
+        new_identity = IdentityModel(**identity.model_dump(exclude={"agent_ids", "block_ids"}, exclude_unset=True))
+        new_identity.organization_id = actor.organization_id
+        await self._process_relationship_async(
+            db_session=db_session,
+            identity=new_identity,
+            relationship_name="agents",
+            model_class=AgentModel,
+            item_ids=identity.agent_ids,
+            allow_partial=False,
+        )
+        await self._process_relationship_async(
+            db_session=db_session,
+            identity=new_identity,
+            relationship_name="blocks",
+            model_class=BlockModel,
+            item_ids=identity.block_ids,
+            allow_partial=False,
+        )
+        await new_identity.create_async(db_session=db_session, actor=actor)
+        return new_identity.to_pydantic()
 
     @enforce_types
     @trace_method
@@ -93,19 +95,19 @@ class IdentityManager:
                 actor=actor,
             )
 
-        if existing_identity is None:
-            return await self.create_identity_async(identity=IdentityCreate(**identity.model_dump()), actor=actor)
-        else:
-            identity_update = IdentityUpdate(
-                name=identity.name,
-                identifier_key=identity.identifier_key,
-                identity_type=identity.identity_type,
-                agent_ids=identity.agent_ids,
-                properties=identity.properties,
-            )
-            return await self._update_identity_async(
-                session=session, existing_identity=existing_identity, identity=identity_update, actor=actor, replace=True
-            )
+            if existing_identity is None:
+                return await self._create_identity_async(db_session=session, identity=IdentityCreate(**identity.model_dump()), actor=actor)
+            else:
+                identity_update = IdentityUpdate(
+                    name=identity.name,
+                    identifier_key=identity.identifier_key,
+                    identity_type=identity.identity_type,
+                    agent_ids=identity.agent_ids,
+                    properties=identity.properties,
+                )
+                return await self._update_identity_async(
+                    db_session=session, existing_identity=existing_identity, identity=identity_update, actor=actor, replace=True
+                )
 
     @enforce_types
     @trace_method
@@ -121,12 +123,12 @@ class IdentityManager:
                 raise HTTPException(status_code=403, detail="Forbidden")
 
             return await self._update_identity_async(
-                session=session, existing_identity=existing_identity, identity=identity, actor=actor, replace=replace
+                db_session=session, existing_identity=existing_identity, identity=identity, actor=actor, replace=replace
             )
 
     async def _update_identity_async(
         self,
-        session: Session,
+        db_session,
         existing_identity: IdentityModel,
         identity: IdentityUpdate,
         actor: PydanticUser,
@@ -149,7 +151,7 @@ class IdentityManager:
 
         if identity.agent_ids is not None:
             await self._process_relationship_async(
-                session=session,
+                db_session=db_session,
                 identity=existing_identity,
                 relationship_name="agents",
                 model_class=AgentModel,
@@ -159,7 +161,7 @@ class IdentityManager:
             )
         if identity.block_ids is not None:
             await self._process_relationship_async(
-                session=session,
+                db_session=db_session,
                 identity=existing_identity,
                 relationship_name="blocks",
                 model_class=BlockModel,
@@ -167,7 +169,7 @@ class IdentityManager:
                 allow_partial=False,
                 replace=replace,
             )
-        await existing_identity.update_async(session, actor=actor)
+        await existing_identity.update_async(db_session=db_session, actor=actor)
         return existing_identity.to_pydantic()
 
     @enforce_types
@@ -180,7 +182,7 @@ class IdentityManager:
             if existing_identity is None:
                 raise HTTPException(status_code=404, detail="Identity not found")
             return await self._update_identity_async(
-                session=session,
+                db_session=session,
                 existing_identity=existing_identity,
                 identity=IdentityUpdate(properties=properties),
                 actor=actor,
@@ -213,7 +215,7 @@ class IdentityManager:
 
     async def _process_relationship_async(
         self,
-        session: Session,
+        db_session,
         identity: PydanticIdentity,
         relationship_name: str,
         model_class,
@@ -228,7 +230,7 @@ class IdentityManager:
             return
 
         # Retrieve models for the provided IDs
-        found_items = (await session.execute(select(model_class).where(model_class.id.in_(item_ids)))).scalars().all()
+        found_items = (await db_session.execute(select(model_class).where(model_class.id.in_(item_ids)))).scalars().all()
 
         # Validate all items are found if allow_partial is False
         if not allow_partial and len(found_items) != len(item_ids):

letta/services/job_manager.py

@@ -169,6 +169,7 @@ class JobManager:
         statuses: Optional[List[JobStatus]] = None,
         job_type: JobType = JobType.JOB,
         ascending: bool = True,
+        source_id: Optional[str] = None,
     ) -> List[PydanticJob]:
         """List all jobs with optional pagination and status filter."""
         async with db_registry.async_session() as session:
@@ -178,6 +179,9 @@ class JobManager:
             if statuses:
                 filter_kwargs["status"] = statuses
 
+            if source_id:
+                filter_kwargs["metadata_.source_id"] = source_id
+
             jobs = await JobModel.list_async(
                 db_session=session,
                 before=before,
@@ -197,6 +201,15 @@ class JobManager:
             job.hard_delete(db_session=session, actor=actor)
             return job.to_pydantic()
 
+    @enforce_types
+    @trace_method
+    async def delete_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob:
+        """Delete a job by its ID."""
+        async with db_registry.async_session() as session:
+            job = await self._verify_job_access_async(session=session, job_id=job_id, actor=actor)
+            await job.hard_delete_async(db_session=session, actor=actor)
+            return job.to_pydantic()
+
     @enforce_types
     @trace_method
     def get_job_messages(
@@ -599,25 +612,23 @@ class JobManager:
 
     async def _dispatch_callback_async(self, session, job: JobModel) -> None:
         """
-        POST a standard JSON payload to job.callback_url
-        and record timestamp + HTTP status asynchronously.
+        POST a standard JSON payload to job.callback_url and record timestamp + HTTP status asynchronously.
         """
-
         payload = {
             "job_id": job.id,
             "status": job.status,
-            "completed_at": job.completed_at.isoformat(),
+            "completed_at": job.completed_at.isoformat() if job.completed_at else None,
         }
+
         try:
             import httpx
 
             async with httpx.AsyncClient() as client:
                 resp = await client.post(job.callback_url, json=payload, timeout=5.0)
-                job.callback_sent_at = get_utc_time()
+                # Ensure timestamp is timezone-naive for DB compatibility
+                job.callback_sent_at = get_utc_time().replace(tzinfo=None)
                 job.callback_status_code = resp.status_code
-
         except Exception:
-            return
-
-        session.add(job)
-        await session.commit()
+            # Silently fail on callback errors - job updates should still succeed
+            # In production, this would include proper error logging
+            pass

letta/services/mcp/base_client.py

@@ -30,7 +30,7 @@ class AsyncBaseMCPClient:
             )
             raise e
 
-    async def _initialize_connection(self, exit_stack: AsyncExitStack[bool | None], server_config: BaseServerConfig) -> None:
+    async def _initialize_connection(self, server_config: BaseServerConfig) -> None:
         raise NotImplementedError("Subclasses must implement _initialize_connection")
 
     async def list_tools(self) -> list[MCPTool]:
@@ -55,7 +55,7 @@ class AsyncBaseMCPClient:
             # TODO move hardcoding to constants
             final_content = "Empty response from tool"
 
-        return final_content, result.isError
+        return final_content, not result.isError
 
     def _check_initialized(self):
         if not self.initialized:
@@ -65,3 +65,6 @@ class AsyncBaseMCPClient:
     async def cleanup(self):
         """Clean up resources"""
         await self.exit_stack.aclose()
+
+    def to_sync_client(self):
+        raise NotImplementedError("Subclasses must implement to_sync_client")

letta/services/mcp/sse_client.py

@@ -1,5 +1,3 @@
-from contextlib import AsyncExitStack
-
 from mcp import ClientSession
 from mcp.client.sse import sse_client
 
@@ -15,11 +13,11 @@ logger = get_logger(__name__)
 
 # TODO: Get rid of Async prefix on this class name once we deprecate old sync code
 class AsyncSSEMCPClient(AsyncBaseMCPClient):
-    async def _initialize_connection(self, exit_stack: AsyncExitStack[bool | None], server_config: SSEServerConfig) -> None:
+    async def _initialize_connection(self, server_config: SSEServerConfig) -> None:
         sse_cm = sse_client(url=server_config.server_url)
-        sse_transport = await exit_stack.enter_async_context(sse_cm)
+        sse_transport = await self.exit_stack.enter_async_context(sse_cm)
         self.stdio, self.write = sse_transport
 
         # Create and enter the ClientSession context manager
         session_cm = ClientSession(self.stdio, self.write)
-        self.session = await exit_stack.enter_async_context(session_cm)
+        self.session = await self.exit_stack.enter_async_context(session_cm)

letta/services/mcp/stdio_client.py

@@ -1,5 +1,3 @@
-from contextlib import AsyncExitStack
-
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 
@@ -12,8 +10,8 @@ logger = get_logger(__name__)
 
 # TODO: Get rid of Async prefix on this class name once we deprecate old sync code
 class AsyncStdioMCPClient(AsyncBaseMCPClient):
-    async def _initialize_connection(self, exit_stack: AsyncExitStack[bool | None], server_config: StdioServerConfig) -> None:
+    async def _initialize_connection(self, server_config: StdioServerConfig) -> None:
         server_params = StdioServerParameters(command=server_config.command, args=server_config.args)
-        stdio_transport = await exit_stack.enter_async_context(stdio_client(server_params))
+        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
         self.stdio, self.write = stdio_transport
-        self.session = await exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
+        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

letta/services/mcp_manager.py

@@ -0,0 +1,281 @@
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import letta.constants as constants
+from letta.functions.mcp_client.types import MCPServerType, MCPTool, SSEServerConfig, StdioServerConfig
+from letta.log import get_logger
+from letta.orm.errors import NoResultFound
+from letta.orm.mcp_server import MCPServer as MCPServerModel
+from letta.schemas.mcp import MCPServer, UpdateMCPServer, UpdateSSEMCPServer, UpdateStdioMCPServer
+from letta.schemas.tool import Tool as PydanticTool
+from letta.schemas.tool import ToolCreate
+from letta.schemas.user import User as PydanticUser
+from letta.server.db import db_registry
+from letta.services.mcp.sse_client import MCP_CONFIG_TOPLEVEL_KEY, AsyncSSEMCPClient
+from letta.services.mcp.stdio_client import AsyncStdioMCPClient
+from letta.services.tool_manager import ToolManager
+from letta.utils import enforce_types, printd
+
+logger = get_logger(__name__)
+
+
+class MCPManager:
+    """Manager class to handle business logic related to MCP."""
+
+    def __init__(self):
+        # TODO: timeouts?
+        self.tool_manager = ToolManager()
+        self.cached_mcp_servers = {}  # maps id -> async connection
+
+    @enforce_types
+    async def list_mcp_server_tools(self, mcp_server_name: str, actor: PydanticUser) -> List[MCPTool]:
+        """Get a list of all tools for a specific MCP server."""
+        print("mcp_server_name", mcp_server_name)
+        mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name, actor=actor)
+        mcp_config = await self.get_mcp_server_by_id_async(mcp_server_id, actor=actor)
+        server_config = mcp_config.to_config()
+
+        if mcp_config.server_type == MCPServerType.SSE:
+            mcp_client = AsyncSSEMCPClient(server_config=server_config)
+        elif mcp_config.server_type == MCPServerType.STDIO:
+            mcp_client = AsyncStdioMCPClient(server_config=server_config)
+        await mcp_client.connect_to_server()
+
+        # list tools
+        tools = await mcp_client.list_tools()
+        # TODO: change to pydantic tools
+
+        await mcp_client.cleanup()
+
+        return tools
+
+    @enforce_types
+    async def execute_mcp_server_tool(
+        self, mcp_server_name: str, tool_name: str, tool_args: Optional[Dict[str, Any]], actor: PydanticUser
+    ) -> Tuple[str, bool]:
+        """Call a specific tool from a specific MCP server."""
+
+        from letta.settings import tool_settings
+
+        if not tool_settings.mcp_read_from_config:
+            # read from DB
+            mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name, actor=actor)
+            mcp_config = await self.get_mcp_server_by_id_async(mcp_server_id, actor=actor)
+            server_config = mcp_config.to_config()
+        else:
+            # read from config file
+            mcp_config = self.read_mcp_config()
+            if mcp_server_name not in mcp_config:
+                print("MCP server not found in config.", mcp_config)
+                raise ValueError(f"MCP server {mcp_server_name} not found in config.")
+            server_config = mcp_config[mcp_server_name]
+
+        if isinstance(server_config, SSEServerConfig):
+            mcp_client = AsyncSSEMCPClient(server_config=server_config)
+        elif isinstance(server_config, StdioServerConfig):
+            mcp_client = AsyncStdioMCPClient(server_config=server_config)
+        await mcp_client.connect_to_server()
+
+        # call tool
+        result, success = await mcp_client.execute_tool(tool_name, tool_args)
+        logger.info(f"MCP Result: {result}, Success: {success}")
+        # TODO: change to pydantic tool
+
+        await mcp_client.cleanup()
+
+        return result, success
+
+    @enforce_types
+    async def add_tool_from_mcp_server(self, mcp_server_name: str, mcp_tool_name: str, actor: PydanticUser) -> PydanticTool:
+        """Add a tool from an MCP server to the Letta tool registry."""
+        mcp_tools = await self.list_mcp_server_tools(mcp_server_name, actor=actor)
+
+        for mcp_tool in mcp_tools:
+            if mcp_tool.name == mcp_tool_name:
+                tool_create = ToolCreate.from_mcp(mcp_server_name=mcp_server_name, mcp_tool=mcp_tool)
+                return await self.tool_manager.create_mcp_tool_async(tool_create=tool_create, mcp_server_name=mcp_server_name, actor=actor)
+
+        # failed to add - handle error?
+        return None
+
+    @enforce_types
+    async def list_mcp_servers(self, actor: PydanticUser) -> List[MCPServer]:
+        """List all MCP servers available"""
+        async with db_registry.async_session() as session:
+            mcp_servers = await MCPServerModel.list_async(
+                db_session=session,
+                organization_id=actor.organization_id,
+            )
+
+            return [mcp_server.to_pydantic() for mcp_server in mcp_servers]
+
+    @enforce_types
+    async def create_or_update_mcp_server(self, pydantic_mcp_server: MCPServer, actor: PydanticUser) -> MCPServer:
+        """Create a new tool based on the ToolCreate schema."""
+        mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name=pydantic_mcp_server.server_name, actor=actor)
+        print("FOUND SERVER", mcp_server_id, pydantic_mcp_server.server_name)
+        if mcp_server_id:
+            # Put to dict and remove fields that should not be reset
+            update_data = pydantic_mcp_server.model_dump(exclude_unset=True, exclude_none=True)
+
+            # If there's anything to update (can only update the configs, not the name)
+            if update_data:
+                if pydantic_mcp_server.server_type == MCPServerType.SSE:
+                    update_request = UpdateSSEMCPServer(server_url=pydantic_mcp_server.server_url)
+                elif pydantic_mcp_server.server_type == MCPServerType.STDIO:
+                    update_request = UpdateStdioMCPServer(stdio_config=pydantic_mcp_server.stdio_config)
+                mcp_server = await self.update_mcp_server_by_id(mcp_server_id, update_request, actor)
+                print("RETURN", mcp_server)
+            else:
+                printd(
+                    f"`create_or_update_mcp_server` was called with user_id={actor.id}, organization_id={actor.organization_id}, name={pydantic_mcp_server.server_name}, but found existing mcp server with nothing to update."
+                )
+                mcp_server = await self.get_mcp_server_by_id_async(mcp_server_id, actor=actor)
+        else:
+            mcp_server = await self.create_mcp_server(pydantic_mcp_server, actor=actor)
+
+        return mcp_server
+
+    @enforce_types
+    async def create_mcp_server(self, pydantic_mcp_server: MCPServer, actor: PydanticUser) -> PydanticTool:
+        """Create a new tool based on the ToolCreate schema."""
+        with db_registry.session() as session:
+            # Set the organization id at the ORM layer
+            pydantic_mcp_server.organization_id = actor.organization_id
+            mcp_server_data = pydantic_mcp_server.model_dump(to_orm=True)
+
+            mcp_server = MCPServerModel(**mcp_server_data)
+            mcp_server.create(session, actor=actor)  # Re-raise other database-related errors
+            return mcp_server.to_pydantic()
+
+    @enforce_types
+    async def update_mcp_server_by_id(self, mcp_server_id: str, mcp_server_update: UpdateMCPServer, actor: PydanticUser) -> PydanticTool:
+        """Update a tool by its ID with the given ToolUpdate object."""
+        async with db_registry.async_session() as session:
+            # Fetch the tool by ID
+            mcp_server = await MCPServerModel.read_async(db_session=session, identifier=mcp_server_id, actor=actor)
+
+            # Update tool attributes with only the fields that were explicitly set
+            update_data = mcp_server_update.model_dump(to_orm=True, exclude_none=True)
+            for key, value in update_data.items():
+                setattr(mcp_server, key, value)
+
+            mcp_server = await mcp_server.update_async(db_session=session, actor=actor)
+
+            # Save the updated tool to the database mcp_server = await mcp_server.update_async(db_session=session, actor=actor)
+            return mcp_server.to_pydantic()
+
+    @enforce_types
+    async def get_mcp_server_id_by_name(self, mcp_server_name: str, actor: PydanticUser) -> Optional[str]:
+        """Retrieve a MCP server by its name and a user"""
+        try:
+            async with db_registry.async_session() as session:
+                mcp_server = await MCPServerModel.read_async(db_session=session, server_name=mcp_server_name, actor=actor)
+                return mcp_server.id
+        except NoResultFound:
+            return None
+
+    @enforce_types
+    async def get_mcp_server_by_id_async(self, mcp_server_id: str, actor: PydanticUser) -> MCPServer:
+        """Fetch a tool by its ID."""
+        async with db_registry.async_session() as session:
+            # Retrieve tool by id using the Tool model's read method
+            mcp_server = await MCPServerModel.read_async(db_session=session, identifier=mcp_server_id, actor=actor)
+            # Convert the SQLAlchemy Tool object to PydanticTool
+            return mcp_server.to_pydantic()
+
+    @enforce_types
+    async def get_mcp_server(self, mcp_server_name: str, actor: PydanticUser) -> PydanticTool:
+        """Get a tool by name."""
+        async with db_registry.async_session() as session:
+            mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name, actor)
+            mcp_server = await MCPServerModel.read_async(db_session=session, identifier=mcp_server_id, actor=actor)
+            if not mcp_server:
+                raise HTTPException(
+                    status_code=404,  # Not Found
+                    detail={
+                        "code": "MCPServerNotFoundError",
+                        "message": f"MCP server {mcp_server_name} not found",
+                        "mcp_server_name": mcp_server_name,
+                    },
+                )
+            return mcp_server.to_pydantic()
+
+    # @enforce_types
+    # async def delete_mcp_server(self, mcp_server_name: str, actor: PydanticUser) -> None:
+    #    """Delete an existing tool."""
+    #    with db_registry.session() as session:
+    #        mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name, actor)
+    #        mcp_server = await MCPServerModel.read_async(db_session=session, identifier=mcp_server_id, actor=actor)
+    #        if not mcp_server:
+    #            raise HTTPException(
+    #                status_code=404,  # Not Found
+    #                detail={
+    #                    "code": "MCPServerNotFoundError",
+    #                    "message": f"MCP server {mcp_server_name} not found",
+    #                    "mcp_server_name": mcp_server_name,
+    #                },
+    #            )
+    #        mcp_server.delete(session, actor=actor)  # Re-raise other database-related errors
+
+    @enforce_types
+    def delete_mcp_server_by_id(self, mcp_server_id: str, actor: PydanticUser) -> None:
+        """Delete a tool by its ID."""
+        with db_registry.session() as session:
+            try:
+                mcp_server = MCPServerModel.read(db_session=session, identifier=mcp_server_id, actor=actor)
+                mcp_server.hard_delete(db_session=session, actor=actor)
+            except NoResultFound:
+                raise ValueError(f"MCP server with id {mcp_server_id} not found.")
+
+    def read_mcp_config(self) -> dict[str, Union[SSEServerConfig, StdioServerConfig]]:
+        mcp_server_list = {}
+
+        # Attempt to read from ~/.letta/mcp_config.json
+        mcp_config_path = os.path.join(constants.LETTA_DIR, constants.MCP_CONFIG_NAME)
+        if os.path.exists(mcp_config_path):
+            with open(mcp_config_path, "r") as f:
+
+                try:
+                    mcp_config = json.load(f)
+                except Exception as e:
+                    logger.error(f"Failed to parse MCP config file ({mcp_config_path}) as json: {e}")
+                    return mcp_server_list
+
+                # Proper formatting is "mcpServers" key at the top level,
+                # then a dict with the MCP server name as the key,
+                # with the value being the schema from StdioServerParameters
+                if MCP_CONFIG_TOPLEVEL_KEY in mcp_config:
+                    for server_name, server_params_raw in mcp_config[MCP_CONFIG_TOPLEVEL_KEY].items():
+
+                        # No support for duplicate server names
+                        if server_name in mcp_server_list:
+                            logger.error(f"Duplicate MCP server name found (skipping): {server_name}")
+                            continue
+
+                        if "url" in server_params_raw:
+                            # Attempt to parse the server params as an SSE server
+                            try:
+                                server_params = SSEServerConfig(
+                                    server_name=server_name,
+                                    server_url=server_params_raw["url"],
+                                )
+                                mcp_server_list[server_name] = server_params
+                            except Exception as e:
+                                logger.error(f"Failed to parse server params for MCP server {server_name} (skipping): {e}")
+                                continue
+                        else:
+                            # Attempt to parse the server params as a StdioServerParameters
+                            try:
+                                server_params = StdioServerConfig(
+                                    server_name=server_name,
+                                    command=server_params_raw["command"],
+                                    args=server_params_raw.get("args", []),
+                                    env=server_params_raw.get("env", {}),
+                                )
+                                mcp_server_list[server_name] = server_params
+                            except Exception as e:
+                                logger.error(f"Failed to parse server params for MCP server {server_name} (skipping): {e}")
+                                continue
+        return mcp_server_list