letta-nightly 0.6.4.dev20241217104233__py3-none-any.whl → 0.6.5.dev20241218055539__py3-none-any.whl

letta/orm/sqlalchemy_base.py

@@ -1,14 +1,16 @@
 from datetime import datetime
 from enum import Enum
+from functools import wraps
 from typing import TYPE_CHECKING, List, Literal, Optional
 
 from sqlalchemy import String, desc, func, or_, select
-from sqlalchemy.exc import DBAPIError, IntegrityError
+from sqlalchemy.exc import DBAPIError, IntegrityError, TimeoutError
 from sqlalchemy.orm import Mapped, Session, mapped_column
 
 from letta.log import get_logger
 from letta.orm.base import Base, CommonSqlalchemyMetaMixins
 from letta.orm.errors import (
+    DatabaseTimeoutError,
     ForeignKeyConstraintViolationError,
     NoResultFound,
     UniqueConstraintViolationError,
@@ -23,6 +25,20 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def handle_db_timeout(func):
+    """Decorator to handle SQLAlchemy TimeoutError and wrap it in a custom exception."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except TimeoutError as e:
+            logger.error(f"Timeout while executing {func.__name__} with args {args} and kwargs {kwargs}: {e}")
+            raise DatabaseTimeoutError(message=f"Timeout occurred in {func.__name__}.", original_exception=e)
+
+    return wrapper
+
+
 class AccessType(str, Enum):
     ORGANIZATION = "organization"
     USER = "user"
@@ -36,22 +52,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
     id: Mapped[str] = mapped_column(String, primary_key=True)
 
     @classmethod
-    def get(cls, *, db_session: Session, id: str) -> Optional["SqlalchemyBase"]:
-        """Get a record by ID.
-
-        Args:
-            db_session: SQLAlchemy session
-            id: Record ID to retrieve
-
-        Returns:
-            Optional[SqlalchemyBase]: The record if found, None otherwise
-        """
-        try:
-            return db_session.query(cls).filter(cls.id == id).first()
-        except DBAPIError:
-            return None
-
-    @classmethod
+    @handle_db_timeout
     def list(
         cls,
         *,
@@ -180,6 +181,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             return list(session.execute(query).scalars())
 
     @classmethod
+    @handle_db_timeout
     def read(
         cls,
         db_session: "Session",
@@ -231,6 +233,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         conditions_str = ", ".join(query_conditions) if query_conditions else "no specific conditions"
         raise NoResultFound(f"{cls.__name__} not found with {conditions_str}")
 
+    @handle_db_timeout
     def create(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Creating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
 
@@ -245,6 +248,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         except (DBAPIError, IntegrityError) as e:
             self._handle_dbapi_error(e)
 
+    @handle_db_timeout
     def delete(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Soft deleting {self.__class__.__name__} with ID: {self.id} with actor={actor}")
 
@@ -254,6 +258,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
         self.is_deleted = True
         return self.update(db_session)
 
+    @handle_db_timeout
     def hard_delete(self, db_session: "Session", actor: Optional["User"] = None) -> None:
         """Permanently removes the record from the database."""
         logger.debug(f"Hard deleting {self.__class__.__name__} with ID: {self.id} with actor={actor}")
@@ -269,6 +274,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             else:
                 logger.debug(f"{self.__class__.__name__} with ID {self.id} successfully hard deleted")
 
+    @handle_db_timeout
     def update(self, db_session: "Session", actor: Optional["User"] = None) -> "SqlalchemyBase":
         logger.debug(f"Updating {self.__class__.__name__} with ID: {self.id} with actor={actor}")
         if actor:
@@ -281,6 +287,7 @@ class SqlalchemyBase(CommonSqlalchemyMetaMixins, Base):
             return self
 
     @classmethod
+    @handle_db_timeout
     def size(
         cls,
         *,

letta/providers.py

@@ -13,6 +13,7 @@ from letta.schemas.llm_config import LLMConfig
 
 
 class Provider(BaseModel):
+    name: str = Field(..., description="The name of the provider")
 
     def list_llm_models(self) -> List[LLMConfig]:
         return []
@@ -465,6 +466,7 @@ class TogetherProvider(OpenAIProvider):
 
 class GoogleAIProvider(Provider):
     # gemini
+    name: str = "google_ai"
     api_key: str = Field(..., description="API key for the Google AI API.")
     base_url: str = "https://generativelanguage.googleapis.com"
 

letta/schemas/agent.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Optional
 
 from pydantic import BaseModel, Field, field_validator
 
+from letta.constants import DEFAULT_EMBEDDING_CHUNK_SIZE
 from letta.schemas.block import CreateBlock
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.letta_base import OrmMetadataBase
@@ -107,6 +108,16 @@ class CreateAgent(BaseModel, validate_assignment=True):  #
     include_base_tools: bool = Field(True, description="The LLM configuration used by the agent.")
     description: Optional[str] = Field(None, description="The description of the agent.")
     metadata_: Optional[Dict] = Field(None, description="The metadata of the agent.", alias="metadata_")
+    llm: Optional[str] = Field(
+        None,
+        description="The LLM configuration handle used by the agent, specified in the format "
+        "provider/model-name, as an alternative to specifying llm_config.",
+    )
+    embedding: Optional[str] = Field(
+        None, description="The embedding configuration handle used by the agent, specified in the format provider/model-name."
+    )
+    context_window_limit: Optional[int] = Field(None, description="The context window limit used by the agent.")
+    embedding_chunk_size: Optional[int] = Field(DEFAULT_EMBEDDING_CHUNK_SIZE, description="The embedding chunk size used by the agent.")
 
     @field_validator("name")
     @classmethod
@@ -133,6 +144,30 @@ class CreateAgent(BaseModel, validate_assignment=True):  #
 
         return name
 
+    @field_validator("llm")
+    @classmethod
+    def validate_llm(cls, llm: Optional[str]) -> Optional[str]:
+        if not llm:
+            return llm
+
+        provider_name, model_name = llm.split("/", 1)
+        if not provider_name or not model_name:
+            raise ValueError("The llm config handle should be in the format provider/model-name")
+
+        return llm
+
+    @field_validator("embedding")
+    @classmethod
+    def validate_embedding(cls, embedding: Optional[str]) -> Optional[str]:
+        if not embedding:
+            return embedding
+
+        provider_name, model_name = embedding.split("/", 1)
+        if not provider_name or not model_name:
+            raise ValueError("The embedding config handle should be in the format provider/model-name")
+
+        return embedding
+
 
 class UpdateAgent(BaseModel):
     name: Optional[str] = Field(None, description="The name of the agent.")

letta/schemas/sandbox_config.py

@@ -1,7 +1,7 @@
 import hashlib
 import json
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel, Field, model_validator
 
@@ -21,6 +21,7 @@ class SandboxRunResult(BaseModel):
     agent_state: Optional[AgentState] = Field(None, description="The agent state")
     stdout: Optional[List[str]] = Field(None, description="Captured stdout (e.g. prints, logs) from the function invocation")
     stderr: Optional[List[str]] = Field(None, description="Captured stderr from the function invocation")
+    status: Literal["success", "error"] = Field(..., description="The status of the tool execution and return object")
     sandbox_config_fingerprint: str = Field(None, description="The fingerprint of the config for the sandbox")
 
 

letta/server/rest_api/app.py

@@ -15,7 +15,12 @@ from letta.__init__ import __version__
 from letta.constants import ADMIN_PREFIX, API_PREFIX, OPENAI_API_PREFIX
 from letta.errors import LettaAgentNotFoundError, LettaUserNotFoundError
 from letta.log import get_logger
-from letta.orm.errors import NoResultFound
+from letta.orm.errors import (
+    DatabaseTimeoutError,
+    ForeignKeyConstraintViolationError,
+    NoResultFound,
+    UniqueConstraintViolationError,
+)
 from letta.schemas.letta_response import LettaResponse
 from letta.server.constants import REST_DEFAULT_PORT
 
@@ -175,7 +180,6 @@ def create_application() -> "FastAPI":
 
     @app.exception_handler(NoResultFound)
     async def no_result_found_handler(request: Request, exc: NoResultFound):
-        logger.error(f"NoResultFound request: {request}")
         logger.error(f"NoResultFound: {exc}")
 
         return JSONResponse(
@@ -183,6 +187,32 @@ def create_application() -> "FastAPI":
             content={"detail": str(exc)},
         )
 
+    @app.exception_handler(ForeignKeyConstraintViolationError)
+    async def foreign_key_constraint_handler(request: Request, exc: ForeignKeyConstraintViolationError):
+        logger.error(f"ForeignKeyConstraintViolationError: {exc}")
+
+        return JSONResponse(
+            status_code=409,
+            content={"detail": str(exc)},
+        )
+
+    @app.exception_handler(UniqueConstraintViolationError)
+    async def unique_key_constraint_handler(request: Request, exc: UniqueConstraintViolationError):
+        logger.error(f"UniqueConstraintViolationError: {exc}")
+
+        return JSONResponse(
+            status_code=409,
+            content={"detail": str(exc)},
+        )
+
+    @app.exception_handler(DatabaseTimeoutError)
+    async def database_timeout_error_handler(request: Request, exc: DatabaseTimeoutError):
+        logger.error(f"Timeout occurred: {exc}. Original exception: {exc.original_exception}")
+        return JSONResponse(
+            status_code=503,
+            content={"detail": "The database is temporarily unavailable. Please try again later."},
+        )
+
     @app.exception_handler(ValueError)
     async def value_error_handler(request: Request, exc: ValueError):
         return JSONResponse(status_code=400, content={"detail": str(exc)})
@@ -235,11 +265,6 @@ def create_application() -> "FastAPI":
 
     @app.on_event("startup")
     def on_startup():
-        # load the default tools
-        # from letta.orm.tool import Tool
-
-        # Tool.load_default_tools(get_db_session())
-
         generate_openapi_schema(app)
 
     @app.on_event("shutdown")

letta/server/rest_api/routers/v1/tools.py

@@ -195,7 +195,7 @@ def run_tool_from_source(
             tool_source_type=request.source_type,
             tool_args=request.args,
             tool_name=request.name,
-            user_id=actor.id,
+            actor=actor,
         )
     except LettaToolCreateError as e:
         # HTTP 400 == Bad Request

letta/server/server.py

@@ -4,7 +4,6 @@ import os
 import traceback
 import warnings
 from abc import abstractmethod
-from asyncio import Lock
 from datetime import datetime
 from typing import Callable, List, Optional, Tuple, Union
 
@@ -75,7 +74,7 @@ from letta.services.source_manager import SourceManager
 from letta.services.tool_execution_sandbox import ToolExecutionSandbox
 from letta.services.tool_manager import ToolManager
 from letta.services.user_manager import UserManager
-from letta.utils import get_utc_time, json_dumps, json_loads
+from letta.utils import get_friendly_error_msg, get_utc_time, json_dumps, json_loads
 
 logger = get_logger(__name__)
 
@@ -191,7 +190,14 @@ if settings.letta_pg_uri_no_default:
     config.archival_storage_uri = settings.letta_pg_uri_no_default
 
     # create engine
-    engine = create_engine(settings.letta_pg_uri)
+    engine = create_engine(
+        settings.letta_pg_uri,
+        pool_size=settings.pg_pool_size,
+        max_overflow=settings.pg_max_overflow,
+        pool_timeout=settings.pg_pool_timeout,
+        pool_recycle=settings.pg_pool_recycle,
+        echo=settings.pg_echo,
+    )
 else:
     # TODO: don't rely on config storage
     engine = create_engine("sqlite:///" + os.path.join(config.recall_storage_path, "sqlite.db"))
@@ -265,9 +271,6 @@ class SyncServer(Server):
 
         self.credentials = LettaCredentials.load()
 
-        # Locks
-        self.send_message_lock = Lock()
-
         # Initialize the metadata store
         config = LettaConfig.load()
         if settings.letta_pg_uri_no_default:
@@ -773,6 +776,18 @@ class SyncServer(Server):
         # interface
         interface: Union[AgentInterface, None] = None,
     ) -> AgentState:
+        if request.llm_config is None:
+            if request.llm is None:
+                raise ValueError("Must specify either llm or llm_config in request")
+            request.llm_config = self.get_llm_config_from_handle(handle=request.llm, context_window_limit=request.context_window_limit)
+
+        if request.embedding_config is None:
+            if request.embedding is None:
+                raise ValueError("Must specify either embedding or embedding_config in request")
+            request.embedding_config = self.get_embedding_config_from_handle(
+                handle=request.embedding, embedding_chunk_size=request.embedding_chunk_size or constants.DEFAULT_EMBEDDING_CHUNK_SIZE
+            )
+
         """Create a new agent using a config"""
         # Invoke manager
         agent_state = self.agent_manager.create_agent(
@@ -821,7 +836,7 @@ class SyncServer(Server):
     ) -> AgentState:
         """Update the agents core memory block, return the new state"""
         # Update agent state in the db first
-        self.agent_manager.update_agent(agent_id=agent_id, agent_update=request, actor=actor)
+        agent_state = self.agent_manager.update_agent(agent_id=agent_id, agent_update=request, actor=actor)
 
         # Get the agent object (loaded in memory)
         letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
@@ -838,13 +853,9 @@ class SyncServer(Server):
             # then (2) setting the attributes ._messages and .state.message_ids
             letta_agent.set_message_buffer(message_ids=request.message_ids)
 
-        # tools
-        if request.tool_ids:
-            letta_agent.link_tools(letta_agent.agent_state.tools)
-
         letta_agent.update_state()
 
-        return letta_agent.agent_state
+        return agent_state
 
     def get_tools_from_agent(self, agent_id: str, user_id: Optional[str]) -> List[Tool]:
         """Get tools from an existing agent"""
@@ -867,11 +878,6 @@ class SyncServer(Server):
 
         agent_state = self.agent_manager.attach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
 
-        # TODO: This is very redundant, and should probably be simplified
-        # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        letta_agent.link_tools(agent_state.tools)
-
         return agent_state
 
     def remove_tool_from_agent(
@@ -885,10 +891,6 @@ class SyncServer(Server):
         actor = self.user_manager.get_user_or_default(user_id=user_id)
         agent_state = self.agent_manager.detach_tool(agent_id=agent_id, tool_id=tool_id, actor=actor)
 
-        # Get the agent object (loaded in memory)
-        letta_agent = self.load_agent(agent_id=agent_id, actor=actor)
-        letta_agent.link_tools(agent_state.tools)
-
         return agent_state
 
     # convert name->id
@@ -1280,6 +1282,55 @@ class SyncServer(Server):
                 warnings.warn(f"An error occurred while listing embedding models for provider {provider}: {e}")
         return embedding_models
 
+    def get_llm_config_from_handle(self, handle: str, context_window_limit: Optional[int] = None) -> LLMConfig:
+        provider_name, model_name = handle.split("/", 1)
+        provider = self.get_provider_from_name(provider_name)
+
+        llm_configs = [config for config in provider.list_llm_models() if config.model == model_name]
+        if not llm_configs:
+            raise ValueError(f"LLM model {model_name} is not supported by {provider_name}")
+        elif len(llm_configs) > 1:
+            raise ValueError(f"Multiple LLM models with name {model_name} supported by {provider_name}")
+        else:
+            llm_config = llm_configs[0]
+
+        if context_window_limit:
+            if context_window_limit > llm_config.context_window:
+                raise ValueError(f"Context window limit ({context_window_limit}) is greater than maximum of ({llm_config.context_window})")
+            llm_config.context_window = context_window_limit
+
+        return llm_config
+
+    def get_embedding_config_from_handle(
+        self, handle: str, embedding_chunk_size: int = constants.DEFAULT_EMBEDDING_CHUNK_SIZE
+    ) -> EmbeddingConfig:
+        provider_name, model_name = handle.split("/", 1)
+        provider = self.get_provider_from_name(provider_name)
+
+        embedding_configs = [config for config in provider.list_embedding_models() if config.embedding_model == model_name]
+        if not embedding_configs:
+            raise ValueError(f"Embedding model {model_name} is not supported by {provider_name}")
+        elif len(embedding_configs) > 1:
+            raise ValueError(f"Multiple embedding models with name {model_name} supported by {provider_name}")
+        else:
+            embedding_config = embedding_configs[0]
+
+        if embedding_chunk_size:
+            embedding_config.embedding_chunk_size = embedding_chunk_size
+
+        return embedding_config
+
+    def get_provider_from_name(self, provider_name: str) -> Provider:
+        providers = [provider for provider in self._enabled_providers if provider.name == provider_name]
+        if not providers:
+            raise ValueError(f"Provider {provider_name} is not supported")
+        elif len(providers) > 1:
+            raise ValueError(f"Multiple providers with name {provider_name} supported")
+        else:
+            provider = providers[0]
+
+        return provider
+
     def add_llm_model(self, request: LLMConfig) -> LLMConfig:
         """Add a new LLM model"""
 
@@ -1300,7 +1351,7 @@ class SyncServer(Server):
 
     def run_tool_from_source(
         self,
-        user_id: str,
+        actor: User,
         tool_args: str,
         tool_source: str,
         tool_source_type: Optional[str] = None,
@@ -1328,56 +1379,29 @@ class SyncServer(Server):
 
         # Next, attempt to run the tool with the sandbox
         try:
-            sandbox_run_result = ToolExecutionSandbox(tool.name, tool_args_dict, user_id, tool_object=tool).run(agent_state=agent_state)
-            function_response = str(sandbox_run_result.func_return)
-            stdout = [s for s in sandbox_run_result.stdout if s.strip()]
-            stderr = [s for s in sandbox_run_result.stderr if s.strip()]
-
-            # expected error
-            if stderr:
-                error_msg = self.get_error_msg_for_func_return(tool.name, stderr[-1])
-                return FunctionReturn(
-                    id="null",
-                    function_call_id="null",
-                    date=get_utc_time(),
-                    status="error",
-                    function_return=error_msg,
-                    stdout=stdout,
-                    stderr=stderr,
-                )
-
+            sandbox_run_result = ToolExecutionSandbox(tool.name, tool_args_dict, actor, tool_object=tool).run(agent_state=agent_state)
             return FunctionReturn(
                 id="null",
                 function_call_id="null",
                 date=get_utc_time(),
-                status="success",
-                function_return=function_response,
-                stdout=stdout,
-                stderr=stderr,
+                status=sandbox_run_result.status,
+                function_return=str(sandbox_run_result.func_return),
+                stdout=sandbox_run_result.stdout,
+                stderr=sandbox_run_result.stderr,
             )
 
-        # unexpected error TODO(@cthomas): consolidate error handling
         except Exception as e:
-            error_msg = self.get_error_msg_for_func_return(tool.name, e)
+            func_return = get_friendly_error_msg(function_name=tool.name, exception_name=type(e).__name__, exception_message=str(e))
             return FunctionReturn(
                 id="null",
                 function_call_id="null",
                 date=get_utc_time(),
                 status="error",
-                function_return=error_msg,
-                stdout=[""],
+                function_return=func_return,
+                stdout=[],
                 stderr=[traceback.format_exc()],
             )
 
-    def get_error_msg_for_func_return(self, tool_name, exception_message):
-        # same as agent.py
-        from letta.constants import MAX_ERROR_MESSAGE_CHAR_LIMIT
-
-        error_msg = f"Error executing tool {tool_name}: {exception_message}"
-        if len(error_msg) > MAX_ERROR_MESSAGE_CHAR_LIMIT:
-            error_msg = error_msg[:MAX_ERROR_MESSAGE_CHAR_LIMIT]
-        return error_msg
-
     # Composio wrappers
     def get_composio_client(self, api_key: Optional[str] = None):
         if api_key:

letta/services/agent_manager.py

@@ -61,6 +61,9 @@ class AgentManager:
     ) -> PydanticAgentState:
         system = derive_system_message(agent_type=agent_create.agent_type, system=agent_create.system)
 
+        if not agent_create.llm_config or not agent_create.embedding_config:
+            raise ValueError("llm_config and embedding_config are required")
+
         # create blocks (note: cannot be linked into the agent_id is created)
         block_ids = list(agent_create.block_ids or [])  # Create a local copy to avoid modifying the original
         for create_block in agent_create.memory_blocks: