letta-nightly 0.6.53.dev20250418104238__py3-none-any.whl → 0.6.54.dev20250419104029__py3-none-any.whl

letta/schemas/providers.py

@@ -228,63 +228,6 @@ class OpenAIProvider(Provider):
             return LLM_MAX_TOKENS["DEFAULT"]
 
 
-class xAIProvider(OpenAIProvider):
-    """https://docs.x.ai/docs/api-reference"""
-
-    name: str = "xai"
-    api_key: str = Field(..., description="API key for the xAI/Grok API.")
-    base_url: str = Field("https://api.x.ai/v1", description="Base URL for the xAI/Grok API.")
-
-    def get_model_context_window_size(self, model_name: str) -> Optional[int]:
-        # xAI doesn't return context window in the model listing,
-        # so these are hardcoded from their website
-        if model_name == "grok-2-1212":
-            return 131072
-        else:
-            return None
-
-    def list_llm_models(self) -> List[LLMConfig]:
-        from letta.llm_api.openai import openai_get_model_list
-
-        response = openai_get_model_list(self.base_url, api_key=self.api_key)
-
-        if "data" in response:
-            data = response["data"]
-        else:
-            data = response
-
-        configs = []
-        for model in data:
-            assert "id" in model, f"xAI/Grok model missing 'id' field: {model}"
-            model_name = model["id"]
-
-            # In case xAI starts supporting it in the future:
-            if "context_length" in model:
-                context_window_size = model["context_length"]
-            else:
-                context_window_size = self.get_model_context_window_size(model_name)
-
-            if not context_window_size:
-                warnings.warn(f"Couldn't find context window size for model {model_name}")
-                continue
-
-            configs.append(
-                LLMConfig(
-                    model=model_name,
-                    model_endpoint_type="xai",
-                    model_endpoint=self.base_url,
-                    context_window=context_window_size,
-                    handle=self.get_handle(model_name),
-                )
-            )
-
-        return configs
-
-    def list_embedding_models(self) -> List[EmbeddingConfig]:
-        # No embeddings supported
-        return []
-
-
 class DeepSeekProvider(OpenAIProvider):
     """
     DeepSeek ChatCompletions API is similar to OpenAI's reasoning API,
@@ -478,7 +421,7 @@ class LMStudioOpenAIProvider(OpenAIProvider):
         return configs
 
 
-class xAIProvider(OpenAIProvider):
+class XAIProvider(OpenAIProvider):
     """https://docs.x.ai/docs/api-reference"""
 
     name: str = "xai"
@@ -490,6 +433,15 @@ class xAIProvider(OpenAIProvider):
         # so these are hardcoded from their website
         if model_name == "grok-2-1212":
             return 131072
+        # NOTE: disabling the minis for now since they return weird MM parts
+        # elif model_name == "grok-3-mini-fast-beta":
+        #     return 131072
+        # elif model_name == "grok-3-mini-beta":
+        #     return 131072
+        elif model_name == "grok-3-fast-beta":
+            return 131072
+        elif model_name == "grok-3-beta":
+            return 131072
         else:
             return None
 

letta/serialize_schemas/marshmallow_agent.py

@@ -4,6 +4,7 @@ from marshmallow import fields, post_dump, pre_load
 
 import letta
 from letta.orm import Agent
+from letta.orm import Message as MessageModel
 from letta.schemas.agent import AgentState as PydanticAgentState
 from letta.schemas.user import User
 from letta.serialize_schemas.marshmallow_agent_environment_variable import SerializedAgentEnvironmentVariableSchema
@@ -35,7 +36,6 @@ class MarshmallowAgentSchema(BaseSchema):
 
     tool_rules = ToolRulesField()
 
-    messages = fields.List(fields.Nested(SerializedMessageSchema))
     core_memory = fields.List(fields.Nested(SerializedBlockSchema))
     tools = fields.List(fields.Nested(SerializedToolSchema))
     tool_exec_environment_variables = fields.List(fields.Nested(SerializedAgentEnvironmentVariableSchema))
@@ -54,6 +54,30 @@ class MarshmallowAgentSchema(BaseSchema):
                 field.schema.session = session
                 field.schema.actor = actor
 
+    @post_dump
+    def attach_messages(self, data: Dict, **kwargs):
+        """
+        After dumping the agent, load all its Message rows and serialize them here.
+        """
+        # TODO: This is hacky, but want to move fast, please refactor moving forward
+        from letta.server.db import db_context as session_maker
+
+        with session_maker() as session:
+            agent_id = data.get("id")
+            msgs = (
+                session.query(MessageModel)
+                .filter(
+                    MessageModel.agent_id == agent_id,
+                    MessageModel.organization_id == self.actor.organization_id,
+                )
+                .order_by(MessageModel.sequence_id.asc())
+                .all()
+            )
+            # overwrite the “messages” key with a fully serialized list
+            data[self.FIELD_MESSAGES] = [SerializedMessageSchema(session=self.session, actor=self.actor).dump(m) for m in msgs]
+
+        return data
+
     @post_dump
     def sanitize_ids(self, data: Dict, **kwargs):
         """
@@ -101,25 +125,6 @@ class MarshmallowAgentSchema(BaseSchema):
         del data[self.FIELD_VERSION]
         return data
 
-    @pre_load
-    def remap_in_context_messages(self, data, **kwargs):
-        """
-        Restores `message_ids` by collecting message IDs where `in_context` is True,
-        generates new IDs for all messages, and removes `in_context` from all messages.
-        """
-        messages = data.get(self.FIELD_MESSAGES, [])
-        for msg in messages:
-            msg[self.FIELD_ID] = SerializedMessageSchema.generate_id()  # Generate new ID
-
-        message_ids = []
-        in_context_message_indices = data.pop(self.FIELD_IN_CONTEXT_INDICES)
-        for idx in in_context_message_indices:
-            message_ids.append(messages[idx][self.FIELD_ID])
-
-        data[self.FIELD_MESSAGE_IDS] = message_ids
-
-        return data
-
     class Meta(BaseSchema.Meta):
         model = Agent
         exclude = BaseSchema.Meta.exclude + (
@@ -127,8 +132,6 @@ class MarshmallowAgentSchema(BaseSchema):
             "template_id",
             "base_template_id",
             "sources",
-            "source_passages",
-            "agent_passages",
             "identities",
             "is_deleted",
             "groups",

letta/serialize_schemas/marshmallow_message.py

@@ -39,4 +39,4 @@ class SerializedMessageSchema(BaseSchema):
 
     class Meta(BaseSchema.Meta):
         model = Message
-        exclude = BaseSchema.Meta.exclude + ("step", "job_message", "agent", "otid", "is_deleted")
+        exclude = BaseSchema.Meta.exclude + ("step", "job_message", "otid", "is_deleted")

letta/server/db.py

@@ -1,4 +1,5 @@
 import os
+import threading
 from contextlib import contextmanager
 
 from rich.console import Console
@@ -10,13 +11,17 @@ from sqlalchemy.orm import sessionmaker
 from letta.config import LettaConfig
 from letta.log import get_logger
 from letta.orm import Base
-
-# NOTE: hack to see if single session management works
 from letta.settings import settings
 
-config = LettaConfig.load()
+# Use globals for the lock and initialization flag
+_engine_lock = threading.Lock()
+_engine_initialized = False
 
+# Create variables in global scope but don't initialize them yet
+config = LettaConfig.load()
 logger = get_logger(__name__)
+engine = None
+SessionLocal = None
 
 
 def print_sqlite_schema_error():
@@ -49,59 +54,80 @@ def db_error_handler():
         exit(1)
 
 
-if settings.letta_pg_uri_no_default:
-    print("Creating postgres engine")
-    config.recall_storage_type = "postgres"
-    config.recall_storage_uri = settings.letta_pg_uri_no_default
-    config.archival_storage_type = "postgres"
-    config.archival_storage_uri = settings.letta_pg_uri_no_default
-
-    # create engine
-    engine = create_engine(
-        settings.letta_pg_uri,
-        # f"{settings.letta_pg_uri}?options=-c%20client_encoding=UTF8",
-        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,
-        # connect_args={"client_encoding": "utf8"},
-    )
-else:
-    # TODO: don't rely on config storage
-    engine_path = "sqlite:///" + os.path.join(config.recall_storage_path, "sqlite.db")
-    logger.info("Creating sqlite engine " + engine_path)
-
-    engine = create_engine(engine_path)
-
-    # Store the original connect method
-    original_connect = engine.connect
-
-    def wrapped_connect(*args, **kwargs):
-        with db_error_handler():
-            # Get the connection
-            connection = original_connect(*args, **kwargs)
-
-            # Store the original execution method
-            original_execute = connection.execute
-
-            # Wrap the execute method of the connection
-            def wrapped_execute(*args, **kwargs):
+def initialize_engine():
+    """Initialize the database engine only when needed."""
+    global engine, SessionLocal, _engine_initialized
+
+    with _engine_lock:
+        # Check again inside the lock to prevent race conditions
+        if _engine_initialized:
+            return
+
+        if settings.letta_pg_uri_no_default:
+            logger.info("Creating postgres engine")
+            config.recall_storage_type = "postgres"
+            config.recall_storage_uri = settings.letta_pg_uri_no_default
+            config.archival_storage_type = "postgres"
+            config.archival_storage_uri = settings.letta_pg_uri_no_default
+
+            # create engine
+            engine = create_engine(
+                settings.letta_pg_uri,
+                # f"{settings.letta_pg_uri}?options=-c%20client_encoding=UTF8",
+                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,
+                # connect_args={"client_encoding": "utf8"},
+            )
+        else:
+            # TODO: don't rely on config storage
+            engine_path = "sqlite:///" + os.path.join(config.recall_storage_path, "sqlite.db")
+            logger.info("Creating sqlite engine " + engine_path)
+
+            engine = create_engine(engine_path)
+
+            # Store the original connect method
+            original_connect = engine.connect
+
+            def wrapped_connect(*args, **kwargs):
                 with db_error_handler():
-                    return original_execute(*args, **kwargs)
+                    # Get the connection
+                    connection = original_connect(*args, **kwargs)
+
+                    # Store the original execution method
+                    original_execute = connection.execute
 
-            # Replace the connection's execute method
-            connection.execute = wrapped_execute
+                    # Wrap the execute method of the connection
+                    def wrapped_execute(*args, **kwargs):
+                        with db_error_handler():
+                            return original_execute(*args, **kwargs)
 
-            return connection
+                    # Replace the connection's execute method
+                    connection.execute = wrapped_execute
 
-    # Replace the engine's connect method
-    engine.connect = wrapped_connect
+                    return connection
 
-    Base.metadata.create_all(bind=engine)
+            # Replace the engine's connect method
+            engine.connect = wrapped_connect
+
+            Base.metadata.create_all(bind=engine)
+
+        # Create the session factory
+        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+        _engine_initialized = True
 
 
 def get_db():
+    """Get a database session, initializing the engine if needed."""
+    global engine, SessionLocal
+
+    # Make sure engine is initialized
+    if not _engine_initialized:
+        initialize_engine()
+
+    # Now SessionLocal should be defined and callable
     db = SessionLocal()
     try:
         yield db
@@ -109,5 +135,5 @@ def get_db():
         db.close()
 
 
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+# Define db_context as a context manager that uses get_db
 db_context = contextmanager(get_db)

letta/server/rest_api/app.py

@@ -174,6 +174,7 @@ def create_application() -> "FastAPI":
     async def generic_error_handler(request: Request, exc: Exception):
         # Log the actual error for debugging
         log.error(f"Unhandled error: {exc}", exc_info=True)
+        print(f"Unhandled error: {exc}")
 
         # Print the stack trace
         print(f"Stack trace: {exc}")

letta/server/rest_api/interface.py

@@ -6,6 +6,8 @@ from collections import deque
 from datetime import datetime
 from typing import AsyncGenerator, Literal, Optional, Union
 
+import demjson3 as demjson
+
 from letta.constants import DEFAULT_MESSAGE_TOOL, DEFAULT_MESSAGE_TOOL_KWARG
 from letta.helpers.datetime_helpers import is_utc_datetime
 from letta.interface import AgentInterface
@@ -502,7 +504,7 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
                 date=message_date,
                 reasoning=message_delta.reasoning_content,
                 signature=message_delta.reasoning_content_signature,
-                source="reasoner_model" if message_delta.reasoning_content_signature else "non_reasoner_model",
+                source="reasoner_model" if message_delta.reasoning_content else "non_reasoner_model",
                 name=name,
                 otid=otid,
             )
@@ -530,7 +532,6 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
             try:
                 # NOTE: this is hardcoded for our DeepSeek API integration
                 json_reasoning_content = parse_json(self.expect_reasoning_content_buffer)
-                print(f"json_reasoning_content: {json_reasoning_content}")
 
                 processed_chunk = ToolCallMessage(
                     id=message_id,
@@ -547,6 +548,10 @@ class StreamingServerInterface(AgentChunkStreamingInterface):
             except json.JSONDecodeError as e:
                 print(f"Failed to interpret reasoning content ({self.expect_reasoning_content_buffer}) as JSON: {e}")
 
+                return None
+            except demjson.JSONDecodeError as e:
+                print(f"Failed to interpret reasoning content ({self.expect_reasoning_content_buffer}) as JSON: {e}")
+
                 return None
             # Else,
             # return None

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

@@ -5,6 +5,7 @@ from letta.server.rest_api.routers.v1.health import router as health_router
 from letta.server.rest_api.routers.v1.identities import router as identities_router
 from letta.server.rest_api.routers.v1.jobs import router as jobs_router
 from letta.server.rest_api.routers.v1.llms import router as llm_router
+from letta.server.rest_api.routers.v1.messages import router as messages_router
 from letta.server.rest_api.routers.v1.providers import router as providers_router
 from letta.server.rest_api.routers.v1.runs import router as runs_router
 from letta.server.rest_api.routers.v1.sandbox_configs import router as sandbox_configs_router
@@ -29,5 +30,6 @@ ROUTERS = [
     runs_router,
     steps_router,
     tags_router,
+    messages_router,
     voice_router,
 ]

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

@@ -1,6 +1,6 @@
 import json
 import traceback
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Annotated, Any, List, Optional
 
 from fastapi import APIRouter, BackgroundTasks, Body, Depends, File, Header, HTTPException, Query, UploadFile, status
@@ -17,6 +17,7 @@ from letta.log import get_logger
 from letta.orm.errors import NoResultFound
 from letta.schemas.agent import AgentState, AgentType, CreateAgent, UpdateAgent
 from letta.schemas.block import Block, BlockUpdate
+from letta.schemas.group import Group
 from letta.schemas.job import JobStatus, JobUpdate, LettaRequestConfig
 from letta.schemas.letta_message import LettaMessageUnion, LettaMessageUpdateUnion
 from letta.schemas.letta_request import LettaRequest, LettaStreamingRequest
@@ -173,7 +174,7 @@ async def import_agent_serialized(
         raise HTTPException(status_code=400, detail="Corrupted agent file format.")
 
     except ValidationError as e:
-        raise HTTPException(status_code=422, detail=f"Invalid agent schema: {e.errors()}")
+        raise HTTPException(status_code=422, detail=f"Invalid agent schema: {str(e)}")
 
     except IntegrityError as e:
         raise HTTPException(status_code=409, detail=f"Database integrity error: {str(e)}")
@@ -282,6 +283,7 @@ def detach_tool(
 def attach_source(
     agent_id: str,
     source_id: str,
+    background_tasks: BackgroundTasks,
     server: "SyncServer" = Depends(get_letta_server),
     actor_id: Optional[str] = Header(None, alias="user_id"),
 ):
@@ -289,7 +291,11 @@ def attach_source(
     Attach a source to an agent.
     """
     actor = server.user_manager.get_user_or_default(user_id=actor_id)
-    return server.agent_manager.attach_source(agent_id=agent_id, source_id=source_id, actor=actor)
+    agent = server.agent_manager.attach_source(agent_id=agent_id, source_id=source_id, actor=actor)
+    if agent.enable_sleeptime:
+        source = server.source_manager.get_source_by_id(source_id=source_id)
+        background_tasks.add_task(server.sleeptime_document_ingest, agent, source, actor)
+    return agent
 
 
 @router.patch("/{agent_id}/sources/detach/{source_id}", response_model=AgentState, operation_id="detach_source_from_agent")
@@ -303,7 +309,15 @@ def detach_source(
     Detach a source from an agent.
     """
     actor = server.user_manager.get_user_or_default(user_id=actor_id)
-    return server.agent_manager.detach_source(agent_id=agent_id, source_id=source_id, actor=actor)
+    agent = server.agent_manager.detach_source(agent_id=agent_id, source_id=source_id, actor=actor)
+    if agent.enable_sleeptime:
+        try:
+            source = server.source_manager.get_source_by_id(source_id=source_id)
+            block = server.agent_manager.get_block_with_label(agent_id=agent.id, block_label=source.name, actor=actor)
+            server.block_manager.delete_block(block.id, actor)
+        except:
+            pass
+    return agent
 
 
 @router.get("/{agent_id}", response_model=AgentState, operation_id="retrieve_agent")
@@ -728,7 +742,7 @@ async def process_message_background(
         # Update job status to completed
         job_update = JobUpdate(
             status=JobStatus.completed,
-            completed_at=datetime.utcnow(),
+            completed_at=datetime.now(timezone.utc),
             metadata={"result": result.model_dump(mode="json")},  # Store the result in metadata
         )
         server.job_manager.update_job_by_id(job_id=job_id, job_update=job_update, actor=actor)
@@ -737,7 +751,7 @@ async def process_message_background(
         # Update job status to failed
         job_update = JobUpdate(
             status=JobStatus.failed,
-            completed_at=datetime.utcnow(),
+            completed_at=datetime.now(timezone.utc),
             metadata={"error": str(e)},
         )
         server.job_manager.update_job_by_id(job_id=job_id, job_update=job_update, actor=actor)
@@ -804,3 +818,16 @@ def reset_messages(
     """Resets the messages for an agent"""
     actor = server.user_manager.get_user_or_default(user_id=actor_id)
     return server.agent_manager.reset_messages(agent_id=agent_id, actor=actor, add_default_initial_messages=add_default_initial_messages)
+
+
+@router.get("/{agent_id}/groups", response_model=List[Group], operation_id="list_agent_groups")
+async def list_agent_groups(
+    agent_id: str,
+    manager_type: Optional[str] = Query(None, description="Manager type to filter groups by"),
+    server: "SyncServer" = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),  # Extract user_id from header, default to None if not present
+):
+    """Lists the groups for an agent"""
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+    print("in list agents with manager_type", manager_type)
+    return server.agent_manager.list_groups(agent_id=agent_id, manager_type=manager_type, actor=actor)

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

@@ -0,0 +1,132 @@
+from typing import List, Optional
+
+from fastapi import APIRouter, Body, Depends, Header
+from fastapi.exceptions import HTTPException
+from starlette.requests import Request
+
+from letta.agents.letta_agent_batch import LettaAgentBatch
+from letta.log import get_logger
+from letta.orm.errors import NoResultFound
+from letta.schemas.job import BatchJob, JobStatus, JobType
+from letta.schemas.letta_request import CreateBatch
+from letta.server.rest_api.utils import get_letta_server
+from letta.server.server import SyncServer
+
+router = APIRouter(prefix="/messages", tags=["messages"])
+
+logger = get_logger(__name__)
+
+
+# Batch APIs
+
+
+@router.post(
+    "/batches",
+    response_model=BatchJob,
+    operation_id="create_messages_batch",
+)
+async def create_messages_batch(
+    request: Request,
+    payload: CreateBatch = Body(..., description="Messages and config for all agents"),
+    server: SyncServer = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Submit a batch of agent messages for asynchronous processing.
+    Creates a job that will fan out messages to all listed agents and process them in parallel.
+    """
+    # Reject requests greater than 256Mbs
+    max_bytes = 256 * 1024 * 1024
+    content_length = request.headers.get("content-length")
+    if content_length:
+        length = int(content_length)
+        if length > max_bytes:
+            raise HTTPException(status_code=413, detail=f"Request too large ({length} bytes). Max is {max_bytes} bytes.")
+
+    try:
+        actor = server.user_manager.get_user_or_default(user_id=actor_id)
+
+        # Create a new job
+        batch_job = BatchJob(
+            user_id=actor.id,
+            status=JobStatus.created,
+            metadata={
+                "job_type": "batch_messages",
+            },
+            callback_url=str(payload.callback_url),
+        )
+
+        # create the batch runner
+        batch_runner = LettaAgentBatch(
+            message_manager=server.message_manager,
+            agent_manager=server.agent_manager,
+            block_manager=server.block_manager,
+            passage_manager=server.passage_manager,
+            batch_manager=server.batch_manager,
+            sandbox_config_manager=server.sandbox_config_manager,
+            job_manager=server.job_manager,
+            actor=actor,
+        )
+        llm_batch_job = await batch_runner.step_until_request(batch_requests=payload.requests, letta_batch_job_id=batch_job.id)
+
+        # TODO: update run metadata
+        batch_job = server.job_manager.create_job(pydantic_job=batch_job, actor=actor)
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        raise
+    return batch_job
+
+
+@router.get("/batches/{batch_id}", response_model=BatchJob, operation_id="retrieve_batch_run")
+async def retrieve_batch_run(
+    batch_id: str,
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+    server: "SyncServer" = Depends(get_letta_server),
+):
+    """
+    Get the status of a batch run.
+    """
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+
+    try:
+        job = server.job_manager.get_job_by_id(job_id=batch_id, actor=actor)
+        return BatchJob.from_job(job)
+    except NoResultFound:
+        raise HTTPException(status_code=404, detail="Batch not found")
+
+
+@router.get("/batches", response_model=List[BatchJob], operation_id="list_batch_runs")
+async def list_batch_runs(
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+    server: "SyncServer" = Depends(get_letta_server),
+):
+    """
+    List all batch runs.
+    """
+    # TODO: filter
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+
+    jobs = server.job_manager.list_jobs(actor=actor, statuses=[JobStatus.created, JobStatus.running], job_type=JobType.BATCH)
+    return [BatchJob.from_job(job) for job in jobs]
+
+
+@router.patch("/batches/{batch_id}/cancel", operation_id="cancel_batch_run")
+async def cancel_batch_run(
+    batch_id: str,
+    server: "SyncServer" = Depends(get_letta_server),
+    actor_id: Optional[str] = Header(None, alias="user_id"),
+):
+    """
+    Cancel a batch run.
+    """
+    actor = server.user_manager.get_user_or_default(user_id=actor_id)
+
+    try:
+        job = server.job_manager.get_job_by_id(job_id=batch_id, actor=actor)
+        job.status = JobStatus.cancelled
+        server.job_manager.update_job_by_id(job_id=job, job=job)
+        # TODO: actually cancel it
+    except NoResultFound:
+        raise HTTPException(status_code=404, detail="Run not found")