letta-nightly 0.12.1.dev20251023104211__py3-none-any.whl → 0.13.0.dev20251024223017__py3-none-any.whl

letta/services/helpers/agent_manager_helper.py

@@ -2,6 +2,10 @@ import uuid
 from datetime import datetime
 from typing import List, Literal, Optional, Set
 
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
 import numpy as np
 from sqlalchemy import Select, and_, asc, desc, func, literal, nulls_last, or_, select, union_all
 from sqlalchemy.orm import noload
@@ -38,7 +42,7 @@ from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import AgentType, MessageRole
 from letta.schemas.letta_message_content import TextContent
 from letta.schemas.memory import Memory
-from letta.schemas.message import Message, MessageCreate
+from letta.schemas.message import Message, MessageCreate, ToolReturn
 from letta.schemas.tool_rule import ToolRule
 from letta.schemas.user import User
 from letta.settings import DatabaseChoice, settings
@@ -256,6 +260,7 @@ def compile_system_message(
     tool_rules_solver: Optional[ToolRulesSolver] = None,
     sources: Optional[List] = None,
     max_files_open: Optional[int] = None,
+    llm_config: Optional[object] = None,
 ) -> str:
     """Prepare the final/full system message that will be fed into the LLM API
 
@@ -289,7 +294,7 @@ def compile_system_message(
         )
 
         memory_with_sources = in_context_memory.compile(
-            tool_usage_rules=tool_constraint_block, sources=sources, max_files_open=max_files_open
+            tool_usage_rules=tool_constraint_block, sources=sources, max_files_open=max_files_open, llm_config=llm_config
         )
         full_memory_string = memory_with_sources + "\n\n" + memory_metadata_string
 
@@ -303,7 +308,7 @@ def compile_system_message(
         if append_icm_if_missing:
             if memory_variable_string not in system_prompt:
                 # In this case, append it to the end to make sure memory is still injected
-                # warnings.warn(f"{IN_CONTEXT_MEMORY_KEYWORD} variable was missing from system prompt, appending instead")
+                # logger.warning(f"{IN_CONTEXT_MEMORY_KEYWORD} variable was missing from system prompt, appending instead")
                 system_prompt += "\n\n" + memory_variable_string
 
         # render the variables using the built-in templater
@@ -536,6 +541,13 @@ def package_initial_message_sequence(
                     agent_id=agent_id,
                     model=model,
                     tool_call_id=tool_call_id,
+                    tool_returns=[
+                        ToolReturn(
+                            tool_call_id=tool_call_id,
+                            status="success",
+                            func_response=function_response,
+                        )
+                    ],
                 )
             )
         else:
@@ -767,24 +779,66 @@ def _apply_filters(
     return query
 
 
-def _apply_relationship_filters(query, include_relationships: Optional[List[str]] = None):
-    if include_relationships is None:
+def _apply_relationship_filters(
+    query,
+    include_relationships: Optional[List[str]] = None,
+    include: Optional[List[str]] = None,
+):
+    # legacy include_relationships
+    if include_relationships is None and not include:
         return query
 
-    if "memory" not in include_relationships:
-        query = query.options(noload(AgentModel.core_memory), noload(AgentModel.file_agents))
-    if "identity_ids" not in include_relationships:
-        query = query.options(noload(AgentModel.identities))
+    column_names = get_column_names_from_includes_params(include_relationships, include)
 
-    relationships = ["tool_exec_environment_variables", "tools", "sources", "tags", "multi_agent_group"]
+    relationships = [
+        "core_memory",
+        "file_agents",
+        "identities",
+        "tool_exec_environment_variables",
+        "tools",
+        "sources",
+        "tags",
+        "multi_agent_group",
+    ]
 
     for rel in relationships:
-        if rel not in include_relationships:
+        if rel not in column_names:
             query = query.options(noload(getattr(AgentModel, rel)))
 
     return query
 
 
+def get_column_names_from_includes_params(
+    include_relationships: Optional[List[str]] = None, includes: Optional[List[str]] = None
+) -> Set[str]:
+    include_mapping = {
+        "agent.blocks": ["core_memory"],
+        "agent.identities": ["identities"],
+        "agent.managed_group": ["multi_agent_group"],
+        "agent.secrets": ["tool_exec_environment_variables"],
+        "agent.sources": ["sources"],
+        "agent.tags": ["tags"],
+        "agent.tools": ["tools"],
+        # legacy
+        "memory": ["core_memory", "file_agents"],
+        "identity_ids": ["identities"],
+        "multi_agent_group": ["multi_agent_group"],
+        "tool_exec_environment_variables": ["tool_exec_environment_variables"],
+        "secrets": ["tool_exec_environment_variables"],
+        "sources": ["sources"],
+        "tags": ["tags"],
+        "tools": ["tools"],
+    }
+    column_names = set()
+    if includes:
+        for include in includes:
+            column_names.update(include_mapping.get(include, []))
+    else:
+        for include_relationship in include_relationships:
+            column_names.update(include_mapping.get(include_relationship, []))
+    return column_names
+
+
 async def build_passage_query(
     actor: User,
     agent_id: Optional[str] = None,

letta/services/identity_manager.py

@@ -12,6 +12,7 @@ from letta.orm.identity import Identity as IdentityModel
 from letta.otel.tracing import trace_method
 from letta.schemas.agent import AgentState
 from letta.schemas.block import Block
+from letta.schemas.enums import PrimitiveType
 from letta.schemas.identity import (
     Identity as PydanticIdentity,
     IdentityCreate,
@@ -24,6 +25,7 @@ from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.settings import DatabaseChoice, settings
 from letta.utils import enforce_types
+from letta.validators import raise_on_invalid_id
 
 
 class IdentityManager:
@@ -40,7 +42,13 @@ class IdentityManager:
         limit: Optional[int] = 50,
         ascending: bool = False,
         actor: PydanticUser = None,
-    ) -> list[PydanticIdentity]:
+    ) -> tuple[list[PydanticIdentity], Optional[str], bool]:
+        """
+        List identities with pagination metadata.
+
+        Returns:
+            Tuple of (identities, next_cursor, has_more)
+        """
         async with db_registry.async_session() as session:
             filters = {"organization_id": actor.organization_id}
             if project_id:
@@ -49,19 +57,34 @@ class IdentityManager:
                 filters["identifier_key"] = identifier_key
             if identity_type:
                 filters["identity_type"] = identity_type
+
+            # Request one more than limit to check if there are more pages
+            query_limit = limit + 1 if limit else None
+
             identities = await IdentityModel.list_async(
                 db_session=session,
                 query_text=name,
                 before=before,
                 after=after,
-                limit=limit,
+                limit=query_limit,
                 ascending=ascending,
                 **filters,
             )
-            return [identity.to_pydantic() for identity in identities]
+
+            # Check if we got more records than requested (meaning there are more pages)
+            has_more = len(identities) > limit if limit else False
+            if has_more:
+                # Trim back to the requested limit
+                identities = identities[:limit]
+
+            # Get cursor for next page (ID of last item in current page)
+            next_cursor = identities[-1].id if identities else None
+
+            return [identity.to_pydantic() for identity in identities], next_cursor, has_more
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def get_identity_async(self, identity_id: str, actor: PydanticUser) -> PydanticIdentity:
         async with db_registry.async_session() as session:
             identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
@@ -143,6 +166,7 @@ class IdentityManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def update_identity_async(
         self, identity_id: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False
     ) -> PydanticIdentity:
@@ -206,6 +230,7 @@ class IdentityManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def upsert_identity_properties_async(
         self, identity_id: str, properties: List[IdentityProperty], actor: PydanticUser
     ) -> PydanticIdentity:
@@ -223,6 +248,7 @@ class IdentityManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def delete_identity_async(self, identity_id: str, actor: PydanticUser) -> None:
         async with db_registry.async_session() as session:
             identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
@@ -280,6 +306,7 @@ class IdentityManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def list_agents_for_identity_async(
         self,
         identity_id: str,
@@ -287,6 +314,7 @@ class IdentityManager:
         after: Optional[str] = None,
         limit: Optional[int] = 50,
         ascending: bool = False,
+        include: List[str] = [],
         actor: PydanticUser = None,
     ) -> List[AgentState]:
         """
@@ -307,10 +335,11 @@ class IdentityManager:
                 ascending=ascending,
                 identity_id=identity.id,
             )
-            return await asyncio.gather(*[agent.to_pydantic_async() for agent in agents])
+            return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=[], include=include) for agent in agents])
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
     async def list_blocks_for_identity_async(
         self,
         identity_id: str,
@@ -339,3 +368,75 @@ class IdentityManager:
                 identity_id=identity.id,
             )
             return [block.to_pydantic() for block in blocks]
+
+    @enforce_types
+    @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    async def attach_agent_async(self, identity_id: str, agent_id: str, actor: PydanticUser) -> None:
+        """
+        Attach an agent to an identity.
+        """
+        async with db_registry.async_session() as session:
+            identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
+
+            agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
+
+            # Add agent to identity if not already attached
+            if agent not in identity.agents:
+                identity.agents.append(agent)
+                await identity.update_async(db_session=session, actor=actor)
+
+    @enforce_types
+    @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
+    @raise_on_invalid_id(param_name="agent_id", expected_prefix=PrimitiveType.AGENT)
+    async def detach_agent_async(self, identity_id: str, agent_id: str, actor: PydanticUser) -> None:
+        """
+        Detach an agent from an identity.
+        """
+        async with db_registry.async_session() as session:
+            identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
+
+            agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
+
+            # Remove agent from identity if attached
+            if agent in identity.agents:
+                identity.agents.remove(agent)
+                await identity.update_async(db_session=session, actor=actor)
+
+    @enforce_types
+    @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
+    @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK)
+    async def attach_block_async(self, identity_id: str, block_id: str, actor: PydanticUser) -> None:
+        """
+        Attach a block to an identity.
+        """
+        async with db_registry.async_session() as session:
+            identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
+
+            block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
+
+            # Add block to identity if not already attached
+            if block not in identity.blocks:
+                identity.blocks.append(block)
+                await identity.update_async(db_session=session, actor=actor)
+
+    @enforce_types
+    @trace_method
+    @raise_on_invalid_id(param_name="identity_id", expected_prefix=PrimitiveType.IDENTITY)
+    @raise_on_invalid_id(param_name="block_id", expected_prefix=PrimitiveType.BLOCK)
+    async def detach_block_async(self, identity_id: str, block_id: str, actor: PydanticUser) -> None:
+        """
+        Detach a block from an identity.
+        """
+        async with db_registry.async_session() as session:
+            identity = await IdentityModel.read_async(db_session=session, identifier=identity_id, actor=actor)
+
+            block = await BlockModel.read_async(db_session=session, identifier=block_id, actor=actor)
+
+            # Remove block from identity if attached
+            if block in identity.blocks:
+                identity.blocks.remove(block)
+                await identity.update_async(db_session=session, actor=actor)

letta/services/job_manager.py

@@ -14,7 +14,7 @@ from letta.orm.message import Message as MessageModel
 from letta.orm.sqlalchemy_base import AccessType
 from letta.orm.step import Step, Step as StepModel
 from letta.otel.tracing import log_event, trace_method
-from letta.schemas.enums import JobStatus, JobType, MessageRole
+from letta.schemas.enums import JobStatus, JobType, MessageRole, PrimitiveType
 from letta.schemas.job import BatchJob as PydanticBatchJob, Job as PydanticJob, JobUpdate, LettaRequestConfig
 from letta.schemas.letta_message import LettaMessage
 from letta.schemas.letta_stop_reason import StopReasonType
@@ -26,6 +26,7 @@ from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.services.helpers.agent_manager_helper import validate_agent_exists_async
 from letta.utils import enforce_types
+from letta.validators import raise_on_invalid_id
 
 logger = get_logger(__name__)
 
@@ -70,6 +71,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def update_job_by_id_async(
         self, job_id: str, job_update: JobUpdate, actor: PydanticUser, safe_update: bool = False
     ) -> PydanticJob:
@@ -147,6 +149,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def safe_update_job_status_async(
         self,
         job_id: str,
@@ -187,6 +190,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def get_job_by_id_async(self, job_id: str, actor: PydanticUser) -> PydanticJob:
         """Fetch a job by its ID asynchronously."""
         async with db_registry.async_session() as session:
@@ -301,6 +305,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     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:
@@ -310,6 +315,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
     async def get_run_messages(
         self,
         run_id: str,
@@ -367,6 +373,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="run_id", expected_prefix=PrimitiveType.RUN)
     async def get_step_messages(
         self,
         run_id: str,
@@ -447,6 +454,7 @@ class JobManager:
         return job
 
     @enforce_types
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def record_ttft(self, job_id: str, ttft_ns: int, actor: PydanticUser) -> None:
         """Record time to first token for a run"""
         try:
@@ -459,6 +467,7 @@ class JobManager:
             logger.warning(f"Failed to record TTFT for job {job_id}: {e}")
 
     @enforce_types
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def record_response_duration(self, job_id: str, total_duration_ns: int, actor: PydanticUser) -> None:
         """Record total response duration for a run"""
         try:
@@ -529,6 +538,7 @@ class JobManager:
 
     @enforce_types
     @trace_method
+    @raise_on_invalid_id(param_name="job_id", expected_prefix=PrimitiveType.JOB)
     async def get_job_steps(
         self,
         job_id: str,

letta/services/mcp/base_client.py

@@ -83,10 +83,10 @@ class AsyncBaseMCPClient:
         for content_piece in result.content:
             if isinstance(content_piece, TextContent):
                 parsed_content.append(content_piece.text)
-                print("parsed_content (text)", parsed_content)
+                logger.debug(f"MCP tool result parsed content (text): {parsed_content}")
             else:
                 parsed_content.append(str(content_piece))
-                print("parsed_content (other)", parsed_content)
+                logger.debug(f"MCP tool result parsed content (other): {parsed_content}")
         if len(parsed_content) > 0:
             final_content = " ".join(parsed_content)
         else:

letta/services/mcp/oauth_utils.py

@@ -34,12 +34,21 @@ class DatabaseTokenStorage(TokenStorage):
     async def get_tokens(self) -> Optional[OAuthToken]:
         """Retrieve tokens from database."""
         oauth_session = await self.mcp_manager.get_oauth_session_by_id(self.session_id, self.actor)
-        if not oauth_session or not oauth_session.access_token:
+        if not oauth_session:
             return None
 
+        # Decrypt tokens using getter methods
+        access_token_secret = oauth_session.get_access_token_secret()
+        access_token = access_token_secret.get_plaintext()
+        if not access_token:
+            return None
+
+        refresh_token_secret = oauth_session.get_refresh_token_secret()
+        refresh_token = refresh_token_secret.get_plaintext()
+
         return OAuthToken(
-            access_token=oauth_session.access_token,
-            refresh_token=oauth_session.refresh_token,
+            access_token=access_token,
+            refresh_token=refresh_token,
             token_type=oauth_session.token_type,
             expires_in=int(oauth_session.expires_at.timestamp() - time.time()),
             scope=oauth_session.scope,
@@ -63,9 +72,13 @@ class DatabaseTokenStorage(TokenStorage):
         if not oauth_session or not oauth_session.client_id:
             return None
 
+        # Decrypt client secret using getter method
+        client_secret_secret = oauth_session.get_client_secret_secret()
+        client_secret = client_secret_secret.get_plaintext()
+
         return OAuthClientInformationFull(
             client_id=oauth_session.client_id,
-            client_secret=oauth_session.client_secret,
+            client_secret=client_secret,
             redirect_uris=[oauth_session.redirect_uri] if oauth_session.redirect_uri else [],
         )
 
@@ -134,13 +147,23 @@ class MCPOAuthSession:
 
     async def store_authorization_code(self, code: str, state: str) -> Optional[MCPOAuth]:
         """Store the authorization code from OAuth callback."""
+        # Use mcp_manager to ensure proper encryption
+        from letta.schemas.mcp import MCPOAuthSessionUpdate
+        from letta.schemas.secret import Secret
+
         async with db_registry.async_session() as session:
             try:
                 oauth_record = await MCPOAuth.read_async(db_session=session, identifier=self.session_id, actor=None)
-                oauth_record.authorization_code = code
-                oauth_record.state = state
+
+                # Encrypt the authorization_code before storing
+                if code is not None:
+                    oauth_record.authorization_code_enc = Secret.from_plaintext(code).get_encrypted()
+                    # Keep plaintext for dual-write during migration
+                    oauth_record.authorization_code = code
+
                 oauth_record.status = OAuthSessionStatus.AUTHORIZED
-                oauth_record.updated_at = datetime.now()
+                oauth_record.state = state
+
                 return await oauth_record.update_async(db_session=session, actor=None)
             except Exception:
                 return None
@@ -212,7 +235,9 @@ async def create_oauth_provider(
         while time.time() - start_time < timeout:
             oauth_session = await mcp_manager.get_oauth_session_by_id(session_id, actor)
             if oauth_session and oauth_session.authorization_code:
-                return oauth_session.authorization_code, oauth_session.state
+                # Decrypt the authorization code before returning
+                auth_code_secret = oauth_session.get_authorization_code_secret()
+                return auth_code_secret.get_plaintext(), oauth_session.state
             elif oauth_session and oauth_session.status == OAuthSessionStatus.ERROR:
                 raise Exception("OAuth authorization failed")
             await asyncio.sleep(1)