letta-nightly 0.6.37.dev20250311104150__py3-none-any.whl → 0.6.39.dev20250313104142__py3-none-any.whl

letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.37"
+__version__ = "0.6.39"
 
 # import clients
 from letta.client.client import LocalClient, RESTClient, create_client

letta/agent.py

@@ -3,7 +3,7 @@ import time
 import traceback
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from openai.types.beta.function_tool import FunctionTool as OpenAITool
 
@@ -26,9 +26,11 @@ from letta.helpers import ToolRulesSolver
 from letta.helpers.composio_helpers import get_composio_api_key
 from letta.helpers.datetime_helpers import get_utc_time
 from letta.helpers.json_helpers import json_dumps, json_loads
+from letta.helpers.mcp_helpers import BaseMCPClient
 from letta.interface import AgentInterface
 from letta.llm_api.helpers import calculate_summarizer_cutoff, get_token_counts_for_messages, is_context_overflow_error
 from letta.llm_api.llm_api_tools import create
+from letta.llm_api.llm_client import LLMClient
 from letta.local_llm.utils import num_tokens_from_functions, num_tokens_from_messages
 from letta.log import get_logger
 from letta.memory import summarize_messages
@@ -37,7 +39,7 @@ from letta.orm.enums import ToolType
 from letta.schemas.agent import AgentState, AgentStepResponse, UpdateAgent
 from letta.schemas.block import BlockUpdate
 from letta.schemas.embedding_config import EmbeddingConfig
-from letta.schemas.enums import MessageRole
+from letta.schemas.enums import MessageContentType, MessageRole
 from letta.schemas.memory import ContextWindowOverview, Memory
 from letta.schemas.message import Message, ToolReturn
 from letta.schemas.openai.chat_completion_response import ChatCompletionResponse
@@ -91,6 +93,8 @@ class Agent(BaseAgent):
         user: User,
         # extras
         first_message_verify_mono: bool = True,  # TODO move to config?
+        # MCP sessions, state held in-memory in the server
+        mcp_clients: Optional[Dict[str, BaseMCPClient]] = None,
     ):
         assert isinstance(agent_state.memory, Memory), f"Memory object is not of type Memory: {type(agent_state.memory)}"
         # Hold a copy of the state that was used to init the agent
@@ -148,18 +152,22 @@ class Agent(BaseAgent):
         # Logger that the Agent specifically can use, will also report the agent_state ID with the logs
         self.logger = get_logger(agent_state.id)
 
+        # MCPClient, state/sessions managed by the server
+        self.mcp_clients = mcp_clients
+
     def load_last_function_response(self):
         """Load the last function response from message history"""
         in_context_messages = self.agent_manager.get_in_context_messages(agent_id=self.agent_state.id, actor=self.user)
         for i in range(len(in_context_messages) - 1, -1, -1):
             msg = in_context_messages[i]
-            if msg.role == MessageRole.tool and msg.text:
+            if msg.role == MessageRole.tool and msg.content and len(msg.content) == 1 and msg.content[0].type == MessageContentType.text:
+                text_content = msg.content[0].text
                 try:
-                    response_json = json.loads(msg.text)
+                    response_json = json.loads(text_content)
                     if response_json.get("message"):
                         return response_json["message"]
                 except (json.JSONDecodeError, KeyError):
-                    raise ValueError(f"Invalid JSON format in message: {msg.text}")
+                    raise ValueError(f"Invalid JSON format in message: {text_content}")
         return None
 
     def update_memory_if_changed(self, new_memory: Memory) -> bool:
@@ -196,6 +204,7 @@ class Agent(BaseAgent):
             return True
         return False
 
+    # TODO: Refactor into separate class v.s. large if/elses here
     def execute_tool_and_persist_state(
         self, function_name: str, function_args: dict, target_letta_tool: Tool
     ) -> tuple[Any, Optional[SandboxRunResult]]:
@@ -236,6 +245,32 @@ class Agent(BaseAgent):
                 function_response = execute_composio_action(
                     action_name=action_name, args=function_args, api_key=composio_api_key, entity_id=entity_id
                 )
+            elif target_letta_tool.tool_type == ToolType.EXTERNAL_MCP:
+                # Get the server name from the tool tag
+                # TODO make a property instead?
+                server_name = target_letta_tool.tags[0].split(":")[1]
+
+                # Get the MCPClient from the server's handle
+                # TODO these don't get raised properly
+                if not self.mcp_clients:
+                    raise ValueError(f"No MCP client available to use")
+                if server_name not in self.mcp_clients:
+                    raise ValueError(f"Unknown MCP server name: {server_name}")
+                mcp_client = self.mcp_clients[server_name]
+                if not isinstance(mcp_client, BaseMCPClient):
+                    raise RuntimeError(f"Expected an MCPClient, but got: {type(mcp_client)}")
+
+                # Check that tool exists
+                available_tools = mcp_client.list_tools()
+                available_tool_names = [t.name for t in available_tools]
+                if function_name not in available_tool_names:
+                    raise ValueError(
+                        f"{function_name} is not available in MCP server {server_name}. Please check your `~/.letta/mcp_config.json` file."
+                    )
+
+                function_response, is_error = mcp_client.execute_tool(tool_name=function_name, tool_args=function_args)
+                sandbox_run_result = SandboxRunResult(status="error" if is_error else "success")
+                return function_response, sandbox_run_result
             else:
                 try:
                     # Parse the source code to extract function annotations
@@ -266,6 +301,7 @@ class Agent(BaseAgent):
             function_response = get_friendly_error_msg(
                 function_name=function_name, exception_name=type(e).__name__, exception_message=str(e)
             )
+            return function_response, SandboxRunResult(status="error")
 
         return function_response, None
 
@@ -356,19 +392,38 @@ class Agent(BaseAgent):
         for attempt in range(1, empty_response_retry_limit + 1):
             try:
                 log_telemetry(self.logger, "_get_ai_reply create start")
-                response = create(
+                # New LLM client flow
+                llm_client = LLMClient.create(
+                    agent_id=self.agent_state.id,
                     llm_config=self.agent_state.llm_config,
-                    messages=message_sequence,
-                    user_id=self.agent_state.created_by_id,
-                    functions=allowed_functions,
-                    # functions_python=self.functions_python, do we need this?
-                    function_call=function_call,
-                    first_message=first_message,
-                    force_tool_call=force_tool_call,
-                    stream=stream,
-                    stream_interface=self.interface,
                     put_inner_thoughts_first=put_inner_thoughts_first,
+                    actor_id=self.agent_state.created_by_id,
                 )
+
+                if llm_client and not stream:
+                    response = llm_client.send_llm_request(
+                        messages=message_sequence,
+                        tools=allowed_functions,
+                        tool_call=function_call,
+                        stream=stream,
+                        first_message=first_message,
+                        force_tool_call=force_tool_call,
+                    )
+                else:
+                    # Fallback to existing flow
+                    response = create(
+                        llm_config=self.agent_state.llm_config,
+                        messages=message_sequence,
+                        user_id=self.agent_state.created_by_id,
+                        functions=allowed_functions,
+                        # functions_python=self.functions_python, do we need this?
+                        function_call=function_call,
+                        first_message=first_message,
+                        force_tool_call=force_tool_call,
+                        stream=stream,
+                        stream_interface=self.interface,
+                        put_inner_thoughts_first=put_inner_thoughts_first,
+                    )
                 log_telemetry(self.logger, "_get_ai_reply create finish")
 
                 # These bottom two are retryable
@@ -632,7 +687,7 @@ class Agent(BaseAgent):
                     function_args,
                     function_response,
                     messages,
-                    [tool_return] if tool_return else None,
+                    [tool_return],
                     include_function_failed_message=True,
                 )
                 return messages, False, True  # force a heartbeat to allow agent to handle error
@@ -659,7 +714,7 @@ class Agent(BaseAgent):
                         "content": function_response,
                         "tool_call_id": tool_call_id,
                     },
-                    tool_returns=[tool_return] if tool_return else None,
+                    tool_returns=[tool_return] if sandbox_run_result else None,
                 )
             )  # extend conversation with function response
             self.interface.function_message(f"Ran {function_name}({function_args})", msg_obj=messages[-1])
@@ -909,6 +964,7 @@ class Agent(BaseAgent):
             # Log step - this must happen before messages are persisted
             step = self.step_manager.log_step(
                 actor=self.user,
+                agent_id=self.agent_state.id,
                 provider_name=self.agent_state.llm_config.model_endpoint_type,
                 model=self.agent_state.llm_config.model,
                 model_endpoint=self.agent_state.llm_config.model_endpoint,
@@ -989,7 +1045,7 @@ class Agent(BaseAgent):
                         err_msg,
                         details={
                             "num_in_context_messages": len(self.agent_state.message_ids),
-                            "in_context_messages_text": [m.text for m in in_context_messages],
+                            "in_context_messages_text": [m.content for m in in_context_messages],
                             "token_counts": token_counts,
                         },
                     )
@@ -1143,14 +1199,17 @@ class Agent(BaseAgent):
         if (
             len(in_context_messages) > 1
             and in_context_messages[1].role == MessageRole.user
-            and isinstance(in_context_messages[1].text, str)
+            and in_context_messages[1].content
+            and len(in_context_messages[1].content) == 1
+            and in_context_messages[1].content[0].type == MessageContentType.text
             # TODO remove hardcoding
-            and "The following is a summary of the previous " in in_context_messages[1].text
+            and "The following is a summary of the previous " in in_context_messages[1].content[0].text
         ):
             # Summary message exists
-            assert in_context_messages[1].text is not None
-            summary_memory = in_context_messages[1].text
-            num_tokens_summary_memory = count_tokens(in_context_messages[1].text)
+            text_content = in_context_messages[1].content[0].text
+            assert text_content is not None
+            summary_memory = text_content
+            num_tokens_summary_memory = count_tokens(text_content)
             # with a summary message, the real messages start at index 2
             num_tokens_messages = (
                 num_tokens_from_messages(messages=in_context_messages_openai[2:], model=self.model)
@@ -1174,6 +1233,7 @@ class Agent(BaseAgent):
             memory_edit_timestamp=get_utc_time(),
             previous_message_count=self.message_manager.size(actor=self.user, agent_id=self.agent_state.id),
             archival_memory_size=self.agent_manager.passage_size(actor=self.user, agent_id=self.agent_state.id),
+            recent_passages=self.agent_manager.list_passages(actor=self.user, agent_id=self.agent_state.id, ascending=False, limit=10),
         )
         num_tokens_external_memory_summary = count_tokens(external_memory_summary)
 

letta/agents/low_latency_agent.py

@@ -237,7 +237,8 @@ class LowLatencyAgent(BaseAgent):
         # TODO: This is a pretty brittle pattern established all over our code, need to get rid of this
         curr_system_message = in_context_messages[0]
         curr_memory_str = agent_state.memory.compile()
-        if curr_memory_str in curr_system_message.text:
+        curr_system_message_text = curr_system_message.content[0].text
+        if curr_memory_str in curr_system_message_text:
             # NOTE: could this cause issues if a block is removed? (substring match would still work)
             logger.debug(
                 f"Memory hasn't changed for agent id={agent_state.id} and actor=({self.actor.id}, {self.actor.name}), skipping system prompt rebuild"
@@ -251,7 +252,7 @@ class LowLatencyAgent(BaseAgent):
             in_context_memory_last_edit=memory_edit_timestamp,
         )
 
-        diff = united_diff(curr_system_message.text, new_system_message_str)
+        diff = united_diff(curr_system_message_text, new_system_message_str)
         if len(diff) > 0:
             logger.info(f"Rebuilding system with new memory...\nDiff:\n{diff}")
 

letta/client/client.py

@@ -4,7 +4,6 @@ import time
 from typing import Callable, Dict, Generator, List, Optional, Union
 
 import requests
-from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall as OpenAIToolCall
 
 import letta.utils
 from letta.constants import ADMIN_PREFIX, BASE_MEMORY_TOOLS, BASE_TOOLS, DEFAULT_HUMAN, DEFAULT_PERSONA, FUNCTION_RETURN_CHAR_LIMIT
@@ -29,7 +28,7 @@ from letta.schemas.letta_request import LettaRequest, LettaStreamingRequest
 from letta.schemas.letta_response import LettaResponse, LettaStreamingResponse
 from letta.schemas.llm_config import LLMConfig
 from letta.schemas.memory import ArchivalMemorySummary, ChatMemory, CreateArchivalMemory, Memory, RecallMemorySummary
-from letta.schemas.message import Message, MessageCreate, MessageUpdate
+from letta.schemas.message import Message, MessageCreate
 from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.organization import Organization
 from letta.schemas.passage import Passage
@@ -640,30 +639,6 @@ class RESTClient(AbstractClient):
         # refresh and return agent
         return self.get_agent(agent_state.id)
 
-    def update_message(
-        self,
-        agent_id: str,
-        message_id: str,
-        role: Optional[MessageRole] = None,
-        text: Optional[str] = None,
-        name: Optional[str] = None,
-        tool_calls: Optional[List[OpenAIToolCall]] = None,
-        tool_call_id: Optional[str] = None,
-    ) -> Message:
-        request = MessageUpdate(
-            role=role,
-            content=text,
-            name=name,
-            tool_calls=tool_calls,
-            tool_call_id=tool_call_id,
-        )
-        response = requests.patch(
-            f"{self.base_url}/{self.api_prefix}/agents/{agent_id}/messages/{message_id}", json=request.model_dump(), headers=self.headers
-        )
-        if response.status_code != 200:
-            raise ValueError(f"Failed to update message: {response.text}")
-        return Message(**response.json())
-
     def update_agent(
         self,
         agent_id: str,
@@ -2436,30 +2411,6 @@ class LocalClient(AbstractClient):
         # TODO: get full agent state
         return self.server.agent_manager.get_agent_by_id(agent_state.id, actor=self.user)
 
-    def update_message(
-        self,
-        agent_id: str,
-        message_id: str,
-        role: Optional[MessageRole] = None,
-        text: Optional[str] = None,
-        name: Optional[str] = None,
-        tool_calls: Optional[List[OpenAIToolCall]] = None,
-        tool_call_id: Optional[str] = None,
-    ) -> Message:
-        message = self.server.update_agent_message(
-            agent_id=agent_id,
-            message_id=message_id,
-            request=MessageUpdate(
-                role=role,
-                content=text,
-                name=name,
-                tool_calls=tool_calls,
-                tool_call_id=tool_call_id,
-            ),
-            actor=self.user,
-        )
-        return message
-
     def update_agent(
         self,
         agent_id: str,

letta/constants.py

@@ -11,6 +11,9 @@ OPENAI_API_PREFIX = "/openai"
 COMPOSIO_ENTITY_ENV_VAR_KEY = "COMPOSIO_ENTITY"
 COMPOSIO_TOOL_TAG_NAME = "composio"
 
+MCP_CONFIG_NAME = "mcp_config.json"
+MCP_TOOL_TAG_NAME_PREFIX = "mcp"  # full format, mcp:server_name
+
 LETTA_CORE_TOOL_MODULE_NAME = "letta.functions.function_sets.base"
 LETTA_MULTI_AGENT_TOOL_MODULE_NAME = "letta.functions.function_sets.multi_agent"
 
@@ -50,7 +53,7 @@ BASE_TOOLS = ["send_message", "conversation_search", "archival_memory_insert", "
 # Base memory tools CAN be edited, and are added by default by the server
 BASE_MEMORY_TOOLS = ["core_memory_append", "core_memory_replace"]
 # Multi agent tools
-MULTI_AGENT_TOOLS = ["send_message_to_agent_and_wait_for_reply", "send_message_to_agents_matching_all_tags", "send_message_to_agent_async"]
+MULTI_AGENT_TOOLS = ["send_message_to_agent_and_wait_for_reply", "send_message_to_agents_matching_tags", "send_message_to_agent_async"]
 # Set of all built-in Letta tools
 LETTA_TOOL_SET = set(BASE_TOOLS + BASE_MEMORY_TOOLS + MULTI_AGENT_TOOLS)
 

letta/functions/function_sets/base.py

@@ -56,7 +56,7 @@ def conversation_search(self: "Agent", query: str, page: Optional[int] = 0) -> O
         results_str = f"No results found."
     else:
         results_pref = f"Showing {len(messages)} of {total} results (page {page}/{num_pages}):"
-        results_formatted = [message.text for message in messages]
+        results_formatted = [message.content[0].text for message in messages]
         results_str = f"{results_pref} {json_dumps(results_formatted)}"
     return results_str
 

letta/functions/function_sets/multi_agent.py

@@ -2,7 +2,7 @@ import asyncio
 from typing import TYPE_CHECKING, List
 
 from letta.functions.helpers import (
-    _send_message_to_agents_matching_all_tags_async,
+    _send_message_to_agents_matching_tags_async,
     execute_send_message_to_agent,
     fire_and_forget_send_to_agent,
 )
@@ -70,18 +70,19 @@ def send_message_to_agent_async(self: "Agent", message: str, other_agent_id: str
     return "Successfully sent message"
 
 
-def send_message_to_agents_matching_all_tags(self: "Agent", message: str, tags: List[str]) -> List[str]:
+def send_message_to_agents_matching_tags(self: "Agent", message: str, match_all: List[str], match_some: List[str]) -> List[str]:
     """
-    Sends a message to all agents within the same organization that match all of the specified tags. Messages are dispatched in parallel for improved performance, with retries to handle transient issues and timeouts to ensure responsiveness. This function enforces a limit of 100 agents and does not support pagination (cursor-based queries). Each agent must match all specified tags (`match_all_tags=True`) to be included.
+    Sends a message to all agents within the same organization that match the specified tag criteria. Agents must possess *all* of the tags in `match_all` and *at least one* of the tags in `match_some` to receive the message.
 
     Args:
         message (str): The content of the message to be sent to each matching agent.
-        tags (List[str]): A list of tags that an agent must possess to receive the message.
+        match_all (List[str]): A list of tags that an agent must possess to receive the message.
+        match_some (List[str]): A list of tags where an agent must have at least one to qualify.
 
     Returns:
-        List[str]: A list of responses from the agents that matched all tags. Each
-        response corresponds to a single agent. Agents that do not respond will not
-        have an entry in the returned list.
+        List[str]: A list of responses from the agents that matched the filtering criteria. Each
+        response corresponds to a single agent. Agents that do not respond will not have an entry
+        in the returned list.
     """
 
-    return asyncio.run(_send_message_to_agents_matching_all_tags_async(self, message, tags))
+    return asyncio.run(_send_message_to_agents_matching_tags_async(self, message, match_all, match_some))

letta/functions/helpers.py

@@ -48,6 +48,20 @@ def generate_composio_action_from_func_name(func_name: str) -> str:
     return func_name.upper()
 
 
+# TODO needed?
+def generate_mcp_tool_wrapper(mcp_tool_name: str) -> tuple[str, str]:
+
+    wrapper_function_str = f"""\
+def {mcp_tool_name}(**kwargs):
+    raise RuntimeError("Something went wrong - we should never be using the persisted source code for MCP. Please reach out to Letta team")
+"""
+
+    # Compile safety check
+    assert_code_gen_compilable(wrapper_function_str.strip())
+
+    return mcp_tool_name, wrapper_function_str.strip()
+
+
 def generate_composio_tool_wrapper(action_name: str) -> tuple[str, str]:
     # Generate func name
     func_name = generate_func_name_from_composio_action(action_name)
@@ -518,8 +532,16 @@ def fire_and_forget_send_to_agent(
         run_in_background_thread(background_task())
 
 
-async def _send_message_to_agents_matching_all_tags_async(sender_agent: "Agent", message: str, tags: List[str]) -> List[str]:
-    log_telemetry(sender_agent.logger, "_send_message_to_agents_matching_all_tags_async start", message=message, tags=tags)
+async def _send_message_to_agents_matching_tags_async(
+    sender_agent: "Agent", message: str, match_all: List[str], match_some: List[str]
+) -> List[str]:
+    log_telemetry(
+        sender_agent.logger,
+        "_send_message_to_agents_matching_tags_async start",
+        message=message,
+        match_all=match_all,
+        match_some=match_some,
+    )
     server = get_letta_server()
 
     augmented_message = (
@@ -529,9 +551,22 @@ async def _send_message_to_agents_matching_all_tags_async(sender_agent: "Agent",
     )
 
     # Retrieve up to 100 matching agents
-    log_telemetry(sender_agent.logger, "_send_message_to_agents_matching_all_tags_async listing agents start", message=message, tags=tags)
-    matching_agents = server.agent_manager.list_agents(actor=sender_agent.user, tags=tags, match_all_tags=True, limit=100)
-    log_telemetry(sender_agent.logger, "_send_message_to_agents_matching_all_tags_async  listing agents finish", message=message, tags=tags)
+    log_telemetry(
+        sender_agent.logger,
+        "_send_message_to_agents_matching_tags_async listing agents start",
+        message=message,
+        match_all=match_all,
+        match_some=match_some,
+    )
+    matching_agents = server.agent_manager.list_agents_matching_tags(actor=sender_agent.user, match_all=match_all, match_some=match_some)
+
+    log_telemetry(
+        sender_agent.logger,
+        "_send_message_to_agents_matching_tags_async  listing agents finish",
+        message=message,
+        match_all=match_all,
+        match_some=match_some,
+    )
 
     # Create a system message
     messages = [MessageCreate(role=MessageRole.system, content=augmented_message, name=sender_agent.agent_state.name)]
@@ -559,7 +594,13 @@ async def _send_message_to_agents_matching_all_tags_async(sender_agent: "Agent",
         else:
             final.append(r)
 
-    log_telemetry(sender_agent.logger, "_send_message_to_agents_matching_all_tags_async finish", message=message, tags=tags)
+    log_telemetry(
+        sender_agent.logger,
+        "_send_message_to_agents_matching_tags_async finish",
+        message=message,
+        match_all=match_all,
+        match_some=match_some,
+    )
     return final
 
 

letta/functions/schema_generator.py

@@ -6,6 +6,8 @@ from composio.client.collections import ActionParametersModel
 from docstring_parser import parse
 from pydantic import BaseModel
 
+from letta.helpers.mcp_helpers import MCPTool
+
 
 def is_optional(annotation):
     # Check if the annotation is a Union
@@ -447,6 +449,51 @@ def generate_schema_from_args_schema_v2(
     return function_call_json
 
 
+def generate_tool_schema_for_mcp(
+    mcp_tool: MCPTool,
+    append_heartbeat: bool = True,
+    strict: bool = False,
+) -> Dict[str, Any]:
+
+    # MCP tool.inputSchema is a JSON schema
+    # https://github.com/modelcontextprotocol/python-sdk/blob/775f87981300660ee957b63c2a14b448ab9c3675/src/mcp/types.py#L678
+    parameters_schema = mcp_tool.inputSchema
+    name = mcp_tool.name
+    description = mcp_tool.description
+
+    assert "type" in parameters_schema
+    assert "required" in parameters_schema
+    assert "properties" in parameters_schema
+
+    # Add the optional heartbeat parameter
+    if append_heartbeat:
+        parameters_schema["properties"]["request_heartbeat"] = {
+            "type": "boolean",
+            "description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
+        }
+        parameters_schema["required"].append("request_heartbeat")
+
+    # Return the final schema
+    if strict:
+        # https://platform.openai.com/docs/guides/function-calling#strict-mode
+
+        # Add additionalProperties: False
+        parameters_schema["additionalProperties"] = False
+
+        return {
+            "strict": True,  # NOTE
+            "name": name,
+            "description": description,
+            "parameters": parameters_schema,
+        }
+    else:
+        return {
+            "name": name,
+            "description": description,
+            "parameters": parameters_schema,
+        }
+
+
 def generate_tool_schema_for_composio(
     parameters_model: ActionParametersModel,
     name: str,

letta/helpers/mcp_helpers.py

@@ -0,0 +1,108 @@
+import asyncio
+from enum import Enum
+from typing import List, Optional, Tuple
+
+from mcp import ClientSession, StdioServerParameters, Tool
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from pydantic import BaseModel, Field
+
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
+
+class MCPTool(Tool):
+    """A simple wrapper around MCP's tool definition (to avoid conflict with our own)"""
+
+
+class MCPServerType(str, Enum):
+    SSE = "sse"
+    LOCAL = "local"
+
+
+class BaseServerConfig(BaseModel):
+    server_name: str = Field(..., description="The name of the server")
+    type: MCPServerType
+
+
+class SSEServerConfig(BaseServerConfig):
+    type: MCPServerType = MCPServerType.SSE
+    server_url: str = Field(..., description="The URL of the server (MCP SSE client will connect to this URL)")
+
+
+class LocalServerConfig(BaseServerConfig):
+    type: MCPServerType = MCPServerType.LOCAL
+    command: str = Field(..., description="The command to run (MCP 'local' client will run this command)")
+    args: List[str] = Field(..., description="The arguments to pass to the command")
+
+
+class BaseMCPClient:
+    def __init__(self):
+        self.session: Optional[ClientSession] = None
+        self.stdio = None
+        self.write = None
+        self.initialized = False
+        self.loop = asyncio.new_event_loop()
+        self.cleanup_funcs = []
+
+    def connect_to_server(self, server_config: BaseServerConfig):
+        asyncio.set_event_loop(self.loop)
+        self._initialize_connection(server_config)
+        self.loop.run_until_complete(self.session.initialize())
+        self.initialized = True
+
+    def _initialize_connection(self, server_config: BaseServerConfig):
+        raise NotImplementedError("Subclasses must implement _initialize_connection")
+
+    def list_tools(self) -> List[Tool]:
+        self._check_initialized()
+        response = self.loop.run_until_complete(self.session.list_tools())
+        return response.tools
+
+    def execute_tool(self, tool_name: str, tool_args: dict) -> Tuple[str, bool]:
+        self._check_initialized()
+        result = self.loop.run_until_complete(self.session.call_tool(tool_name, tool_args))
+        return str(result.content), result.isError
+
+    def _check_initialized(self):
+        if not self.initialized:
+            logger.error("MCPClient has not been initialized")
+            raise RuntimeError("MCPClient has not been initialized")
+
+    def cleanup(self):
+        try:
+            for cleanup_func in self.cleanup_funcs:
+                cleanup_func()
+            self.initialized = False
+            if not self.loop.is_closed():
+                self.loop.close()
+        except Exception as e:
+            logger.warning(e)
+        finally:
+            logger.info("Cleaned up MCP clients on shutdown.")
+
+
+class LocalMCPClient(BaseMCPClient):
+    def _initialize_connection(self, server_config: LocalServerConfig):
+        server_params = StdioServerParameters(command=server_config.command, args=server_config.args)
+        stdio_cm = stdio_client(server_params)
+        stdio_transport = self.loop.run_until_complete(stdio_cm.__aenter__())
+        self.stdio, self.write = stdio_transport
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(stdio_cm.__aexit__(None, None, None)))
+
+        session_cm = ClientSession(self.stdio, self.write)
+        self.session = self.loop.run_until_complete(session_cm.__aenter__())
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))
+
+
+class SSEMCPClient(BaseMCPClient):
+    def _initialize_connection(self, server_config: SSEServerConfig):
+        sse_cm = sse_client(url=server_config.server_url)
+        sse_transport = self.loop.run_until_complete(sse_cm.__aenter__())
+        self.stdio, self.write = sse_transport
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(sse_cm.__aexit__(None, None, None)))
+
+        session_cm = ClientSession(self.stdio, self.write)
+        self.session = self.loop.run_until_complete(session_cm.__aenter__())
+        self.cleanup_funcs.append(lambda: self.loop.run_until_complete(session_cm.__aexit__(None, None, None)))

letta/llm_api/cohere.py

@@ -321,7 +321,7 @@ def cohere_chat_completions_request(
     # See: https://docs.cohere.com/reference/chat
     # The chat_history parameter should not be used for SYSTEM messages in most cases. Instead, to add a SYSTEM role message at the beginning of a conversation, the preamble parameter should be used.
     assert msg_objs[0].role == "system", msg_objs[0]
-    preamble = msg_objs[0].text
+    preamble = msg_objs[0].content[0].text
 
     # data["messages"] = [m.to_cohere_dict() for m in msg_objs[1:]]
     data["messages"] = []