letta-nightly 0.8.4.dev20250618104304__py3-none-any.whl → 0.8.5.dev20250619180801__py3-none-any.whl

letta/services/file_processor/chunker/llama_index_chunker.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 
 from mistralai import OCRPageObject
 
@@ -27,3 +27,91 @@ class LlamaIndexChunker:
         except Exception as e:
             logger.error(f"Chunking failed: {str(e)}")
             raise
+
+
+class MarkdownChunker:
+    """Markdown-specific chunker that preserves line numbers for citation purposes"""
+
+    def __init__(self, chunk_size: int = 2048):
+        self.chunk_size = chunk_size
+        # No overlap for line-based citations to avoid ambiguity
+
+        from llama_index.core.node_parser import MarkdownNodeParser
+
+        self.parser = MarkdownNodeParser()
+
+    def chunk_markdown_with_line_numbers(self, markdown_content: str) -> List[Tuple[str, int, int]]:
+        """
+        Chunk markdown content while preserving line number mappings.
+
+        Returns:
+            List of tuples: (chunk_text, start_line, end_line)
+        """
+        try:
+            # Split content into lines for line number tracking
+            lines = markdown_content.split("\n")
+
+            # Create nodes using MarkdownNodeParser
+            from llama_index.core import Document
+
+            document = Document(text=markdown_content)
+            nodes = self.parser.get_nodes_from_documents([document])
+
+            chunks_with_line_numbers = []
+
+            for node in nodes:
+                chunk_text = node.text
+
+                # Find the line numbers for this chunk
+                start_line, end_line = self._find_line_numbers(chunk_text, lines)
+
+                chunks_with_line_numbers.append((chunk_text, start_line, end_line))
+
+            return chunks_with_line_numbers
+
+        except Exception as e:
+            logger.error(f"Markdown chunking failed: {str(e)}")
+            # Fallback to simple line-based chunking
+            return self._fallback_line_chunking(markdown_content)
+
+    def _find_line_numbers(self, chunk_text: str, lines: List[str]) -> Tuple[int, int]:
+        """Find the start and end line numbers for a given chunk of text."""
+        chunk_lines = chunk_text.split("\n")
+
+        # Find the first line of the chunk in the original document
+        start_line = 1
+        for i, line in enumerate(lines):
+            if chunk_lines[0].strip() in line.strip() and len(chunk_lines[0].strip()) > 10:  # Avoid matching short lines
+                start_line = i + 1
+                break
+
+        # Calculate end line
+        end_line = start_line + len(chunk_lines) - 1
+
+        return start_line, min(end_line, len(lines))
+
+    def _fallback_line_chunking(self, markdown_content: str) -> List[Tuple[str, int, int]]:
+        """Fallback chunking method that simply splits by lines with no overlap."""
+        lines = markdown_content.split("\n")
+        chunks = []
+
+        i = 0
+        while i < len(lines):
+            chunk_lines = []
+            start_line = i + 1
+            char_count = 0
+
+            # Build chunk until we hit size limit
+            while i < len(lines) and char_count < self.chunk_size:
+                line = lines[i]
+                chunk_lines.append(line)
+                char_count += len(line) + 1  # +1 for newline
+                i += 1
+
+            end_line = i
+            chunk_text = "\n".join(chunk_lines)
+            chunks.append((chunk_text, start_line, end_line))
+
+            # No overlap - continue from where we left off
+
+        return chunks

letta/services/file_processor/embedder/openai_embedder.py

@@ -16,7 +16,12 @@ 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")
+        self.default_embedding_config = (
+            EmbeddingConfig.default_config(model_name="text-embedding-3-small", provider="openai")
+            if model_settings.openai_api_key
+            else EmbeddingConfig.default_config(model_name="letta")
+        )
+        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)

letta/services/file_processor/parser/mistral_parser.py

@@ -20,11 +20,10 @@ class MistralFileParser(FileParser):
     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 is_simple_text_mime_type(mime_type):
+                logger.info(f"Extracting text directly (no Mistral): {self.model}")
                 text = content.decode("utf-8", errors="replace")
                 return OCRResponse(
                     model=self.model,
@@ -43,6 +42,7 @@ class MistralFileParser(FileParser):
             base64_encoded_content = base64.b64encode(content).decode("utf-8")
             document_url = f"data:{mime_type};base64,{base64_encoded_content}"
 
+            logger.info(f"Extracting text using Mistral OCR model: {self.model}")
             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

letta/services/helpers/agent_manager_helper.py

@@ -449,41 +449,69 @@ def _cursor_filter(created_at_col, id_col, ref_created_at, ref_id, forward: bool
         )
 
 
-def _apply_pagination(query, before: Optional[str], after: Optional[str], session, ascending: bool = True) -> any:
+def _apply_pagination(
+    query, before: Optional[str], after: Optional[str], session, ascending: bool = True, sort_by: str = "created_at"
+) -> any:
+    # Determine the sort column
+    if sort_by == "last_run_completion":
+        sort_column = AgentModel.last_run_completion
+    else:
+        sort_column = AgentModel.created_at
+
     if after:
-        result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after)).first()
+        if sort_by == "last_run_completion":
+            result = session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == after)).first()
+        else:
+            result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after)).first()
         if result:
-            after_created_at, after_id = result
-            query = query.where(_cursor_filter(AgentModel.created_at, AgentModel.id, after_created_at, after_id, forward=ascending))
+            after_sort_value, after_id = result
+            query = query.where(_cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending))
 
     if before:
-        result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before)).first()
+        if sort_by == "last_run_completion":
+            result = session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == before)).first()
+        else:
+            result = session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before)).first()
         if result:
-            before_created_at, before_id = result
-            query = query.where(_cursor_filter(AgentModel.created_at, AgentModel.id, before_created_at, before_id, forward=not ascending))
+            before_sort_value, before_id = result
+            query = query.where(_cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending))
 
     # Apply ordering
     order_fn = asc if ascending else desc
-    query = query.order_by(order_fn(AgentModel.created_at), order_fn(AgentModel.id))
+    query = query.order_by(order_fn(sort_column), order_fn(AgentModel.id))
     return query
 
 
-async def _apply_pagination_async(query, before: Optional[str], after: Optional[str], session, ascending: bool = True) -> any:
+async def _apply_pagination_async(
+    query, before: Optional[str], after: Optional[str], session, ascending: bool = True, sort_by: str = "created_at"
+) -> any:
+    # Determine the sort column
+    if sort_by == "last_run_completion":
+        sort_column = AgentModel.last_run_completion
+    else:
+        sort_column = AgentModel.created_at
+
     if after:
-        result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after))).first()
+        if sort_by == "last_run_completion":
+            result = (await session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == after))).first()
+        else:
+            result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == after))).first()
         if result:
-            after_created_at, after_id = result
-            query = query.where(_cursor_filter(AgentModel.created_at, AgentModel.id, after_created_at, after_id, forward=ascending))
+            after_sort_value, after_id = result
+            query = query.where(_cursor_filter(sort_column, AgentModel.id, after_sort_value, after_id, forward=ascending))
 
     if before:
-        result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before))).first()
+        if sort_by == "last_run_completion":
+            result = (await session.execute(select(AgentModel.last_run_completion, AgentModel.id).where(AgentModel.id == before))).first()
+        else:
+            result = (await session.execute(select(AgentModel.created_at, AgentModel.id).where(AgentModel.id == before))).first()
         if result:
-            before_created_at, before_id = result
-            query = query.where(_cursor_filter(AgentModel.created_at, AgentModel.id, before_created_at, before_id, forward=not ascending))
+            before_sort_value, before_id = result
+            query = query.where(_cursor_filter(sort_column, AgentModel.id, before_sort_value, before_id, forward=not ascending))
 
     # Apply ordering
     order_fn = asc if ascending else desc
-    query = query.order_by(order_fn(AgentModel.created_at), order_fn(AgentModel.id))
+    query = query.order_by(order_fn(sort_column), order_fn(AgentModel.id))
     return query
 
 

letta/services/job_manager.py

@@ -6,6 +6,7 @@ from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from letta.helpers.datetime_helpers import get_utc_time
+from letta.log import get_logger
 from letta.orm.enums import JobType
 from letta.orm.errors import NoResultFound
 from letta.orm.job import Job as JobModel
@@ -28,6 +29,8 @@ from letta.schemas.user import User as PydanticUser
 from letta.server.db import db_registry
 from letta.utils import enforce_types
 
+logger = get_logger(__name__)
+
 
 class JobManager:
     """Manager class to handle business logic related to Jobs."""
@@ -67,18 +70,22 @@ class JobManager:
         with db_registry.session() as session:
             # Fetch the job by ID
             job = self._verify_job_access(session=session, job_id=job_id, actor=actor, access=["write"])
+            not_completed_before = not bool(job.completed_at)
 
             # Update job attributes with only the fields that were explicitly set
             update_data = job_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
 
             # Automatically update the completion timestamp if status is set to 'completed'
             for key, value in update_data.items():
+                # Ensure completed_at is timezone-naive for database compatibility
+                if key == "completed_at" and value is not None and hasattr(value, "replace"):
+                    value = value.replace(tzinfo=None)
                 setattr(job, key, value)
 
-            if update_data.get("status") == JobStatus.completed and not job.completed_at:
+            if job_update.status in {JobStatus.completed, JobStatus.failed} and not_completed_before:
                 job.completed_at = get_utc_time().replace(tzinfo=None)
                 if job.callback_url:
-                    self._dispatch_callback(session, job)
+                    self._dispatch_callback(job)
 
             # Save the updated job to the database
             job.update(db_session=session, actor=actor)
@@ -92,18 +99,22 @@ class JobManager:
         async with db_registry.async_session() as session:
             # Fetch the job by ID
             job = await self._verify_job_access_async(session=session, job_id=job_id, actor=actor, access=["write"])
+            not_completed_before = not bool(job.completed_at)
 
             # Update job attributes with only the fields that were explicitly set
             update_data = job_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
 
             # Automatically update the completion timestamp if status is set to 'completed'
             for key, value in update_data.items():
+                # Ensure completed_at is timezone-naive for database compatibility
+                if key == "completed_at" and value is not None and hasattr(value, "replace"):
+                    value = value.replace(tzinfo=None)
                 setattr(job, key, value)
 
-            if update_data.get("status") == JobStatus.completed and not job.completed_at:
+            if job_update.status in {JobStatus.completed, JobStatus.failed} and not_completed_before:
                 job.completed_at = get_utc_time().replace(tzinfo=None)
                 if job.callback_url:
-                    await self._dispatch_callback_async(session, job)
+                    await self._dispatch_callback_async(job)
 
             # Save the updated job to the database
             await job.update_async(db_session=session, actor=actor)
@@ -586,7 +597,7 @@ class JobManager:
             request_config = job.request_config or LettaRequestConfig()
         return request_config
 
-    def _dispatch_callback(self, session: Session, job: JobModel) -> None:
+    def _dispatch_callback(self, job: JobModel) -> None:
         """
         POST a standard JSON payload to job.callback_url
         and record timestamp + HTTP status.
@@ -595,22 +606,25 @@ class JobManager:
         payload = {
             "job_id": job.id,
             "status": job.status,
-            "completed_at": job.completed_at.isoformat(),
+            "completed_at": job.completed_at.isoformat() if job.completed_at else None,
+            "metadata": job.metadata_,
         }
         try:
             import httpx
 
             resp = httpx.post(job.callback_url, json=payload, timeout=5.0)
-            job.callback_sent_at = get_utc_time()
+            job.callback_sent_at = get_utc_time().replace(tzinfo=None)
             job.callback_status_code = resp.status_code
 
-        except Exception:
-            return
-
-        session.add(job)
-        session.commit()
+        except Exception as e:
+            error_message = f"Failed to dispatch callback for job {job.id} to {job.callback_url}: {str(e)}"
+            logger.error(error_message)
+            # Record the failed attempt
+            job.callback_sent_at = get_utc_time().replace(tzinfo=None)
+            job.callback_error = error_message
+            # Continue silently - callback failures should not affect job completion
 
-    async def _dispatch_callback_async(self, session, job: JobModel) -> None:
+    async def _dispatch_callback_async(self, job: JobModel) -> None:
         """
         POST a standard JSON payload to job.callback_url and record timestamp + HTTP status asynchronously.
         """
@@ -618,6 +632,7 @@ class JobManager:
             "job_id": job.id,
             "status": job.status,
             "completed_at": job.completed_at.isoformat() if job.completed_at else None,
+            "metadata": job.metadata_,
         }
 
         try:
@@ -628,7 +643,10 @@ class JobManager:
                 # Ensure timestamp is timezone-naive for DB compatibility
                 job.callback_sent_at = get_utc_time().replace(tzinfo=None)
                 job.callback_status_code = resp.status_code
-        except Exception:
-            # Silently fail on callback errors - job updates should still succeed
-            # In production, this would include proper error logging
-            pass
+        except Exception as e:
+            error_message = f"Failed to dispatch callback for job {job.id} to {job.callback_url}: {str(e)}"
+            logger.error(error_message)
+            # Record the failed attempt
+            job.callback_sent_at = get_utc_time().replace(tzinfo=None)
+            job.callback_error = error_message
+            # Continue silently - callback failures should not affect job completion

letta/services/mcp/base_client.py

@@ -1,3 +1,4 @@
+import asyncio
 from contextlib import AsyncExitStack
 from typing import Optional, Tuple
 
@@ -18,6 +19,9 @@ class AsyncBaseMCPClient:
         self.exit_stack = AsyncExitStack()
         self.session: Optional[ClientSession] = None
         self.initialized = False
+        # Track the task that created this client
+        self._creation_task = asyncio.current_task()
+        self._cleanup_queue = asyncio.Queue(maxsize=1)
 
     async def connect_to_server(self):
         try:
@@ -74,8 +78,29 @@ class AsyncBaseMCPClient:
             raise RuntimeError("MCPClient has not been initialized")
 
     async def cleanup(self):
-        """Clean up resources"""
+        """Clean up resources - ensure this runs in the same task"""
+        if hasattr(self, "_cleanup_task"):
+            # If we're in a different task, schedule cleanup in original task
+            current_task = asyncio.current_task()
+            if current_task != self._creation_task:
+                # Create a future to signal completion
+                cleanup_done = asyncio.Future()
+                self._cleanup_queue.put_nowait((self.exit_stack, cleanup_done))
+                await cleanup_done
+                return
+
+        # Normal cleanup
         await self.exit_stack.aclose()
 
     def to_sync_client(self):
         raise NotImplementedError("Subclasses must implement to_sync_client")
+
+    async def __aenter__(self):
+        """Enter the async context manager."""
+        await self.connect_to_server()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit the async context manager."""
+        await self.cleanup()
+        return False  # Don't suppress exceptions

letta/services/mcp_manager.py

@@ -75,23 +75,23 @@ class MCPManager:
             server_config = mcp_config[mcp_server_name]
 
         if isinstance(server_config, SSEServerConfig):
-            mcp_client = AsyncSSEMCPClient(server_config=server_config)
+            # mcp_client = AsyncSSEMCPClient(server_config=server_config)
+            async with AsyncSSEMCPClient(server_config=server_config) as mcp_client:
+                result, success = await mcp_client.execute_tool(tool_name, tool_args)
+                logger.info(f"MCP Result: {result}, Success: {success}")
+                return result, success
         elif isinstance(server_config, StdioServerConfig):
-            mcp_client = AsyncStdioMCPClient(server_config=server_config)
+            async with AsyncStdioMCPClient(server_config=server_config) as mcp_client:
+                result, success = await mcp_client.execute_tool(tool_name, tool_args)
+                logger.info(f"MCP Result: {result}, Success: {success}")
+                return result, success
         elif isinstance(server_config, StreamableHTTPServerConfig):
-            mcp_client = AsyncStreamableHTTPMCPClient(server_config=server_config)
+            async with AsyncStreamableHTTPMCPClient(server_config=server_config) as mcp_client:
+                result, success = await mcp_client.execute_tool(tool_name, tool_args)
+                logger.info(f"MCP Result: {result}, Success: {success}")
+                return result, success
         else:
             raise ValueError(f"Unsupported server config type: {type(server_config)}")
-        await mcp_client.connect_to_server()
-
-        # call tool
-        result, success = await mcp_client.execute_tool(tool_name, tool_args)
-        logger.info(f"MCP Result: {result}, Success: {success}")
-        # TODO: change to pydantic tool
-
-        await mcp_client.cleanup()
-
-        return result, success
 
     @enforce_types
     async def add_tool_from_mcp_server(self, mcp_server_name: str, mcp_tool_name: str, actor: PydanticUser) -> PydanticTool:
@@ -149,19 +149,19 @@ class MCPManager:
         return mcp_server
 
     @enforce_types
-    async def create_mcp_server(self, pydantic_mcp_server: MCPServer, actor: PydanticUser) -> PydanticTool:
-        """Create a new tool based on the ToolCreate schema."""
-        with db_registry.session() as session:
+    async def create_mcp_server(self, pydantic_mcp_server: MCPServer, actor: PydanticUser) -> MCPServer:
+        """Create a new MCP server."""
+        async with db_registry.async_session() as session:
             # Set the organization id at the ORM layer
             pydantic_mcp_server.organization_id = actor.organization_id
             mcp_server_data = pydantic_mcp_server.model_dump(to_orm=True)
 
             mcp_server = MCPServerModel(**mcp_server_data)
-            mcp_server.create(session, actor=actor)  # Re-raise other database-related errors
+            mcp_server = await mcp_server.create_async(session, actor=actor)
             return mcp_server.to_pydantic()
 
     @enforce_types
-    async def update_mcp_server_by_id(self, mcp_server_id: str, mcp_server_update: UpdateMCPServer, actor: PydanticUser) -> PydanticTool:
+    async def update_mcp_server_by_id(self, mcp_server_id: str, mcp_server_update: UpdateMCPServer, actor: PydanticUser) -> MCPServer:
         """Update a tool by its ID with the given ToolUpdate object."""
         async with db_registry.async_session() as session:
             # Fetch the tool by ID
@@ -177,6 +177,21 @@ class MCPManager:
             # Save the updated tool to the database mcp_server = await mcp_server.update_async(db_session=session, actor=actor)
             return mcp_server.to_pydantic()
 
+    @enforce_types
+    async def update_mcp_server_by_name(self, mcp_server_name: str, mcp_server_update: UpdateMCPServer, actor: PydanticUser) -> MCPServer:
+        """Update an MCP server by its name."""
+        mcp_server_id = await self.get_mcp_server_id_by_name(mcp_server_name, actor)
+        if not mcp_server_id:
+            raise HTTPException(
+                status_code=404,
+                detail={
+                    "code": "MCPServerNotFoundError",
+                    "message": f"MCP server {mcp_server_name} not found",
+                    "mcp_server_name": mcp_server_name,
+                },
+            )
+        return await self.update_mcp_server_by_id(mcp_server_id, mcp_server_update, actor)
+
     @enforce_types
     async def get_mcp_server_id_by_name(self, mcp_server_name: str, actor: PydanticUser) -> Optional[str]:
         """Retrieve a MCP server by its name and a user"""

letta/services/provider_manager.py

@@ -71,6 +71,25 @@ class ProviderManager:
             existing_provider.update(session, actor=actor)
             return existing_provider.to_pydantic()
 
+    @enforce_types
+    @trace_method
+    async def update_provider_async(self, provider_id: str, provider_update: ProviderUpdate, actor: PydanticUser) -> PydanticProvider:
+        """Update provider details."""
+        async with db_registry.async_session() as session:
+            # Retrieve the existing provider by ID
+            existing_provider = await ProviderModel.read_async(
+                db_session=session, identifier=provider_id, actor=actor, check_is_deleted=True
+            )
+
+            # Update only the fields that are provided in ProviderUpdate
+            update_data = provider_update.model_dump(to_orm=True, exclude_unset=True, exclude_none=True)
+            for key, value in update_data.items():
+                setattr(existing_provider, key, value)
+
+            # Commit the updated provider
+            await existing_provider.update_async(session, actor=actor)
+            return existing_provider.to_pydantic()
+
     @enforce_types
     @trace_method
     def delete_provider_by_id(self, provider_id: str, actor: PydanticUser):
@@ -175,6 +194,15 @@ class ProviderManager:
         providers = await self.list_providers_async(name=provider_name, actor=actor)
         return providers[0].api_key if providers else None
 
+    @enforce_types
+    @trace_method
+    async def get_bedrock_credentials_async(self, provider_name: Union[str, None], actor: PydanticUser) -> Optional[str]:
+        providers = await self.list_providers_async(name=provider_name, actor=actor)
+        access_key = providers[0].api_key if providers else None
+        secret_key = providers[0].api_secret if providers else None
+        region = providers[0].region if providers else None
+        return access_key, secret_key, region
+
     @enforce_types
     @trace_method
     def check_provider_api_key(self, provider_check: ProviderCheck) -> None:
@@ -183,6 +211,8 @@ class ProviderManager:
             provider_type=provider_check.provider_type,
             api_key=provider_check.api_key,
             provider_category=ProviderCategory.byok,
+            secret_key=provider_check.api_secret,
+            region=provider_check.region,
         ).cast_to_subtype()
 
         # TODO: add more string sanity checks here before we hit actual endpoints