letta-nightly 0.8.8.dev20250703104323__py3-none-any.whl → 0.8.8.dev20250703174903__py3-none-any.whl

letta/services/agent_manager.py

@@ -16,9 +16,8 @@ from letta.constants import (
     BASE_VOICE_SLEEPTIME_CHAT_TOOLS,
     BASE_VOICE_SLEEPTIME_TOOLS,
     DEFAULT_TIMEZONE,
-    DEPRECATED_BASE_TOOLS,
+    DEPRECATED_LETTA_TOOLS,
     FILES_TOOLS,
-    MULTI_AGENT_TOOLS,
 )
 from letta.helpers import ToolRulesSolver
 from letta.helpers.datetime_helpers import get_utc_time
@@ -79,6 +78,7 @@ from letta.services.helpers.agent_manager_helper import (
     build_passage_query,
     build_source_passage_query,
     calculate_base_tools,
+    calculate_multi_agent_tools,
     check_supports_structured_output,
     compile_system_message,
     derive_system_message,
@@ -271,7 +271,7 @@ class AgentManager:
             else:
                 tool_names |= calculate_base_tools(is_v2=False)
         if agent_create.include_multi_agent_tools:
-            tool_names |= set(MULTI_AGENT_TOOLS)
+            tool_names |= calculate_multi_agent_tools()
 
         supplied_ids = set(agent_create.tool_ids or [])
 
@@ -294,7 +294,7 @@ class AgentManager:
                 tool_rules = list(agent_create.tool_rules or [])
                 if agent_create.include_base_tool_rules:
                     for tn in tool_names:
-                        if tn in {"send_message", "memory_finish_edits"}:
+                        if tn in {"send_message", "send_message_to_agent_async", "memory_finish_edits"}:
                             tool_rules.append(TerminalToolRule(tool_name=tn))
                         elif tn in (BASE_TOOLS + BASE_MEMORY_TOOLS + BASE_SLEEPTIME_TOOLS):
                             tool_rules.append(ContinueToolRule(tool_name=tn))
@@ -438,10 +438,10 @@ class AgentManager:
             else:
                 tool_names |= calculate_base_tools(is_v2=False)
         if agent_create.include_multi_agent_tools:
-            tool_names |= set(MULTI_AGENT_TOOLS)
+            tool_names |= calculate_multi_agent_tools()
 
         # take out the deprecated tool names
-        tool_names.difference_update(set(DEPRECATED_BASE_TOOLS))
+        tool_names.difference_update(set(DEPRECATED_LETTA_TOOLS))
 
         supplied_ids = set(agent_create.tool_ids or [])
 
@@ -479,7 +479,7 @@ class AgentManager:
                 tool_rules = list(agent_create.tool_rules or [])
                 if agent_create.include_base_tool_rules:
                     for tn in tool_names:
-                        if tn in {"send_message", "memory_finish_edits"}:
+                        if tn in {"send_message", "send_message_to_agent_async", "memory_finish_edits"}:
                             tool_rules.append(TerminalToolRule(tool_name=tn))
                         elif tn in (BASE_TOOLS + BASE_MEMORY_TOOLS + BASE_MEMORY_TOOLS_V2 + BASE_SLEEPTIME_TOOLS):
                             tool_rules.append(ContinueToolRule(tool_name=tn))
@@ -1111,6 +1111,7 @@ class AgentManager:
         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 await agent.to_pydantic_async(include_relationships=include_relationships)
@@ -1434,7 +1435,7 @@ class AgentManager:
 
         # 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()
+        curr_memory_str = agent_state.memory.compile(sources=agent_state.sources)
         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(
@@ -1461,6 +1462,7 @@ class AgentManager:
             timezone=agent_state.timezone,
             previous_message_count=num_messages - len(agent_state.message_ids),
             archival_memory_size=num_archival_memories,
+            sources=agent_state.sources,
         )
 
         diff = united_diff(curr_system_message_openai["content"], new_system_message_str)
@@ -1493,7 +1495,8 @@ class AgentManager:
 
         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)
+        # Get the current agent state
+        agent_state = await self.get_agent_by_id_async(agent_id=agent_id, include_relationships=["memory", "sources"], actor=actor)
         if not tool_rules_solver:
             tool_rules_solver = ToolRulesSolver(agent_state.tool_rules)
 
@@ -1509,7 +1512,9 @@ class AgentManager:
 
         # 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()
+        curr_memory_str = agent_state.memory.compile(
+            sources=agent_state.sources, tool_usage_rules=tool_rules_solver.compile_tool_rule_prompts()
+        )
         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(
@@ -1529,6 +1534,7 @@ class AgentManager:
         num_archival_memories = await self.passage_manager.agent_passage_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,
@@ -1537,6 +1543,7 @@ class AgentManager:
             previous_message_count=num_messages - len(agent_state.message_ids),
             archival_memory_size=num_archival_memories,
             tool_rules_solver=tool_rules_solver,
+            sources=agent_state.sources,
         )
 
         diff = united_diff(curr_system_message_openai["content"], new_system_message_str)
@@ -1654,7 +1661,7 @@ class AgentManager:
             # Update agent to only keep the system message
             agent.message_ids = [system_message_id]
             await agent.update_async(db_session=session, actor=actor)
-            agent_state = await agent.to_pydantic_async()
+            agent_state = await agent.to_pydantic_async(include_relationships=["sources"])
 
         # Optionally add default initial messages after the system message
         if add_default_initial_messages:
@@ -1688,9 +1695,13 @@ class AgentManager:
         Returns:
             modified (bool): whether the memory was updated
         """
-        agent_state = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor)
+        agent_state = await self.get_agent_by_id_async(agent_id=agent_id, actor=actor, include_relationships=["memory", "sources"])
         system_message = await self.message_manager.get_message_by_id_async(message_id=agent_state.message_ids[0], actor=actor)
-        if new_memory.compile() not in system_message.content[0].text:
+        temp_tool_rules_solver = ToolRulesSolver(agent_state.tool_rules)
+        if (
+            new_memory.compile(sources=agent_state.sources, tool_usage_rules=temp_tool_rules_solver.compile_tool_rule_prompts())
+            not in system_message.content[0].text
+        ):
             # update the blocks (LRW) in the DB
             for label in agent_state.memory.list_block_labels():
                 updated_value = new_memory.get_block(label).value
@@ -1730,7 +1741,9 @@ class AgentManager:
             agent_state.memory.blocks = [b for b in blocks if b is not None]
 
         if file_block_names:
-            file_blocks = await self.file_agent_manager.get_all_file_blocks_by_name(file_names=file_block_names, actor=actor)
+            file_blocks = await self.file_agent_manager.get_all_file_blocks_by_name(
+                file_names=file_block_names, agent_id=agent_state.id, actor=actor
+            )
             agent_state.memory.file_blocks = [b for b in file_blocks if b is not None]
 
         return agent_state
@@ -1772,8 +1785,7 @@ class AgentManager:
                 relationship_name="sources",
                 model_class=SourceModel,
                 item_ids=[source_id],
-                allow_partial=False,
-                replace=False,  # Extend existing sources rather than replace
+                replace=False,
             )
 
             # Commit the changes

letta/services/file_manager.py

@@ -1,11 +1,13 @@
+import os
 from datetime import datetime
 from typing import List, Optional
 
-from sqlalchemy import select, update
+from sqlalchemy import func, select, update
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import selectinload
 
+from letta.constants import MAX_FILENAME_LENGTH
 from letta.orm.errors import NoResultFound
 from letta.orm.file import FileContent as FileContentModel
 from letta.orm.file import FileMetadata as FileMetadataModel
@@ -13,6 +15,7 @@ from letta.orm.sqlalchemy_base import AccessType
 from letta.otel.tracing import trace_method
 from letta.schemas.enums import FileProcessingStatus
 from letta.schemas.file import FileMetadata as PydanticFileMetadata
+from letta.schemas.source import Source as PydanticSource
 from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.utils import enforce_types
@@ -58,11 +61,7 @@ class FileManager:
     @enforce_types
     @trace_method
     async def get_file_by_id(
-        self,
-        file_id: str,
-        actor: Optional[PydanticUser] = None,
-        *,
-        include_content: bool = False,
+        self, file_id: str, actor: Optional[PydanticUser] = None, *, include_content: bool = False, strip_directory_prefix: bool = False
     ) -> Optional[PydanticFileMetadata]:
         """Retrieve a file by its ID.
 
@@ -96,7 +95,7 @@ class FileManager:
                         actor=actor,
                     )
 
-                return await file_orm.to_pydantic_async(include_content=include_content)
+                return await file_orm.to_pydantic_async(include_content=include_content, strip_directory_prefix=strip_directory_prefix)
 
             except NoResultFound:
                 return None
@@ -193,7 +192,13 @@ class FileManager:
     @enforce_types
     @trace_method
     async def list_files(
-        self, source_id: str, actor: PydanticUser, after: Optional[str] = None, limit: Optional[int] = 50, include_content: bool = False
+        self,
+        source_id: str,
+        actor: PydanticUser,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+        include_content: bool = False,
+        strip_directory_prefix: bool = False,
     ) -> List[PydanticFileMetadata]:
         """List all files with optional pagination."""
         async with db_registry.async_session() as session:
