letta-nightly 0.8.15.dev20250719104256__py3-none-any.whl → 0.8.16.dev20250721070720__py3-none-any.whl

letta/schemas/source.py

@@ -9,11 +9,17 @@ from letta.schemas.letta_base import LettaBase
 
 class BaseSource(LettaBase):
     """
-    Shared attributes accourss all source schemas.
+    Shared attributes across all source schemas.
     """
 
     __id_prefix__ = "source"
 
+    # Core source fields
+    name: str = Field(..., description="The name of the source.")
+    description: Optional[str] = Field(None, description="The description of the source.")
+    instructions: Optional[str] = Field(None, description="Instructions for how to use the source.")
+    metadata: Optional[dict] = Field(None, description="Metadata associated with the source.")
+
 
 class Source(BaseSource):
     """
@@ -29,9 +35,6 @@ class Source(BaseSource):
     """
 
     id: str = BaseSource.generate_id_field()
-    name: str = Field(..., description="The name of the source.")
-    description: Optional[str] = Field(None, description="The description of the source.")
-    instructions: Optional[str] = Field(None, description="Instructions for how to use the source.")
     embedding_config: EmbeddingConfig = Field(..., description="The embedding configuration used by the source.")
     organization_id: Optional[str] = Field(None, description="The ID of the organization that created the source.")
     metadata: Optional[dict] = Field(None, validation_alias="metadata_", description="Metadata associated with the source.")
@@ -48,29 +51,24 @@ class SourceCreate(BaseSource):
     Schema for creating a new Source.
     """
 
-    # required
-    name: str = Field(..., description="The name of the source.")
     # TODO: @matt, make this required after shub makes the FE changes
-
-    embedding: Optional[str] = Field(None, description="The hande for the embedding config used by the source.")
+    embedding: Optional[str] = Field(None, description="The handle for the embedding config used by the source.")
     embedding_chunk_size: Optional[int] = Field(None, description="The chunk size of the embedding.")
 
     # TODO: remove (legacy config)
     embedding_config: Optional[EmbeddingConfig] = Field(None, description="(Legacy) The embedding configuration used by the source.")
 
-    # optional
-    description: Optional[str] = Field(None, description="The description of the source.")
-    instructions: Optional[str] = Field(None, description="Instructions for how to use the source.")
-    metadata: Optional[dict] = Field(None, description="Metadata associated with the source.")
-
 
 class SourceUpdate(BaseSource):
     """
     Schema for updating an existing Source.
     """
 
+    # Override base fields to make them optional for updates
     name: Optional[str] = Field(None, description="The name of the source.")
     description: Optional[str] = Field(None, description="The description of the source.")
     instructions: Optional[str] = Field(None, description="Instructions for how to use the source.")
     metadata: Optional[dict] = Field(None, description="Metadata associated with the source.")
+
+    # Additional update-specific fields
     embedding_config: Optional[EmbeddingConfig] = Field(None, description="The embedding configuration used by the source.")

letta/schemas/step.py

@@ -1,8 +1,10 @@
+from enum import Enum, auto
 from typing import Dict, List, Literal, Optional
 
 from pydantic import Field
 
 from letta.schemas.letta_base import LettaBase
+from letta.schemas.letta_stop_reason import StopReasonType
 from letta.schemas.message import Message
 
 
@@ -28,6 +30,7 @@ class Step(StepBase):
     prompt_tokens: Optional[int] = Field(None, description="The number of tokens in the prompt during this step.")
     total_tokens: Optional[int] = Field(None, description="The total number of tokens processed by the agent during this step.")
     completion_tokens_details: Optional[Dict] = Field(None, description="Metadata for the agent.")
+    stop_reason: Optional[StopReasonType] = Field(None, description="The stop reason associated with the step.")
     tags: List[str] = Field([], description="Metadata tags.")
     tid: Optional[str] = Field(None, description="The unique identifier of the transaction that processed this step.")
     trace_id: Optional[str] = Field(None, description="The trace id of the agent step.")
@@ -36,3 +39,12 @@ class Step(StepBase):
         None, description="The feedback for this step. Must be either 'positive' or 'negative'."
     )
     project_id: Optional[str] = Field(None, description="The project that the agent that executed this step belongs to (cloud only).")
+
+
+class StepProgression(int, Enum):
+    START = auto()
+    STREAM_RECEIVED = auto()
+    RESPONSE_RECEIVED = auto()
+    STEP_LOGGED = auto()
+    LOGGED_TRACE = auto()
+    FINISHED = auto()

letta/schemas/tool.py

@@ -50,7 +50,6 @@ class Tool(BaseTool):
     tool_type: ToolType = Field(ToolType.CUSTOM, description="The type of the tool.")
     description: Optional[str] = Field(None, description="The description of the tool.")
     source_type: Optional[str] = Field(None, description="The type of the source code.")
-    organization_id: Optional[str] = Field(None, description="The unique identifier of the organization associated with the tool.")
     name: Optional[str] = Field(None, description="The name of the function.")
     tags: List[str] = Field([], description="Metadata tags.")
 
@@ -147,9 +146,6 @@ class ToolCreate(LettaBase):
     return_char_limit: int = Field(FUNCTION_RETURN_CHAR_LIMIT, description="The maximum number of characters in the response.")
     pip_requirements: Optional[List[PipRequirement]] = Field(None, description="Optional list of pip packages required by this tool.")
 
-    # TODO should we put the HTTP / API fetch inside from_mcp?
-    # async def from_mcp(cls, mcp_server: str, mcp_tool_name: str) -> "ToolCreate":
-
     @classmethod
     def from_mcp(cls, mcp_server_name: str, mcp_tool: MCPTool) -> "ToolCreate":
         from letta.functions.helpers import generate_mcp_tool_wrapper

letta/serialize_schemas/marshmallow_agent.py

@@ -86,7 +86,9 @@ class MarshmallowAgentSchema(BaseSchema):
         - Marks messages as in-context, preserving the order of the original `message_ids`
         - Removes individual message `id` fields
         """
-        data = super().sanitize_ids(data, **kwargs)
+        del data["id"]
+        del data["_created_by_id"]
+        del data["_last_updated_by_id"]
         data[self.FIELD_VERSION] = letta.__version__
 
         original_message_ids = data.pop(self.FIELD_MESSAGE_IDS, [])
@@ -107,6 +109,15 @@ class MarshmallowAgentSchema(BaseSchema):
 
         return data
 
+    @pre_load
+    def regenerate_ids(self, data: Dict, **kwargs) -> Dict:
+        if self.Meta.model:
+            data["id"] = self.generate_id()
+            data["_created_by_id"] = self.actor.id
+            data["_last_updated_by_id"] = self.actor.id
+
+        return data
+
     @post_dump
     def hide_tool_exec_environment_variables(self, data: Dict, **kwargs):
         """Hide the value of tool_exec_environment_variables"""
@@ -135,4 +146,6 @@ class MarshmallowAgentSchema(BaseSchema):
             "identities",
             "is_deleted",
             "groups",
+            "batch_items",
+            "organization",
         )

letta/serialize_schemas/marshmallow_block.py

@@ -1,3 +1,7 @@
+from typing import Dict
+
+from marshmallow import post_dump, pre_load
+
 from letta.orm.block import Block
 from letta.schemas.block import Block as PydanticBlock
 from letta.serialize_schemas.marshmallow_base import BaseSchema
@@ -10,6 +14,24 @@ class SerializedBlockSchema(BaseSchema):
 
     __pydantic_model__ = PydanticBlock
 
+    @post_dump
+    def sanitize_ids(self, data: Dict, **kwargs) -> Dict:
+        # delete id
+        del data["id"]
+        del data["_created_by_id"]
+        del data["_last_updated_by_id"]
+
+        return data
+
+    @pre_load
+    def regenerate_ids(self, data: Dict, **kwargs) -> Dict:
+        if self.Meta.model:
+            data["id"] = self.generate_id()
+            data["_created_by_id"] = self.actor.id
+            data["_last_updated_by_id"] = self.actor.id
+
+        return data
+
     class Meta(BaseSchema.Meta):
         model = Block
-        exclude = BaseSchema.Meta.exclude + ("agents", "identities", "is_deleted")
+        exclude = BaseSchema.Meta.exclude + ("agents", "identities", "is_deleted", "groups", "organization")

letta/serialize_schemas/marshmallow_message.py

@@ -23,7 +23,6 @@ class SerializedMessageSchema(BaseSchema):
         # agent dump will then get rid of message ids
         del data["_created_by_id"]
         del data["_last_updated_by_id"]
-        del data["organization"]
 
         return data
 
@@ -33,10 +32,9 @@ class SerializedMessageSchema(BaseSchema):
             # Skip regenerating ID, as agent dump will do it
             data["_created_by_id"] = self.actor.id
             data["_last_updated_by_id"] = self.actor.id
-            data["organization"] = self.actor.organization_id
 
         return data
 
     class Meta(BaseSchema.Meta):
         model = Message
-        exclude = BaseSchema.Meta.exclude + ("step", "job_message", "otid", "is_deleted")
+        exclude = BaseSchema.Meta.exclude + ("step", "job_message", "otid", "is_deleted", "organization")

letta/serialize_schemas/marshmallow_tool.py

@@ -1,3 +1,7 @@
+from typing import Dict
+
+from marshmallow import post_dump, pre_load
+
 from letta.orm import Tool
 from letta.schemas.tool import Tool as PydanticTool
 from letta.serialize_schemas.marshmallow_base import BaseSchema
@@ -10,6 +14,24 @@ class SerializedToolSchema(BaseSchema):
 
     __pydantic_model__ = PydanticTool
 
+    @post_dump
+    def sanitize_ids(self, data: Dict, **kwargs) -> Dict:
+        # delete id
+        del data["id"]
+        del data["_created_by_id"]
+        del data["_last_updated_by_id"]
+
+        return data
+
+    @pre_load
+    def regenerate_ids(self, data: Dict, **kwargs) -> Dict:
+        if self.Meta.model:
+            data["id"] = self.generate_id()
+            data["_created_by_id"] = self.actor.id
+            data["_last_updated_by_id"] = self.actor.id
+
+        return data
+
     class Meta(BaseSchema.Meta):
         model = Tool
-        exclude = BaseSchema.Meta.exclude + ("is_deleted",)
+        exclude = BaseSchema.Meta.exclude + ("is_deleted", "organization")

letta/server/db.py

@@ -1,20 +1,22 @@
 import os
 import threading
+import time
 import uuid
 from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncGenerator, Generator
 
+from opentelemetry import trace
 from rich.console import Console
 from rich.panel import Panel
 from rich.text import Text
-from sqlalchemy import Engine, NullPool, QueuePool, create_engine
+from sqlalchemy import Engine, NullPool, QueuePool, create_engine, event
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
 from sqlalchemy.orm import sessionmaker
 
 from letta.config import LettaConfig
 from letta.log import get_logger
 from letta.otel.tracing import trace_method
-from letta.settings import settings
+from letta.settings import DatabaseChoice, settings
 
 logger = get_logger(__name__)
 
@@ -36,6 +38,46 @@ def print_sqlite_schema_error():
     console.print(Panel(error_text, border_style="red"))
 
 
+@event.listens_for(Engine, "connect")
+def enable_sqlite_foreign_keys(dbapi_connection, connection_record):
+    """Enable foreign key constraints for SQLite connections."""
+    if "sqlite" in str(dbapi_connection):
+        cursor = dbapi_connection.cursor()
+        cursor.execute("PRAGMA foreign_keys=ON")
+        cursor.close()
+
+
+def on_connect(dbapi_connection, connection_record):
+    cursor = dbapi_connection.cursor()
+    cursor.execute("SELECT pg_backend_pid()")
+    pid = cursor.fetchone()[0]
+    connection_record.info["pid"] = pid
+    connection_record.info["connect_spawn_time_ms"] = time.perf_counter() * 1000
+    cursor.close()
+
+
+def on_close(dbapi_connection, connection_record):
+    connection_record.info.get("pid")
+    (time.perf_counter() * 1000) - connection_record.info.get("connect_spawn_time_ms")
+    # print(f"Connection closed: {pid}, duration: {duration:.6f}s")
+
+
+def on_checkout(dbapi_connection, connection_record, connection_proxy):
+    connection_record.info.get("pid")
+    connection_record.info["connect_checkout_time_ms"] = time.perf_counter() * 1000
+
+
+def on_checkin(dbapi_connection, connection_record):
+    pid = connection_record.info.get("pid")
+    duration = (time.perf_counter() * 1000) - connection_record.info.get("connect_checkout_time_ms")
+
+    tracer = trace.get_tracer("letta.db.connection")
+    with tracer.start_as_current_span("connect_release") as span:
+        span.set_attribute("db.connection.pid", pid)
+        span.set_attribute("db.connection.duration_ms", duration)
+        span.set_attribute("db.connection.operation", "checkin")
+
+
 @contextmanager
 def db_error_handler():
     """Context manager for handling database errors"""
@@ -43,6 +85,14 @@ def db_error_handler():
         yield
     except Exception as e:
         # Handle other SQLAlchemy errors
+        error_str = str(e)
+
+        # Don't exit for expected constraint violations that should be handled by the application
+        if "UNIQUE constraint failed" in error_str or "FOREIGN KEY constraint failed" in error_str:
+            # These are application-level errors that should be handled by the ORM
+            raise
+
+        # For other database errors, print error and exit
         print(e)
         print_sqlite_schema_error()
         # raise ValueError(f"SQLite DB error: {str(e)}")
@@ -73,7 +123,7 @@ class DatabaseRegistry:
                 return
 
             # Postgres engine
-            if settings.letta_pg_uri_no_default:
+            if settings.database_engine is DatabaseChoice.POSTGRES:
                 self.logger.info("Creating postgres engine")
                 self.config.recall_storage_type = "postgres"
                 self.config.recall_storage_uri = settings.letta_pg_uri_no_default
@@ -99,6 +149,15 @@ class DatabaseRegistry:
                 Base.metadata.create_all(bind=engine)
                 self._engines["default"] = engine
 
+            # Set up connection monitoring
+            if settings.sqlalchemy_tracing and settings.database_engine is DatabaseChoice.POSTGRES:
+                event.listen(engine, "connect", on_connect)
+                event.listen(engine, "close", on_close)
+                event.listen(engine, "checkout", on_checkout)
+                event.listen(engine, "checkin", on_checkin)
+
+            self._setup_pool_monitoring(engine, "default")
+
             # Create session factory
             self._session_factories["default"] = sessionmaker(autocommit=False, autoflush=False, bind=self._engines["default"])
             self._initialized["sync"] = True
@@ -109,7 +168,7 @@ class DatabaseRegistry:
             if self._initialized.get("async") and not force:
                 return
 
-            if settings.letta_pg_uri_no_default:
+            if settings.database_engine is DatabaseChoice.POSTGRES:
                 self.logger.info("Creating async postgres engine")
 
                 # Create async engine - convert URI to async format
@@ -128,10 +187,27 @@ class DatabaseRegistry:
                 self.logger.info("Creating sqlite engine " + engine_path)
                 async_engine = create_async_engine(engine_path, **self._build_sqlalchemy_engine_args(is_async=True))
 
+                # Enable foreign keys for SQLite async connections
+                @event.listens_for(async_engine.sync_engine, "connect")
+                def enable_sqlite_foreign_keys_async(dbapi_connection, connection_record):
+                    cursor = dbapi_connection.cursor()
+                    cursor.execute("PRAGMA foreign_keys=ON")
+                    cursor.close()
+
             # Create async session factory
             self._async_engines["default"] = async_engine
+
+            # Set up connection monitoring for async engine
+            if settings.sqlalchemy_tracing and settings.database_engine is DatabaseChoice.POSTGRES:
+                event.listen(async_engine.sync_engine, "connect", on_connect)
+                event.listen(async_engine.sync_engine, "close", on_close)
+                event.listen(async_engine.sync_engine, "checkout", on_checkout)
+                event.listen(async_engine.sync_engine, "checkin", on_checkin)
+
+            self._setup_pool_monitoring(async_engine, "default_async")
+
             self._async_session_factories["default"] = async_sessionmaker(
-                expire_on_commit=True,
+                expire_on_commit=False,
                 close_resets_only=False,
                 autocommit=False,
                 autoflush=False,
@@ -149,7 +225,10 @@ class DatabaseRegistry:
             pool_cls = NullPool
         else:
             logger.info("Enabling pooling on SqlAlchemy")
-            pool_cls = QueuePool if not is_async else None
+            # AsyncAdaptedQueuePool will be the default if none is provided for async but setting this explicitly.
+            from sqlalchemy import AsyncAdaptedQueuePool
+
+            pool_cls = QueuePool if not is_async else AsyncAdaptedQueuePool
 
         base_args = {
             "echo": settings.pg_echo,
@@ -207,11 +286,31 @@ class DatabaseRegistry:
 
         engine.connect = wrapped_connect
 
+    def _setup_pool_monitoring(self, engine: Engine | AsyncEngine, engine_name: str) -> None:
+        """Set up database pool monitoring for the given engine."""
+        if not settings.enable_db_pool_monitoring:
+            return
+
+        try:
+            from letta.otel.db_pool_monitoring import setup_pool_monitoring
+
+            setup_pool_monitoring(engine, engine_name)
+            self.logger.info(f"Database pool monitoring enabled for {engine_name}")
+        except ImportError:
+            self.logger.warning("Database pool monitoring not available - missing dependencies")
+        except Exception as e:
+            self.logger.warning(f"Failed to setup pool monitoring for {engine_name}: {e}")
+
     def get_engine(self, name: str = "default") -> Engine:
         """Get a database engine by name."""
         self.initialize_sync()
         return self._engines.get(name)
 
+    def get_async_engine(self, name: str = "default") -> Engine:
+        """Get a database engine by name."""
+        self.initialize_async()
+        return self._async_engines.get(name)
+
     def get_session_factory(self, name: str = "default") -> sessionmaker:
         """Get a session factory by name."""
         self.initialize_sync()
@@ -286,6 +385,11 @@ class DatabaseRegistry:
 db_registry = DatabaseRegistry()
 
 
+def get_db_registry() -> DatabaseRegistry:
+    """Get the global database registry instance."""
+    return db_registry
+
+
 def get_db():
     """Get a database session."""
     with db_registry.session() as session: