letta-nightly 0.7.20.dev20250521104258__py3-none-any.whl → 0.7.21.dev20250522104246__py3-none-any.whl

letta/server/rest_api/streaming_response.py

@@ -0,0 +1,105 @@
+# Alternative implementation of StreamingResponse that allows for effectively
+# stremaing HTTP trailers, as we cannot set codes after the initial response.
+# Taken from: https://github.com/fastapi/fastapi/discussions/10138#discussioncomment-10377361
+
+import json
+from collections.abc import AsyncIterator
+
+from fastapi.responses import StreamingResponse
+from starlette.types import Send
+
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
+
+class StreamingResponseWithStatusCode(StreamingResponse):
+    """
+    Variation of StreamingResponse that can dynamically decide the HTTP status code,
+    based on the return value of the content iterator (parameter `content`).
+    Expects the content to yield either just str content as per the original `StreamingResponse`
+    or else tuples of (`content`: `str`, `status_code`: `int`).
+    """
+
+    body_iterator: AsyncIterator[str | bytes]
+    response_started: bool = False
+
+    async def stream_response(self, send: Send) -> None:
+        more_body = True
+        try:
+            first_chunk = await self.body_iterator.__anext__()
+            if isinstance(first_chunk, tuple):
+                first_chunk_content, self.status_code = first_chunk
+            else:
+                first_chunk_content = first_chunk
+            if isinstance(first_chunk_content, str):
+                first_chunk_content = first_chunk_content.encode(self.charset)
+
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": self.status_code,
+                    "headers": self.raw_headers,
+                }
+            )
+            self.response_started = True
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": first_chunk_content,
+                    "more_body": more_body,
+                }
+            )
+
+            async for chunk in self.body_iterator:
+                if isinstance(chunk, tuple):
+                    content, status_code = chunk
+                    if status_code // 100 != 2:
+                        # An error occurred mid-stream
+                        if not isinstance(content, bytes):
+                            content = content.encode(self.charset)
+                        more_body = False
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": content,
+                                "more_body": more_body,
+                            }
+                        )
+                        return
+                else:
+                    content = chunk
+
+                if isinstance(content, str):
+                    content = content.encode(self.charset)
+                more_body = True
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": content,
+                        "more_body": more_body,
+                    }
+                )
+
+        except Exception:
+            logger.exception("unhandled_streaming_error")
+            more_body = False
+            error_resp = {"error": {"message": "Internal Server Error"}}
+            error_event = f"event: error\ndata: {json.dumps(error_resp)}\n\n".encode(self.charset)
+            if not self.response_started:
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 500,
+                        "headers": self.raw_headers,
+                    }
+                )
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": error_event,
+                    "more_body": more_body,
+                }
+            )
+        if more_body:
+            await send({"type": "http.response.body", "body": b"", "more_body": False})

letta/server/rest_api/utils.py

@@ -190,6 +190,7 @@ def create_letta_messages_from_llm_response(
     pre_computed_assistant_message_id: Optional[str] = None,
     pre_computed_tool_message_id: Optional[str] = None,
     llm_batch_item_id: Optional[str] = None,
+    step_id: str | None = None,
 ) -> List[Message]:
     messages = []
 
@@ -244,6 +245,9 @@ def create_letta_messages_from_llm_response(
         )
         messages.append(heartbeat_system_message)
 
+    for message in messages:
+        message.step_id = step_id
+
     return messages
 
 

letta/server/server.py

@@ -94,6 +94,7 @@ from letta.services.provider_manager import ProviderManager
 from letta.services.sandbox_config_manager import SandboxConfigManager
 from letta.services.source_manager import SourceManager
 from letta.services.step_manager import StepManager
+from letta.services.telemetry_manager import TelemetryManager
 from letta.services.tool_executor.tool_execution_sandbox import ToolExecutionSandbox
 from letta.services.tool_manager import ToolManager
 from letta.services.user_manager import UserManager
@@ -213,6 +214,7 @@ class SyncServer(Server):
         self.identity_manager = IdentityManager()
         self.group_manager = GroupManager()
         self.batch_manager = LLMBatchManager()
+        self.telemetry_manager = TelemetryManager()
 
         # A resusable httpx client
         timeout = httpx.Timeout(connect=10.0, read=20.0, write=10.0, pool=10.0)
@@ -1000,6 +1002,30 @@ class SyncServer(Server):
         )
         return records
 
+    async def get_agent_archival_async(
+        self,
+        agent_id: str,
+        actor: User,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+        limit: Optional[int] = 100,
+        order_by: Optional[str] = "created_at",
+        reverse: Optional[bool] = False,
+        query_text: Optional[str] = None,
+        ascending: Optional[bool] = True,
+    ) -> List[Passage]:
+        # iterate over records
+        records = await self.agent_manager.list_passages_async(
+            actor=actor,
+            agent_id=agent_id,
+            after=after,
+            query_text=query_text,
+            before=before,
+            ascending=ascending,
+            limit=limit,
+        )
+        return records
+
     def insert_archival_memory(self, agent_id: str, memory_contents: str, actor: User) -> List[Passage]:
         # Get the agent object (loaded in memory)
         agent_state = self.agent_manager.get_agent_by_id(agent_id=agent_id, actor=actor)
@@ -1070,6 +1096,44 @@ class SyncServer(Server):
 
         return records
 
+    async def get_agent_recall_async(
+        self,
+        agent_id: str,
+        actor: User,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+        limit: Optional[int] = 100,
+        group_id: Optional[str] = None,
+        reverse: Optional[bool] = False,
+        return_message_object: bool = True,
+        use_assistant_message: bool = True,
+        assistant_message_tool_name: str = constants.DEFAULT_MESSAGE_TOOL,
+        assistant_message_tool_kwarg: str = constants.DEFAULT_MESSAGE_TOOL_KWARG,
+    ) -> Union[List[Message], List[LettaMessage]]:
+        records = await self.message_manager.list_messages_for_agent_async(
+            agent_id=agent_id,
+            actor=actor,
+            after=after,
+            before=before,
+            limit=limit,
+            ascending=not reverse,
+            group_id=group_id,
+        )
+
+        if not return_message_object:
+            records = Message.to_letta_messages_from_list(
+                messages=records,
+                use_assistant_message=use_assistant_message,
+                assistant_message_tool_name=assistant_message_tool_name,
+                assistant_message_tool_kwarg=assistant_message_tool_kwarg,
+                reverse=reverse,
+            )
+
+        if reverse:
+            records = records[::-1]
+
+        return records
+
     def get_server_config(self, include_defaults: bool = False) -> dict:
         """Return the base config"""
 
@@ -1301,6 +1365,48 @@ class SyncServer(Server):
 
         return llm_models
 
+    @trace_method
+    async def list_llm_models_async(
+        self,
+        actor: User,
+        provider_category: Optional[List[ProviderCategory]] = None,
+        provider_name: Optional[str] = None,
+        provider_type: Optional[ProviderType] = None,
+    ) -> List[LLMConfig]:
+        """Asynchronously list available models with maximum concurrency"""
+        import asyncio
+
+        providers = self.get_enabled_providers(
+            provider_category=provider_category,
+            provider_name=provider_name,
+            provider_type=provider_type,
+            actor=actor,
+        )
+
+        async def get_provider_models(provider):
+            try:
+                return await provider.list_llm_models_async()
+            except Exception as e:
+                import traceback
+
+                traceback.print_exc()
+                warnings.warn(f"An error occurred while listing LLM models for provider {provider}: {e}")
+                return []
+
+        # Execute all provider model listing tasks concurrently
+        provider_results = await asyncio.gather(*[get_provider_models(provider) for provider in providers])
+
+        # Flatten the results
+        llm_models = []
+        for models in provider_results:
+            llm_models.extend(models)
+
+        # Get local configs - if this is potentially slow, consider making it async too
+        local_configs = self.get_local_llm_configs()
+        llm_models.extend(local_configs)
+
+        return llm_models
+
     def list_embedding_models(self, actor: User) -> List[EmbeddingConfig]:
         """List available embedding models"""
         embedding_models = []
@@ -1311,6 +1417,35 @@ class SyncServer(Server):
                 warnings.warn(f"An error occurred while listing embedding models for provider {provider}: {e}")
         return embedding_models
 
+    async def list_embedding_models_async(self, actor: User) -> List[EmbeddingConfig]:
+        """Asynchronously list available embedding models with maximum concurrency"""
+        import asyncio
+
+        # Get all eligible providers first
+        providers = self.get_enabled_providers(actor=actor)
+
+        # Fetch embedding models from each provider concurrently
+        async def get_provider_embedding_models(provider):
+            try:
+                # All providers now have list_embedding_models_async
+                return await provider.list_embedding_models_async()
+            except Exception as e:
+                import traceback
+
+                traceback.print_exc()
+                warnings.warn(f"An error occurred while listing embedding models for provider {provider}: {e}")
+                return []
+
+        # Execute all provider model listing tasks concurrently
+        provider_results = await asyncio.gather(*[get_provider_embedding_models(provider) for provider in providers])
+
+        # Flatten the results
+        embedding_models = []
+        for models in provider_results:
+            embedding_models.extend(models)
+
+        return embedding_models
+
     def get_enabled_providers(
         self,
         actor: User,
@@ -1482,6 +1617,10 @@ class SyncServer(Server):
         letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
         return letta_agent.get_context_window()
 
+    async def get_agent_context_window_async(self, agent_id: str, actor: User) -> ContextWindowOverview:
+        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
+        return await letta_agent.get_context_window_async()
+
     def run_tool_from_source(
         self,
         actor: User,
@@ -1615,7 +1754,7 @@ class SyncServer(Server):
                                     server_name=server_name,
                                     command=server_params_raw["command"],
                                     args=server_params_raw.get("args", []),
-                                    env=server_params_raw.get("env", {})
+                                    env=server_params_raw.get("env", {}),
                                 )
                                 mcp_server_list[server_name] = server_params
                             except Exception as e:

letta/services/agent_manager.py

@@ -892,7 +892,7 @@ class AgentManager:
             List[PydanticAgentState]: The filtered list of matching agents.
         """
         async with db_registry.async_session() as session:
-            query = select(AgentModel).distinct(AgentModel.created_at, AgentModel.id)
+            query = select(AgentModel)
             query = AgentModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
 
             # Apply filters
@@ -961,6 +961,16 @@ class AgentManager:
         with db_registry.session() as session:
             return AgentModel.size(db_session=session, actor=actor)
 
+    async def size_async(
+        self,
+        actor: PydanticUser,
+    ) -> int:
+        """
+        Get the total count of agents for the given user.
+        """
+        async with db_registry.async_session() as session:
+            return await AgentModel.size_async(db_session=session, actor=actor)
+
     @enforce_types
     def get_agent_by_id(self, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         """Fetch an agent by its ID."""
@@ -969,18 +979,32 @@ class AgentManager:
             return agent.to_pydantic()
 
     @enforce_types
-    async def get_agent_by_id_async(self, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
+    async def get_agent_by_id_async(
+        self,
+        agent_id: str,
+        actor: PydanticUser,
+        include_relationships: Optional[List[str]] = None,
+    ) -> PydanticAgentState:
         """Fetch an agent by its ID."""
         async with db_registry.async_session() as session:
             agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
-            return agent.to_pydantic()
+            return await agent.to_pydantic_async(include_relationships=include_relationships)
 
     @enforce_types
-    async def get_agents_by_ids_async(self, agent_ids: list[str], actor: PydanticUser) -> list[PydanticAgentState]:
+    async def get_agents_by_ids_async(
+        self,
+        agent_ids: list[str],
+        actor: PydanticUser,
+        include_relationships: Optional[List[str]] = None,
+    ) -> list[PydanticAgentState]:
         """Fetch a list of agents by their IDs."""
         async with db_registry.async_session() as session:
-            agents = await AgentModel.read_multiple_async(db_session=session, identifiers=agent_ids, actor=actor)
-            return [await agent.to_pydantic_async() for agent in agents]
+            agents = await AgentModel.read_multiple_async(
+                db_session=session,
+                identifiers=agent_ids,
+                actor=actor,
+            )
+            return await asyncio.gather(*[agent.to_pydantic_async(include_relationships=include_relationships) for agent in agents])
 
     @enforce_types
     def get_agent_by_name(self, agent_name: str, actor: PydanticUser) -> PydanticAgentState:
@@ -1191,7 +1215,7 @@ class AgentManager:
 
     @enforce_types
     async def get_in_context_messages_async(self, agent_id: str, actor: PydanticUser) -> List[PydanticMessage]:
-        agent = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor)
+        agent = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=[], actor=actor)
         return await self.message_manager.get_messages_by_ids_async(message_ids=agent.message_ids, actor=actor)
 
     @enforce_types
@@ -1199,6 +1223,11 @@ class AgentManager:
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
         return self.message_manager.get_message_by_id(message_id=message_ids[0], actor=actor)
 
+    @enforce_types
+    async def get_system_message_async(self, agent_id: str, actor: PydanticUser) -> PydanticMessage:
+        agent = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=[], actor=actor)
+        return await self.message_manager.get_message_by_id_async(message_id=agent.message_ids[0], actor=actor)
+
     # TODO: This is duplicated below
     # TODO: This is legacy code and should be cleaned up
     # TODO: A lot of the memory "compilation" should be offset to a separate class
@@ -1267,10 +1296,81 @@ class AgentManager:
         else:
             return agent_state
 
+    @enforce_types
+    async def rebuild_system_prompt_async(
+        self, agent_id: str, actor: PydanticUser, force=False, update_timestamp=True
+    ) -> PydanticAgentState:
+        """Rebuilds the system message with the latest memory object and any shared memory block updates
+
+        Updates to core memory blocks should trigger a "rebuild", which itself will create a new message object
+
+        Updates to the memory header should *not* trigger a rebuild, since that will simply flood recall storage with excess messages
+        """
+        agent_state = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=["memory"], actor=actor)
+
+        curr_system_message = await self.get_system_message_async(
+            agent_id=agent_id, actor=actor
+        )  # this is the system + memory bank, not just the system prompt
+        curr_system_message_openai = curr_system_message.to_openai_dict()
+
+        # note: we only update the system prompt if the core memory is changed
+        # this means that the archival/recall memory statistics may be someout out of date
+        curr_memory_str = agent_state.memory.compile()
+        if curr_memory_str in curr_system_message_openai["content"] and not force:
+            # 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_id} and actor=({actor.id}, {actor.name}), skipping system prompt rebuild"
+            )
+            return agent_state
+
+        # If the memory didn't update, we probably don't want to update the timestamp inside
+        # For example, if we're doing a system prompt swap, this should probably be False
+        if update_timestamp:
+            memory_edit_timestamp = get_utc_time()
+        else:
+            # NOTE: a bit of a hack - we pull the timestamp from the message created_by
+            memory_edit_timestamp = curr_system_message.created_at
+
+        num_messages = await self.message_manager.size_async(actor=actor, agent_id=agent_id)
+        num_archival_memories = await self.passage_manager.size_async(actor=actor, agent_id=agent_id)
+
+        # update memory (TODO: potentially update recall/archival stats separately)
+        new_system_message_str = compile_system_message(
+            system_prompt=agent_state.system,
+            in_context_memory=agent_state.memory,
+            in_context_memory_last_edit=memory_edit_timestamp,
+            recent_passages=self.list_passages(actor=actor, agent_id=agent_id, ascending=False, limit=10),
+            previous_message_count=num_messages,
+            archival_memory_size=num_archival_memories,
+        )
+
+        diff = united_diff(curr_system_message_openai["content"], new_system_message_str)
+        if len(diff) > 0:  # there was a diff
+            logger.debug(f"Rebuilding system with new memory...\nDiff:\n{diff}")
+
+            # Swap the system message out (only if there is a diff)
+            message = PydanticMessage.dict_to_message(
+                agent_id=agent_id,
+                model=agent_state.llm_config.model,
+                openai_message_dict={"role": "system", "content": new_system_message_str},
+            )
+            message = await self.message_manager.update_message_by_id_async(
+                message_id=curr_system_message.id,
+                message_update=MessageUpdate(**message.model_dump()),
+                actor=actor,
+            )
+            return await self.set_in_context_messages_async(agent_id=agent_id, message_ids=agent_state.message_ids, actor=actor)
+        else:
+            return agent_state
+
     @enforce_types
     def set_in_context_messages(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState:
         return self.update_agent(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor)
 
+    @enforce_types
+    async def set_in_context_messages_async(self, agent_id: str, message_ids: List[str], actor: PydanticUser) -> PydanticAgentState:
+        return await self.update_agent_async(agent_id=agent_id, agent_update=UpdateAgent(message_ids=message_ids), actor=actor)
+
     @enforce_types
     def trim_older_in_context_messages(self, num: int, agent_id: str, actor: PydanticUser) -> PydanticAgentState:
         message_ids = self.get_agent_by_id(agent_id=agent_id, actor=actor).message_ids
@@ -1382,17 +1482,6 @@ class AgentManager:
 
         return agent_state
 
-    @enforce_types
-    def refresh_memory(self, agent_state: PydanticAgentState, actor: PydanticUser) -> PydanticAgentState:
-        block_ids = [b.id for b in agent_state.memory.blocks]
-        if not block_ids:
-            return agent_state
-
-        agent_state.memory.blocks = self.block_manager.get_all_blocks_by_ids(
-            block_ids=[b.id for b in agent_state.memory.blocks], actor=actor
-        )
-        return agent_state
-
     @enforce_types
     async def refresh_memory_async(self, agent_state: PydanticAgentState, actor: PydanticUser) -> PydanticAgentState:
         block_ids = [b.id for b in agent_state.memory.blocks]
@@ -1482,6 +1571,25 @@ class AgentManager:
             # Use the lazy-loaded relationship to get sources
             return [source.to_pydantic() for source in agent.sources]
 
+    @enforce_types
+    async def list_attached_sources_async(self, agent_id: str, actor: PydanticUser) -> List[PydanticSource]:
+        """
+        Lists all sources attached to an agent.
+
+        Args:
+            agent_id: ID of the agent to list sources for
+            actor: User performing the action
+
+        Returns:
+            List[str]: List of source IDs attached to the agent
+        """
+        async with db_registry.async_session() as session:
+            # Verify agent exists and user has permission to access it
+            agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
+
+            # Use the lazy-loaded relationship to get sources
+            return [source.to_pydantic() for source in agent.sources]
+
     @enforce_types
     def detach_source(self, agent_id: str, source_id: str, actor: PydanticUser) -> PydanticAgentState:
         """
@@ -1527,6 +1635,33 @@ class AgentManager:
                     return block.to_pydantic()
             raise NoResultFound(f"No block with label '{block_label}' found for agent '{agent_id}'")
 
+    @enforce_types
+    async def modify_block_by_label_async(
+        self,
+        agent_id: str,
+        block_label: str,
+        block_update: BlockUpdate,
+        actor: PydanticUser,
+    ) -> PydanticBlock:
+        """Gets a block attached to an agent by its label."""
+        async with db_registry.async_session() as session:
+            block = None
+            agent = await AgentModel.read_async(db_session=session, identifier=agent_id, actor=actor)
+            for block in agent.core_memory:
+                if block.label == block_label:
+                    block = block
+                    break
+            if not block:
+                raise NoResultFound(f"No block with label '{block_label}' found for agent '{agent_id}'")
+
+            update_data = block_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
+
+            for key, value in update_data.items():
+                setattr(block, key, value)
+
+            await block.update_async(session, actor=actor)
+            return block.to_pydantic()
+
     @enforce_types
     def update_block_with_label(
         self,
@@ -1848,6 +1983,65 @@ class AgentManager:
 
             return [p.to_pydantic() for p in passages]
 
+    @enforce_types
+    async def list_passages_async(
+        self,
+        actor: PydanticUser,
+        agent_id: Optional[str] = None,
+        file_id: Optional[str] = None,
+        limit: Optional[int] = 50,
+        query_text: Optional[str] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        source_id: Optional[str] = None,
+        embed_query: bool = False,
+        ascending: bool = True,
+        embedding_config: Optional[EmbeddingConfig] = None,
+        agent_only: bool = False,
+    ) -> List[PydanticPassage]:
+        """Lists all passages attached to an agent."""
+        async with db_registry.async_session() as session:
+            main_query = self._build_passage_query(
+                actor=actor,
+                agent_id=agent_id,
+                file_id=file_id,
+                query_text=query_text,
+                start_date=start_date,
+                end_date=end_date,
+                before=before,
+                after=after,
+                source_id=source_id,
+                embed_query=embed_query,
+                ascending=ascending,
+                embedding_config=embedding_config,
+                agent_only=agent_only,
+            )
+
+            # Add limit
+            if limit:
+                main_query = main_query.limit(limit)
+
+            # Execute query
+            result = await session.execute(main_query)
+
+            passages = []
+            for row in result:
+                data = dict(row._mapping)
+                if data["agent_id"] is not None:
+                    # This is an AgentPassage - remove source fields
+                    data.pop("source_id", None)
+                    data.pop("file_id", None)
+                    passage = AgentPassage(**data)
+                else:
+                    # This is a SourcePassage - remove agent field
+                    data.pop("agent_id", None)
+                    passage = SourcePassage(**data)
+                passages.append(passage)
+
+            return [p.to_pydantic() for p in passages]
+
     @enforce_types
     def passage_size(
         self,
@@ -2010,3 +2204,42 @@ class AgentManager:
             query = query.order_by(AgentsTags.tag).limit(limit)
             results = [tag[0] for tag in query.all()]
             return results
+
+    @enforce_types
+    async def list_tags_async(
+        self, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50, query_text: Optional[str] = None
+    ) -> List[str]:
+        """
+        Get all tags a user has created, ordered alphabetically.
+
+        Args:
+            actor: User performing the action.
+            after: Cursor for forward pagination.
+            limit: Maximum number of tags to return.
+            query text to filter tags by.
+
+        Returns:
+            List[str]: List of all tags.
+        """
+        async with db_registry.async_session() as session:
+            # Build the query using select() for async SQLAlchemy
+            query = (
+                select(AgentsTags.tag)
+                .join(AgentModel, AgentModel.id == AgentsTags.agent_id)
+                .where(AgentModel.organization_id == actor.organization_id)
+                .distinct()
+            )
+
+            if query_text:
+                query = query.where(AgentsTags.tag.ilike(f"%{query_text}%"))
+
+            if after:
+                query = query.where(AgentsTags.tag > after)
+
+            query = query.order_by(AgentsTags.tag).limit(limit)
+
+            # Execute the query asynchronously
+            result = await session.execute(query)
+            # Extract the tag values from the result
+            results = [row[0] for row in result.all()]
+        return results