@@ -207,7 +212,10 @@ class FileManager:
                 source_id=source_id,
                 query_options=options,
             )
-            return [await file.to_pydantic_async(include_content=include_content) for file in files]
+            return [
+                await file.to_pydantic_async(include_content=include_content, strip_directory_prefix=strip_directory_prefix)
+                for file in files
+            ]
 
     @enforce_types
     @trace_method
@@ -217,3 +225,44 @@ class FileManager:
             file = await FileMetadataModel.read_async(db_session=session, identifier=file_id)
             await file.hard_delete_async(db_session=session, actor=actor)
             return await file.to_pydantic_async()
+
+    @enforce_types
+    @trace_method
+    async def generate_unique_filename(self, original_filename: str, source: PydanticSource, organization_id: str) -> str:
+        """
+        Generate a unique filename by checking for duplicates and adding a numeric suffix if needed.
+        Similar to how filesystems handle duplicates (e.g., file.txt, file (1).txt, file (2).txt).
+
+        Parameters:
+            original_filename (str): The original filename as uploaded.
+            source_id (str): Source ID to check for duplicates within.
+            organization_id (str): Organization ID to check for duplicates within.
+
+        Returns:
+            str: A unique filename with numeric suffix if needed.
+        """
+        base, ext = os.path.splitext(original_filename)
+
+        # Reserve space for potential suffix: " (999)" = 6 characters
+        max_base_length = MAX_FILENAME_LENGTH - len(ext) - 6
+        if len(base) > max_base_length:
+            base = base[:max_base_length]
+            original_filename = f"{base}{ext}"
+
+        async with db_registry.async_session() as session:
+            # Count existing files with the same original_file_name in this source
+            query = select(func.count(FileMetadataModel.id)).where(
+                FileMetadataModel.original_file_name == original_filename,
+                FileMetadataModel.source_id == source.id,
+                FileMetadataModel.organization_id == organization_id,
+                FileMetadataModel.is_deleted == False,
+            )
+            result = await session.execute(query)
+            count = result.scalar() or 0
+
+            if count == 0:
+                # No duplicates, return original filename with source.name
+                return f"{source.name}/{original_filename}"
+            else:
+                # Add numeric suffix
+                return f"{source.name}/{base}_({count}){ext}"

letta/services/file_processor/chunker/llama_index_chunker.py

@@ -3,6 +3,7 @@ from typing import List, Tuple
 from mistralai import OCRPageObject
 
 from letta.log import get_logger
+from letta.otel.tracing import trace_method
 
 logger = get_logger(__name__)
 
@@ -19,6 +20,7 @@ class LlamaIndexChunker:
         self.parser = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
 
     # TODO: Make this more general beyond Mistral
+    @trace_method
     def chunk_text(self, page: OCRPageObject) -> List[str]:
         """Chunk text using LlamaIndex splitter"""
         try:

letta/services/file_processor/embedder/openai_embedder.py

@@ -1,10 +1,12 @@
 import asyncio
-from typing import List, Optional, Tuple
-
-import openai
+from typing import List, Optional, Tuple, cast
 
+from letta.llm_api.llm_client import LLMClient
+from letta.llm_api.openai_client import OpenAIClient
 from letta.log import get_logger
+from letta.otel.tracing import log_event, trace_method
 from letta.schemas.embedding_config import EmbeddingConfig
+from letta.schemas.enums import ProviderType
 from letta.schemas.passage import Passage
 from letta.schemas.user import User
 from letta.settings import model_settings
@@ -24,45 +26,83 @@ class OpenAIEmbedder:
         self.embedding_config = embedding_config or self.default_embedding_config
 
         # TODO: Unify to global OpenAI client
-        self.client = openai.AsyncOpenAI(api_key=model_settings.openai_api_key)
-        self.max_batch = 1024
+        self.client: OpenAIClient = cast(
+            OpenAIClient,
+            LLMClient.create(
+                provider_type=ProviderType.openai,
+                put_inner_thoughts_first=False,
+                actor=None,  # Not necessary
+            ),
+        )
         self.max_concurrent_requests = 20
 
+    @trace_method
     async def _embed_batch(self, batch: List[str], batch_indices: List[int]) -> List[Tuple[int, List[float]]]:
         """Embed a single batch and return embeddings with their original indices"""
