letta-nightly 0.7.29.dev20250602104315__py3-none-any.whl → 0.8.0.dev20250604104349__py3-none-any.whl

letta/services/block_manager.py

@@ -77,6 +77,30 @@ class BlockManager:
             # Convert back to Pydantic
             return [m.to_pydantic() for m in created_models]
 
+    @trace_method
+    @enforce_types
+    async def batch_create_blocks_async(self, blocks: List[PydanticBlock], actor: PydanticUser) -> List[PydanticBlock]:
+        """
+        Batch-create multiple Blocks in one transaction for better performance.
+        Args:
+            blocks: List of PydanticBlock schemas to create
+            actor:    The user performing the operation
+        Returns:
+            List of created PydanticBlock instances (with IDs, timestamps, etc.)
+        """
+        if not blocks:
+            return []
+
+        async with db_registry.async_session() as session:
+            block_models = [
+                BlockModel(**block.model_dump(to_orm=True, exclude_none=True), organization_id=actor.organization_id) for block in blocks
+            ]
+
+            created_models = await BlockModel.batch_create_async(items=block_models, db_session=session, actor=actor)
+
+            # Convert back to Pydantic
+            return [m.to_pydantic() for m in created_models]
+
     @trace_method
     @enforce_types
     def update_block(self, block_id: str, block_update: BlockUpdate, actor: PydanticUser) -> PydanticBlock:
@@ -238,9 +262,9 @@ class BlockManager:
             if actor:
                 query = BlockModel.apply_access_predicate(query, actor, ["read"], AccessType.ORGANIZATION)
 
-            # Add soft delete filter if applicable
-            if hasattr(BlockModel, "is_deleted"):
-                query = query.where(BlockModel.is_deleted == False)
+            # TODO: Add soft delete filter if applicable
+            # if hasattr(BlockModel, "is_deleted"):
+            #     query = query.where(BlockModel.is_deleted == False)
 
             # Execute the query
             result = await session.execute(query)
@@ -273,15 +297,12 @@ class BlockManager:
 
     @trace_method
     @enforce_types
-    def size(
-        self,
-        actor: PydanticUser,
-    ) -> int:
+    async def size_async(self, actor: PydanticUser) -> int:
         """
         Get the total count of blocks for the given user.
         """
-        with db_registry.session() as session:
-            return BlockModel.size(db_session=session, actor=actor)
+        async with db_registry.async_session() as session:
+            return await BlockModel.size_async(db_session=session, actor=actor)
 
     # Block History Functions
 

letta/services/context_window_calculator/context_window_calculator.py

@@ -0,0 +1,150 @@
+import asyncio
+from typing import Any, List, Optional, Tuple
+
+from openai.types.beta.function_tool import FunctionTool as OpenAITool
+
+from letta.log import get_logger
+from letta.schemas.enums import MessageRole
+from letta.schemas.letta_message_content import TextContent
+from letta.schemas.memory import ContextWindowOverview
+from letta.schemas.user import User as PydanticUser
+from letta.services.context_window_calculator.token_counter import TokenCounter
+
+logger = get_logger(__name__)
+
+
+class ContextWindowCalculator:
+    """Handles context window calculations with different token counting strategies"""
+
+    @staticmethod
+    def extract_system_components(system_message: str) -> Tuple[str, str, str]:
+        """Extract system prompt, core memory, and external memory summary from system message"""
+        base_start = system_message.find("<base_instructions>")
+        memory_blocks_start = system_message.find("<memory_blocks>")
+        metadata_start = system_message.find("<memory_metadata>")
+
+        system_prompt = ""
+        core_memory = ""
+        external_memory_summary = ""
+
+        if base_start != -1 and memory_blocks_start != -1:
+            system_prompt = system_message[base_start:memory_blocks_start].strip()
+
+        if memory_blocks_start != -1 and metadata_start != -1:
+            core_memory = system_message[memory_blocks_start:metadata_start].strip()
+
+        if metadata_start != -1:
+            external_memory_summary = system_message[metadata_start:].strip()
+
+        return system_prompt, core_memory, external_memory_summary
+
+    @staticmethod
+    def extract_summary_memory(messages: List[Any]) -> Tuple[Optional[str], int]:
+        """Extract summary memory if present and return starting index for real messages"""
+        if (
+            len(messages) > 1
+            and messages[1].role == MessageRole.user
+            and messages[1].content
+            and len(messages[1].content) == 1
+            and isinstance(messages[1].content[0], TextContent)
+            and "The following is a summary of the previous " in messages[1].content[0].text
+        ):
+            summary_memory = messages[1].content[0].text
+            start_index = 2
+            return summary_memory, start_index
+
+        return None, 1
+
+    async def calculate_context_window(
+        self, agent_state: Any, actor: PydanticUser, token_counter: TokenCounter, message_manager: Any, passage_manager: Any
+    ) -> ContextWindowOverview:
+        """Calculate context window information using the provided token counter"""
+
+        # Fetch data concurrently
+        (in_context_messages, passage_manager_size, message_manager_size) = await asyncio.gather(
+            message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=actor),
+            passage_manager.size_async(actor=actor, agent_id=agent_state.id),
+            message_manager.size_async(actor=actor, agent_id=agent_state.id),
+        )
+
+        # Convert messages to appropriate format
+        converted_messages = token_counter.convert_messages(in_context_messages)
+
+        # Extract system components
+        system_prompt = ""
+        core_memory = ""
+        external_memory_summary = ""
+
+        if (
+            in_context_messages
+            and in_context_messages[0].role == MessageRole.system
+            and in_context_messages[0].content
+            and len(in_context_messages[0].content) == 1
+            and isinstance(in_context_messages[0].content[0], TextContent)
+        ):
+            system_message = in_context_messages[0].content[0].text
+            system_prompt, core_memory, external_memory_summary = self.extract_system_components(system_message)
+
+        # System prompt
+        system_prompt = system_prompt or agent_state.system
+
+        # Extract summary memory
+        summary_memory, message_start_index = self.extract_summary_memory(in_context_messages)
+
+        # Prepare tool definitions
+        available_functions_definitions = []
+        if agent_state.tools:
+            available_functions_definitions = [OpenAITool(type="function", function=f.json_schema) for f in agent_state.tools]
+
+        # Count tokens concurrently
+        token_counts = await asyncio.gather(
+            token_counter.count_text_tokens(system_prompt),
+            token_counter.count_text_tokens(core_memory),
+            token_counter.count_text_tokens(external_memory_summary),
+            token_counter.count_text_tokens(summary_memory) if summary_memory else asyncio.sleep(0, result=0),
+            (
+                token_counter.count_message_tokens(converted_messages[message_start_index:])
+                if len(converted_messages) > message_start_index
+                else asyncio.sleep(0, result=0)
+            ),
+            (
+                token_counter.count_tool_tokens(available_functions_definitions)
+                if available_functions_definitions
+                else asyncio.sleep(0, result=0)
+            ),
+        )
+
+        (
+            num_tokens_system,
+            num_tokens_core_memory,
+            num_tokens_external_memory_summary,
+            num_tokens_summary_memory,
+            num_tokens_messages,
+            num_tokens_available_functions_definitions,
+        ) = token_counts
+
+        num_tokens_used_total = sum(token_counts)
+
+        return ContextWindowOverview(
+            # context window breakdown (in messages)
+            num_messages=len(in_context_messages),
+            num_archival_memory=passage_manager_size,
+            num_recall_memory=message_manager_size,
+            num_tokens_external_memory_summary=num_tokens_external_memory_summary,
+            external_memory_summary=external_memory_summary,
+            # top-level information
+            context_window_size_max=agent_state.llm_config.context_window,
+            context_window_size_current=num_tokens_used_total,
+            # context window breakdown (in tokens)
+            num_tokens_system=num_tokens_system,
+            system_prompt=system_prompt,
+            num_tokens_core_memory=num_tokens_core_memory,
+            core_memory=core_memory,
+            num_tokens_summary_memory=num_tokens_summary_memory,
+            summary_memory=summary_memory,
+            num_tokens_messages=num_tokens_messages,
+            messages=in_context_messages,
+            # related to functions
+            num_tokens_functions_definitions=num_tokens_available_functions_definitions,
+            functions_definitions=available_functions_definitions,
+        )

letta/services/context_window_calculator/token_counter.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from letta.llm_api.anthropic_client import AnthropicClient
+from letta.utils import count_tokens
+
+
+class TokenCounter(ABC):
+    """Abstract base class for token counting strategies"""
+
+    @abstractmethod
+    async def count_text_tokens(self, text: str) -> int:
+        """Count tokens in a text string"""
+
+    @abstractmethod
+    async def count_message_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        """Count tokens in a list of messages"""
+
+    @abstractmethod
+    async def count_tool_tokens(self, tools: List[Any]) -> int:
+        """Count tokens in tool definitions"""
+
+    @abstractmethod
+    def convert_messages(self, messages: List[Any]) -> List[Dict[str, Any]]:
+        """Convert messages to the appropriate format for this counter"""
+
+
+class AnthropicTokenCounter(TokenCounter):
+    """Token counter using Anthropic's API"""
+
+    def __init__(self, anthropic_client: AnthropicClient, model: str):
+        self.client = anthropic_client
+        self.model = model
+
+    async def count_text_tokens(self, text: str) -> int:
+        if not text:
+            return 0
+        return await self.client.count_tokens(model=self.model, messages=[{"role": "user", "content": text}])
+
+    async def count_message_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        if not messages:
+            return 0
+        return await self.client.count_tokens(model=self.model, messages=messages)
+
+    async def count_tool_tokens(self, tools: List[Any]) -> int:
+        if not tools:
+            return 0
+        return await self.client.count_tokens(model=self.model, tools=tools)
+
+    def convert_messages(self, messages: List[Any]) -> List[Dict[str, Any]]:
+        return [m.to_anthropic_dict() for m in messages]
+
+
+class TiktokenCounter(TokenCounter):
+    """Token counter using tiktoken"""
+
+    def __init__(self, model: str):
+        self.model = model
+
+    async def count_text_tokens(self, text: str) -> int:
+        if not text:
+            return 0
+        return count_tokens(text)
+
+    async def count_message_tokens(self, messages: List[Dict[str, Any]]) -> int:
+        if not messages:
+            return 0
+        from letta.local_llm.utils import num_tokens_from_messages
+
+        return num_tokens_from_messages(messages=messages, model=self.model)
+
+    async def count_tool_tokens(self, tools: List[Any]) -> int:
+        if not tools:
+            return 0
+        from letta.local_llm.utils import num_tokens_from_functions
+
+        # Extract function definitions from OpenAITool objects
+        functions = [t.function.model_dump() for t in tools]
+        return num_tokens_from_functions(functions=functions, model=self.model)
+
+    def convert_messages(self, messages: List[Any]) -> List[Dict[str, Any]]:
+        return [m.to_openai_dict() for m in messages]

letta/services/file_processor/chunker/llama_index_chunker.py

@@ -0,0 +1,29 @@
+from typing import List
+
+from mistralai import OCRPageObject
+
+from letta.log import get_logger
+
+logger = get_logger(__name__)
+
+
+class LlamaIndexChunker:
+    """LlamaIndex-based text chunking"""
+
+    def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+        from llama_index.core.node_parser import SentenceSplitter
+
+        self.parser = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+
+    # TODO: Make this more general beyond Mistral
+    def chunk_text(self, page: OCRPageObject) -> List[str]:
+        """Chunk text using LlamaIndex splitter"""
+        try:
+            return self.parser.split_text(page.markdown)
+
+        except Exception as e:
+            logger.error(f"Chunking failed: {str(e)}")
+            raise

letta/services/file_processor/embedder/openai_embedder.py

@@ -0,0 +1,84 @@
+import asyncio
+from typing import List, Optional, Tuple
+
+import openai
+
+from letta.log import get_logger
+from letta.schemas.embedding_config import EmbeddingConfig
+from letta.schemas.passage import Passage
+from letta.schemas.user import User
+from letta.settings import model_settings
+
+logger = get_logger(__name__)
+
+
+class OpenAIEmbedder:
+    """OpenAI-based embedding generation"""
+
+    def __init__(self, embedding_config: Optional[EmbeddingConfig] = None):
+        self.embedding_config = embedding_config or EmbeddingConfig.default_config(provider="openai")
+
+        # TODO: Unify to global OpenAI client
+        self.client = openai.AsyncOpenAI(api_key=model_settings.openai_api_key)
+        self.max_batch = 1024
+        self.max_concurrent_requests = 20
+
+    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)]
+
+    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}")
+
+        # 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))))
+            batches.append(batch)
+            batch_indices.append(indices)
+
+        logger.info(f"Processing {len(batches)} batches")
+
+        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)}")
+                raise
+
+        # Execute all batches concurrently with semaphore control
+        tasks = [process(batch, indices) for batch, indices in zip(batches, batch_indices)]
+
+        results = await asyncio.gather(*tasks)
+
+        # Flatten results and sort by original index
+        indexed_embeddings = []
+        for batch_result in results:
+            indexed_embeddings.extend(batch_result)
+
+        # Sort by index to maintain original order
+        indexed_embeddings.sort(key=lambda x: x[0])
+
+        # Create Passage objects in original order
+        passages = []
+        for (idx, embedding), text in zip(indexed_embeddings, chunks):
+            passage = Passage(
+                text=text,
+                file_id=file_id,
+                source_id=source_id,
+                embedding=embedding,
+                embedding_config=self.embedding_config,
+                organization_id=actor.organization_id,
+            )
+            passages.append(passage)
+
+        logger.info(f"Successfully generated {len(passages)} embeddings")
+        return passages

letta/services/file_processor/file_processor.py

@@ -0,0 +1,123 @@
+import mimetypes
+from typing import List, Optional
+
+from fastapi import UploadFile
+
+from letta.log import get_logger
+from letta.schemas.agent import AgentState
+from letta.schemas.enums import JobStatus
+from letta.schemas.file import FileMetadata
+from letta.schemas.job import Job, JobUpdate
+from letta.schemas.passage import Passage
+from letta.schemas.user import User
+from letta.server.server import SyncServer
+from letta.services.file_processor.chunker.llama_index_chunker import LlamaIndexChunker
+from letta.services.file_processor.embedder.openai_embedder import OpenAIEmbedder
+from letta.services.file_processor.parser.mistral_parser import MistralFileParser
+from letta.services.job_manager import JobManager
+from letta.services.passage_manager import PassageManager
+from letta.services.source_manager import SourceManager
+
+logger = get_logger(__name__)
+
+
+class FileProcessor:
+    """Main PDF processing orchestrator"""
+
+    def __init__(
+        self,
+        file_parser: MistralFileParser,
+        text_chunker: LlamaIndexChunker,
+        embedder: OpenAIEmbedder,
+        actor: User,
+        max_file_size: int = 50 * 1024 * 1024,  # 50MB default
+    ):
+        self.file_parser = file_parser
+        self.text_chunker = text_chunker
+        self.embedder = embedder
+        self.max_file_size = max_file_size
+        self.source_manager = SourceManager()
+        self.passage_manager = PassageManager()
+        self.job_manager = JobManager()
+        self.actor = actor
+
+    # TODO: Factor this function out of SyncServer
+    async def process(
+        self,
+        server: SyncServer,
+        agent_states: List[AgentState],
+        source_id: str,
+        content: bytes,
+        file: UploadFile,
+        job: Optional[Job] = None,
+    ) -> List[Passage]:
+        file_metadata = self._extract_upload_file_metadata(file, source_id=source_id)
+        file_metadata = await self.source_manager.create_file(file_metadata, self.actor)
+        filename = file_metadata.file_name
+
+        try:
+            # Ensure we're working with bytes
+            if isinstance(content, str):
+                content = content.encode("utf-8")
+
+            if len(content) > 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}")
+            ocr_response = await self.file_parser.extract_text(content, mime_type=file_metadata.file_type)
+
+            if not ocr_response or len(ocr_response.pages) == 0:
+                raise ValueError("No text extracted from PDF")
+
+            logger.info("Chunking extracted text")
+            all_passages = []
+            for page in ocr_response.pages:
+                chunks = self.text_chunker.chunk_text(page)
+
+                if not chunks:
+                    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
+                )
+                all_passages.extend(passages)
+
+            all_passages = await self.passage_manager.create_many_passages_async(all_passages, self.actor)
+
+            logger.info(f"Successfully processed {filename}: {len(all_passages)} passages")
+
+            await server.insert_file_into_context_windows(
+                source_id=source_id,
+                text="".join([ocr_response.pages[i].markdown for i in range(min(3, len(ocr_response.pages)))]),
+                file_id=file_metadata.id,
+                actor=self.actor,
+                agent_states=agent_states,
+            )
+
+            # update job status
+            if job:
+                job.status = JobStatus.completed
+                job.metadata["num_passages"] = len(all_passages)
+                await self.job_manager.update_job_by_id_async(job_id=job.id, job_update=JobUpdate(**job.model_dump()), actor=self.actor)
+
+            return all_passages
+
+        except Exception as e:
+            logger.error(f"PDF processing failed for {filename}: {str(e)}")
+
+            # update job status
+            if job:
+                job.status = JobStatus.failed
+                job.metadata["error"] = str(e)
+                await self.job_manager.update_job_by_id_async(job_id=job.id, job_update=JobUpdate(**job.model_dump()), actor=self.actor)
+
+            return []
+
+    def _extract_upload_file_metadata(self, file: UploadFile, source_id: str) -> FileMetadata:
+        file_metadata = {
+            "file_name": file.filename,
+            "file_path": None,
+            "file_type": mimetypes.guess_type(file.filename)[0] or file.content_type or "unknown",
+            "file_size": file.size if file.size is not None else None,
+        }
+        return FileMetadata(**file_metadata, source_id=source_id)

letta/services/file_processor/parser/base_parser.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+
+
+class FileParser(ABC):
+    """Abstract base class for file parser"""
+
+    @abstractmethod
+    async def extract_text(self, content: bytes, mime_type: str):
+        """Extract text from PDF content"""

letta/services/file_processor/parser/mistral_parser.py

@@ -0,0 +1,54 @@
+import base64
+
+from mistralai import Mistral, OCRPageObject, OCRResponse, OCRUsageInfo
+
+from letta.log import get_logger
+from letta.services.file_processor.parser.base_parser import FileParser
+from letta.settings import settings
+
+logger = get_logger(__name__)
+
+
+class MistralFileParser(FileParser):
+    """Mistral-based OCR extraction"""
+
+    def __init__(self, model: str = "mistral-ocr-latest"):
+        self.model = model
+
+    # TODO: Make this return something general if we add more file parsers
+    async def extract_text(self, content: bytes, mime_type: str) -> OCRResponse:
+        """Extract text using Mistral OCR or shortcut for plain text."""
+        try:
+            logger.info(f"Extracting text using Mistral OCR model: {self.model}")
+
+            # TODO: Kind of hacky...we try to exit early here?
+            # TODO: Create our internal file parser representation we return instead of OCRResponse
+            if mime_type == "text/plain":
+                text = content.decode("utf-8", errors="replace")
+                return OCRResponse(
+                    model=self.model,
+                    pages=[
+                        OCRPageObject(
+                            index=0,
+                            markdown=text,
+                            images=[],
+                            dimensions=None,
+                        )
+                    ],
+                    usage_info=OCRUsageInfo(pages_processed=1),  # You might need to construct this properly
+                    document_annotation=None,
+                )
+
+            base64_encoded_content = base64.b64encode(content).decode("utf-8")
+            document_url = f"data:{mime_type};base64,{base64_encoded_content}"
+
+            async with Mistral(api_key=settings.mistral_api_key) as mistral:
+                ocr_response = await mistral.ocr.process_async(
+                    model="mistral-ocr-latest", document={"type": "document_url", "document_url": document_url}, include_image_base64=False
+                )
+
+            return ocr_response
+
+        except Exception as e:
+            logger.error(f"OCR extraction failed: {str(e)}")
+            raise