letta-nightly 0.6.38.dev20250312104155__py3-none-any.whl → 0.6.39.dev20250313104142__py3-none-any.whl

letta/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.38"
+__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,6 +26,7 @@ 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
@@ -38,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
@@ -92,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
@@ -149,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:
@@ -197,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]]:
@@ -237,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
@@ -267,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
 
@@ -1010,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,
                         },
                     )
@@ -1164,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)

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/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"
 

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/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)

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"] = []

letta/llm_api/helpers.py

@@ -86,9 +86,8 @@ def convert_to_structured_output(openai_function: dict, allow_optional: bool = F
     # but if "type" is "object" we expected "properties", where each property has details
     # and if "type" is "array" we expect "items": <type>
     for param, details in openai_function["parameters"]["properties"].items():
-
         param_type = details["type"]
-        description = details["description"]
+        description = details.get("description", "")
 
         if param_type == "object":
             if "properties" not in details:

letta/llm_api/llm_api_tools.py

@@ -596,7 +596,6 @@ def create(
             messages[0].content[
                 0
             ].text += f'Select best function to call simply by responding with a single json block with the keys "function" and "params". Use double quotes around the arguments.'
-
         return get_chat_completion(
             model=llm_config.model,
             messages=messages,

letta/local_llm/utils.py

@@ -114,26 +114,36 @@ def num_tokens_from_functions(functions: List[dict], model: str = "gpt-4"):
                     function_tokens += len(encoding.encode(propertiesKey))
                     v = parameters["properties"][propertiesKey]
                     for field in v:
-                        if field == "type":
-                            function_tokens += 2
-                            function_tokens += len(encoding.encode(v["type"]))
-                        elif field == "description":
-                            function_tokens += 2
-                            function_tokens += len(encoding.encode(v["description"]))
-                        elif field == "enum":
-                            function_tokens -= 3
-                            for o in v["enum"]:
-                                function_tokens += 3
-                                function_tokens += len(encoding.encode(o))
-                        elif field == "items":
-                            function_tokens += 2
-                            if isinstance(v["items"], dict) and "type" in v["items"]:
-                                function_tokens += len(encoding.encode(v["items"]["type"]))
-                        elif field == "default":
-                            function_tokens += 2
-                            function_tokens += len(encoding.encode(str(v["default"])))
-                        else:
-                            logger.warning(f"num_tokens_from_functions: Unsupported field {field} in function {function}")
+                        try:
+                            if field == "type":
+                                function_tokens += 2
+                                function_tokens += len(encoding.encode(v["type"]))
+                            elif field == "description":
+                                function_tokens += 2
+                                function_tokens += len(encoding.encode(v["description"]))
+                            elif field == "enum":
+                                function_tokens -= 3
+                                for o in v["enum"]:
+                                    function_tokens += 3
+                                    function_tokens += len(encoding.encode(o))
+                            elif field == "items":
+                                function_tokens += 2
+                                if isinstance(v["items"], dict) and "type" in v["items"]:
+                                    function_tokens += len(encoding.encode(v["items"]["type"]))
+                            elif field == "default":
+                                function_tokens += 2
+                                function_tokens += len(encoding.encode(str(v["default"])))
+                            elif field == "title":
+                                # TODO: Is this right? For MCP
+                                continue
+                            else:
+                                # TODO: Handle nesting here properly
+                                # Disable this for now for MCP
+                                continue
+                                # logger.warning(f"num_tokens_from_functions: Unsupported field {field} in function {function}")
+                        except:
+                            logger.error(f"Failed to encode field {field} with value {v}")
+                            raise
                 function_tokens += 11
 
         num_tokens += function_tokens

letta/log.py

@@ -54,7 +54,7 @@ DEVELOPMENT_LOGGING = {
             "propagate": True,  # Let logs bubble up to root
         },
         "uvicorn": {
-            "level": "INFO",
+            "level": "CRITICAL",
             "handlers": ["console"],
             "propagate": True,
         },

letta/memory.py

@@ -36,7 +36,7 @@ def get_memory_functions(cls: Memory) -> Dict[str, Callable]:
 
 def _format_summary_history(message_history: List[Message]):
     # TODO use existing prompt formatters for this (eg ChatML)
-    return "\n".join([f"{m.role}: {m.text}" for m in message_history])
+    return "\n".join([f"{m.role}: {m.content[0].text}" for m in message_history])
 
 
 def summarize_messages(

letta/orm/__init__.py

@@ -5,6 +5,7 @@ from letta.orm.block import Block
 from letta.orm.blocks_agents import BlocksAgents
 from letta.orm.file import FileMetadata
 from letta.orm.identities_agents import IdentitiesAgents
+from letta.orm.identities_blocks import IdentitiesBlocks
 from letta.orm.identity import Identity
 from letta.orm.job import Job
 from letta.orm.job_messages import JobMessage

letta/orm/block.py

@@ -12,6 +12,7 @@ from letta.schemas.block import Human, Persona
 
 if TYPE_CHECKING:
     from letta.orm import Organization
+    from letta.orm.identity import Identity
 
 
 class Block(OrganizationMixin, SqlalchemyBase):
@@ -47,6 +48,13 @@ class Block(OrganizationMixin, SqlalchemyBase):
         back_populates="core_memory",
         doc="Agents associated with this block.",
     )
+    identities: Mapped[List["Identity"]] = relationship(
+        "Identity",
+        secondary="identities_blocks",
+        lazy="selectin",
+        back_populates="blocks",
+        passive_deletes=True,
+    )
 
     def to_pydantic(self) -> Type:
         match self.label:

letta/orm/enums.py

@@ -8,6 +8,8 @@ class ToolType(str, Enum):
     LETTA_MULTI_AGENT_CORE = "letta_multi_agent_core"
     EXTERNAL_COMPOSIO = "external_composio"
     EXTERNAL_LANGCHAIN = "external_langchain"
+    # TODO is "external" the right name here? Since as of now, MCP is local / doesn't support remote?
+    EXTERNAL_MCP = "external_mcp"
 
 
 class JobType(str, Enum):

letta/orm/identities_blocks.py

@@ -0,0 +1,13 @@
+from sqlalchemy import ForeignKey, String
+from sqlalchemy.orm import Mapped, mapped_column
+
+from letta.orm.base import Base
+
+
+class IdentitiesBlocks(Base):
+    """Identities may have one or many blocks associated with them."""
+
+    __tablename__ = "identities_blocks"
+
+    identity_id: Mapped[str] = mapped_column(String, ForeignKey("identities.id", ondelete="CASCADE"), primary_key=True)
+    block_id: Mapped[str] = mapped_column(String, ForeignKey("block.id", ondelete="CASCADE"), primary_key=True)

letta/orm/identity.py

@@ -40,12 +40,20 @@ class Identity(SqlalchemyBase, OrganizationMixin):
     agents: Mapped[List["Agent"]] = relationship(
         "Agent", secondary="identities_agents", lazy="selectin", passive_deletes=True, back_populates="identities"
     )
+    blocks: Mapped[List["Block"]] = relationship(
+        "Block", secondary="identities_blocks", lazy="selectin", passive_deletes=True, back_populates="identities"
+    )
 
     @property
     def agent_ids(self) -> List[str]:
         """Get just the agent IDs without loading the full agent objects"""
         return [agent.id for agent in self.agents]
 
+    @property
+    def block_ids(self) -> List[str]:
+        """Get just the block IDs without loading the full agent objects"""
+        return [block.id for block in self.blocks]
+
     def to_pydantic(self) -> PydanticIdentity:
         state = {
             "id": self.id,
@@ -54,6 +62,7 @@ class Identity(SqlalchemyBase, OrganizationMixin):
             "identity_type": self.identity_type,
             "project_id": self.project_id,
             "agent_ids": self.agent_ids,
+            "block_ids": self.block_ids,
             "organization_id": self.organization_id,
             "properties": self.properties,
         }

letta/orm/sqlalchemy_base.py

@@ -69,7 +69,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         join_model: Optional[Base] = None,
         join_conditions: Optional[Union[Tuple, List]] = None,
         identifier_keys: Optional[List[str]] = None,
-        identifier_id: Optional[str] = None,
+        identity_id: Optional[str] = None,
         **kwargs,
     ) -> List["SqlalchemyBase"]:
         """
@@ -148,9 +148,9 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             if identifier_keys and hasattr(cls, "identities"):
                 query = query.join(cls.identities).filter(cls.identities.property.mapper.class_.identifier_key.in_(identifier_keys))
 
-            # given the identifier_id, we can find within the agents table any agents that have the identifier_id in their identity_ids
-            if identifier_id and hasattr(cls, "identities"):
-                query = query.join(cls.identities).filter(cls.identities.property.mapper.class_.id == identifier_id)
+            # given the identity_id, we can find within the agents table any agents that have the identity_id in their identity_ids
+            if identity_id and hasattr(cls, "identities"):
+                query = query.join(cls.identities).filter(cls.identities.property.mapper.class_.id == identity_id)
 
             # Apply filtering logic from kwargs
             for key, value in kwargs.items():

letta/schemas/identity.py

@@ -46,6 +46,7 @@ class Identity(IdentityBase):
     identity_type: IdentityType = Field(..., description="The type of the identity.")
     project_id: Optional[str] = Field(None, description="The project id of the identity, if applicable.")
     agent_ids: List[str] = Field(..., description="The IDs of the agents associated with the identity.")
+    block_ids: List[str] = Field(..., description="The IDs of the blocks associated with the identity.")
     organization_id: Optional[str] = Field(None, description="The organization id of the user")
     properties: List[IdentityProperty] = Field(default_factory=list, description="List of properties associated with the identity")
 
@@ -56,6 +57,7 @@ class IdentityCreate(LettaBase):
     identity_type: IdentityType = Field(..., description="The type of the identity.")
     project_id: Optional[str] = Field(None, description="The project id of the identity, if applicable.")
     agent_ids: Optional[List[str]] = Field(None, description="The agent ids that are associated with the identity.")
+    block_ids: Optional[List[str]] = Field(None, description="The IDs of the blocks associated with the identity.")
     properties: Optional[List[IdentityProperty]] = Field(None, description="List of properties associated with the identity.")
 
 
@@ -64,4 +66,5 @@ class IdentityUpdate(LettaBase):
     name: Optional[str] = Field(None, description="The name of the identity.")
     identity_type: Optional[IdentityType] = Field(None, description="The type of the identity.")
     agent_ids: Optional[List[str]] = Field(None, description="The agent ids that are associated with the identity.")
+    block_ids: Optional[List[str]] = Field(None, description="The IDs of the blocks associated with the identity.")
     properties: Optional[List[IdentityProperty]] = Field(None, description="List of properties associated with the identity.")