-        response = await self.client.embeddings.create(model=self.embedding_config.embedding_model, input=batch)
-        return [(idx, res.embedding) for idx, res in zip(batch_indices, response.data)]
+        log_event(
+            "embedder.batch_started",
+            {
+                "batch_size": len(batch),
+                "model": self.embedding_config.embedding_model,
+                "embedding_endpoint_type": self.embedding_config.embedding_endpoint_type,
+            },
+        )
+        embeddings = await self.client.request_embeddings(inputs=batch, embedding_config=self.embedding_config)
+        log_event("embedder.batch_completed", {"batch_size": len(batch), "embeddings_generated": len(embeddings)})
+        return [(idx, e) for idx, e in zip(batch_indices, embeddings)]
 
+    @trace_method
     async def generate_embedded_passages(self, file_id: str, source_id: str, chunks: List[str], actor: User) -> List[Passage]:
         """Generate embeddings for chunks with batching and concurrent processing"""
         if not chunks:
             return []
 
         logger.info(f"Generating embeddings for {len(chunks)} chunks using {self.embedding_config.embedding_model}")
+        log_event(
+            "embedder.generation_started",
+            {
+                "total_chunks": len(chunks),
+                "model": self.embedding_config.embedding_model,
+                "embedding_endpoint_type": self.embedding_config.embedding_endpoint_type,
+                "batch_size": self.embedding_config.batch_size,
+                "file_id": file_id,
+                "source_id": source_id,
+            },
+        )
 
         # Create batches with their original indices
         batches = []
         batch_indices = []
 
-        for i in range(0, len(chunks), self.max_batch):
-            batch = chunks[i : i + self.max_batch]
-            indices = list(range(i, min(i + self.max_batch, len(chunks))))
+        for i in range(0, len(chunks), self.embedding_config.batch_size):
+            batch = chunks[i : i + self.embedding_config.batch_size]
+            indices = list(range(i, min(i + self.embedding_config.batch_size, len(chunks))))
             batches.append(batch)
             batch_indices.append(indices)
 
         logger.info(f"Processing {len(batches)} batches")
+        log_event(
+            "embedder.batching_completed",
+            {"total_batches": len(batches), "batch_size": self.embedding_config.batch_size, "total_chunks": len(chunks)},
+        )
 
         async def process(batch: List[str], indices: List[int]):
             try:
                 return await self._embed_batch(batch, indices)
             except Exception as e:
                 logger.error(f"Failed to embed batch of size {len(batch)}: {str(e)}")
+                log_event("embedder.batch_failed", {"batch_size": len(batch), "error": str(e), "error_type": type(e).__name__})
                 raise
 
         # Execute all batches concurrently with semaphore control
         tasks = [process(batch, indices) for batch, indices in zip(batches, batch_indices)]
 
+        log_event(
+            "embedder.concurrent_processing_started",
+            {"concurrent_tasks": len(tasks), "max_concurrent_requests": self.max_concurrent_requests},
+        )
         results = await asyncio.gather(*tasks)
+        log_event("embedder.concurrent_processing_completed", {"batches_processed": len(results)})
 
         # Flatten results and sort by original index
         indexed_embeddings = []
@@ -86,4 +126,8 @@ class OpenAIEmbedder:
             passages.append(passage)
 
         logger.info(f"Successfully generated {len(passages)} embeddings")
+        log_event(
+            "embedder.generation_completed",
+            {"passages_created": len(passages), "total_chunks_processed": len(chunks), "file_id": file_id, "source_id": source_id},
+        )
         return passages

letta/services/file_processor/file_processor.py

@@ -1,6 +1,7 @@
 from typing import List
 
 from letta.log import get_logger
+from letta.otel.tracing import log_event, trace_method
 from letta.schemas.agent import AgentState
 from letta.schemas.enums import FileProcessingStatus
 from letta.schemas.file import FileMetadata
@@ -42,6 +43,7 @@ class FileProcessor:
         self.actor = actor
 
     # TODO: Factor this function out of SyncServer
+    @trace_method
     async def process(
         self, server: SyncServer, agent_states: List[AgentState], source_id: str, content: bytes, file_metadata: FileMetadata
     ) -> List[Passage]:
@@ -50,6 +52,15 @@ class FileProcessor:
         # Create file as early as possible with no content
         file_metadata.processing_status = FileProcessingStatus.PARSING  # Parsing now
         file_metadata = await self.file_manager.create_file(file_metadata, self.actor)
+        log_event(
+            "file_processor.file_created",
+            {
+                "file_id": str(file_metadata.id),
+                "filename": filename,
+                "file_type": file_metadata.file_type,
+                "status": FileProcessingStatus.PARSING.value,
+            },
+        )
 
         try:
             # Ensure we're working with bytes
@@ -57,13 +68,22 @@ class FileProcessor:
                 content = content.encode("utf-8")
 
             if len(content) > self.max_file_size:
+                log_event(
+                    "file_processor.size_limit_exceeded",
+                    {"filename": filename, "file_size": len(content), "max_file_size": self.max_file_size},
+                )
                 raise ValueError(f"PDF size exceeds maximum allowed size of {self.max_file_size} bytes")
 
             logger.info(f"Starting OCR extraction for {filename}")
+            log_event("file_processor.ocr_started", {"filename": filename, "file_size": len(content), "mime_type": file_metadata.file_type})
             ocr_response = await self.file_parser.extract_text(content, mime_type=file_metadata.file_type)
 
             # update file with raw text
             raw_markdown_text = "".join([page.markdown for page in ocr_response.pages])
+            log_event(
+                "file_processor.ocr_completed",
+                {"filename": filename, "pages_extracted": len(ocr_response.pages), "text_length": len(raw_markdown_text)},
+            )
             file_metadata = await self.file_manager.update_file_status(
                 file_id=file_metadata.id, actor=self.actor, processing_status=FileProcessingStatus.EMBEDDING
             )
@@ -77,27 +97,56 @@ class FileProcessor:
             )
 
             if not ocr_response or len(ocr_response.pages) == 0:
+                log_event(
+                    "file_processor.ocr_no_text",
+                    {
+                        "filename": filename,
+                        "ocr_response_empty": not ocr_response,
+                        "pages_count": len(ocr_response.pages) if ocr_response else 0,
+                    },
+                )
                 raise ValueError("No text extracted from PDF")
 
             logger.info("Chunking extracted text")
+            log_event("file_processor.chunking_started", {"filename": filename, "pages_to_process": len(ocr_response.pages)})
             all_passages = []
 
             for page in ocr_response.pages:
                 chunks = self.text_chunker.chunk_text(page)
 
                 if not chunks:
+                    log_event("file_processor.chunking_failed", {"filename": filename, "page_index": ocr_response.pages.index(page)})
                     raise ValueError("No chunks created from text")
 
                 passages = await self.embedder.generate_embedded_passages(
                     file_id=file_metadata.id, source_id=source_id, chunks=chunks, actor=self.actor
                 )
+                log_event(
+                    "file_processor.page_processed",
+                    {
+                        "filename": filename,
+                        "page_index": ocr_response.pages.index(page),
+                        "chunks_created": len(chunks),
+                        "passages_generated": len(passages),
+                    },
+                )
                 all_passages.extend(passages)
 
             all_passages = await self.passage_manager.create_many_source_passages_async(
                 passages=all_passages, file_metadata=file_metadata, actor=self.actor
             )
+            log_event("file_processor.passages_created", {"filename": filename, "total_passages": len(all_passages)})
 
             logger.info(f"Successfully processed {filename}: {len(all_passages)} passages")
+            log_event(
+                "file_processor.processing_completed",
+                {
+                    "filename": filename,
+                    "file_id": str(file_metadata.id),
+                    "total_passages": len(all_passages),
+                    "status": FileProcessingStatus.COMPLETED.value,
+                },
+            )
 
             # update job status
             await self.file_manager.update_file_status(
@@ -108,6 +157,16 @@ class FileProcessor:
 
         except Exception as e:
             logger.error(f"File processing failed for {filename}: {str(e)}")
+            log_event(
+                "file_processor.processing_failed",
+                {
+                    "filename": filename,
+                    "file_id": str(file_metadata.id),
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "status": FileProcessingStatus.ERROR.value,
+                },
+            )
             await self.file_manager.update_file_status(
                 file_id=file_metadata.id, actor=self.actor, processing_status=FileProcessingStatus.ERROR, error_message=str(e)
             )

letta/services/file_processor/parser/mistral_parser.py

@@ -3,6 +3,7 @@ import base64
 from mistralai import Mistral, OCRPageObject, OCRResponse, OCRUsageInfo
 
 from letta.log import get_logger
+from letta.otel.tracing import trace_method
 from letta.services.file_processor.file_types import is_simple_text_mime_type
 from letta.services.file_processor.parser.base_parser import FileParser
 from letta.settings import settings
@@ -17,6 +18,7 @@ class MistralFileParser(FileParser):
         self.model = model
 
     # TODO: Make this return something general if we add more file parsers
+    @trace_method
     async def extract_text(self, content: bytes, mime_type: str) -> OCRResponse:
         """Extract text using Mistral OCR or shortcut for plain text."""
         try: