letta-nightly 0.6.27.dev20250220104103__py3-none-any.whl → 0.6.28.dev20250220163833__py3-none-any.whl

letta/services/agent_manager.py

@@ -27,10 +27,14 @@ from letta.schemas.message import MessageCreate
 from letta.schemas.passage import Passage as PydanticPassage
 from letta.schemas.source import Source as PydanticSource
 from letta.schemas.tool import Tool as PydanticTool
+from letta.schemas.tool_rule import ContinueToolRule as PydanticContinueToolRule
+from letta.schemas.tool_rule import TerminalToolRule as PydanticTerminalToolRule
 from letta.schemas.tool_rule import ToolRule as PydanticToolRule
 from letta.schemas.user import User as PydanticUser
+from letta.serialize_schemas import SerializedAgentSchema
 from letta.services.block_manager import BlockManager
 from letta.services.helpers.agent_manager_helper import (
+    _process_identity,
     _process_relationship,
     _process_tags,
     check_supports_structured_output,
@@ -39,6 +43,7 @@ from letta.services.helpers.agent_manager_helper import (
     initialize_message_sequence,
     package_initial_message_sequence,
 )
+from letta.services.identity_manager import IdentityManager
 from letta.services.message_manager import MessageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
@@ -53,13 +58,14 @@ class AgentManager:
     """Manager class to handle business logic related to Agents."""
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
         self.block_manager = BlockManager()
         self.tool_manager = ToolManager()
         self.source_manager = SourceManager()
         self.message_manager = MessageManager()
+        self.identity_manager = IdentityManager()
 
     # ======================================================================================================================
     # Basic CRUD operations
@@ -75,10 +81,6 @@ class AgentManager:
         if not agent_create.llm_config or not agent_create.embedding_config:
             raise ValueError("llm_config and embedding_config are required")
 
-        # Check tool rules are valid
-        if agent_create.tool_rules:
-            check_supports_structured_output(model=agent_create.llm_config.model, tool_rules=agent_create.tool_rules)
-
         # 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
         if agent_create.memory_blocks:
@@ -98,6 +100,25 @@ class AgentManager:
         # Remove duplicates
         tool_names = list(set(tool_names))
 
+        # add default tool rules
+        if agent_create.include_base_tool_rules:
+            if not agent_create.tool_rules:
+                tool_rules = []
+            else:
+                tool_rules = agent_create.tool_rules
+
+            # apply default tool rules
+            for tool_name in tool_names:
+                if tool_name == "send_message" or tool_name == "send_message_to_agent_async":
+                    tool_rules.append(PydanticTerminalToolRule(tool_name=tool_name))
+                elif tool_name in BASE_TOOLS:
+                    tool_rules.append(PydanticContinueToolRule(tool_name=tool_name))
+        else:
+            tool_rules = agent_create.tool_rules
+        # Check tool rules are valid
+        if agent_create.tool_rules:
+            check_supports_structured_output(model=agent_create.llm_config.model, tool_rules=agent_create.tool_rules)
+
         tool_ids = agent_create.tool_ids or []
         for tool_name in tool_names:
             tool = self.tool_manager.get_tool_by_name(tool_name=tool_name, actor=actor)
@@ -119,11 +140,12 @@ class AgentManager:
             tags=agent_create.tags or [],
             description=agent_create.description,
             metadata=agent_create.metadata,
-            tool_rules=agent_create.tool_rules,
+            tool_rules=tool_rules,
             actor=actor,
             project_id=agent_create.project_id,
             template_id=agent_create.template_id,
             base_template_id=agent_create.base_template_id,
+            identifier_key=agent_create.identifier_key,
             message_buffer_autoclear=agent_create.message_buffer_autoclear,
         )
 
@@ -187,6 +209,7 @@ class AgentManager:
         project_id: Optional[str] = None,
         template_id: Optional[str] = None,
         base_template_id: Optional[str] = None,
+        identifier_key: Optional[str] = None,
         message_buffer_autoclear: bool = False,
     ) -> PydanticAgentState:
         """Create a new agent."""
@@ -214,6 +237,10 @@ class AgentManager:
             _process_relationship(session, new_agent, "sources", SourceModel, source_ids, replace=True)
             _process_relationship(session, new_agent, "core_memory", BlockModel, block_ids, replace=True)
             _process_tags(new_agent, tags, replace=True)
+            if identifier_key is not None:
+                identity = self.identity_manager.get_identity_from_identifier_key(identifier_key)
+                _process_identity(new_agent, identifier_key, identity)
+
             new_agent.create(session, actor=actor)
 
             # Convert to PydanticAgentState and return
@@ -286,6 +313,9 @@ class AgentManager:
                 _process_relationship(session, agent, "core_memory", BlockModel, agent_update.block_ids, replace=True)
             if agent_update.tags is not None:
                 _process_tags(agent, agent_update.tags, replace=True)
+            if agent_update.identifier_key is not None:
+                identity = self.identity_manager.get_identity_from_identifier_key(agent_update.identifier_key)
+                _process_identity(agent, agent_update.identifier_key, identity)
 
             # Commit and refresh the agent
             agent.update(session, actor=actor)
@@ -355,6 +385,24 @@ class AgentManager:
             agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
             agent.hard_delete(session)
 
+    @enforce_types
+    def serialize(self, agent_id: str, actor: PydanticUser) -> dict:
+        with self.session_maker() as session:
+            # Retrieve the agent
+            agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+            schema = SerializedAgentSchema(session=session)
+            return schema.dump(agent)
+
+    @enforce_types
+    def deserialize(self, serialized_agent: dict, actor: PydanticUser) -> PydanticAgentState:
+        # TODO: Use actor to override fields
+        with self.session_maker() as session:
+            schema = SerializedAgentSchema(session=session)
+            agent = schema.load(serialized_agent, session=session)
+            agent.organization_id = actor.organization_id
+            agent = agent.create(session, actor=actor)
+            return agent.to_pydantic()
+
     # ======================================================================================================================
     # Per Agent Environment Variable Management
     # ======================================================================================================================

letta/services/block_manager.py

@@ -16,7 +16,7 @@ class BlockManager:
 
     def __init__(self):
         # Fetching the db_context similarly as in ToolManager
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/helpers/agent_manager_helper.py

@@ -11,6 +11,7 @@ from letta.orm.errors import NoResultFound
 from letta.prompts import gpt_system
 from letta.schemas.agent import AgentState, AgentType
 from letta.schemas.enums import MessageRole
+from letta.schemas.identity import Identity
 from letta.schemas.memory import Memory
 from letta.schemas.message import Message, MessageCreate, TextContent
 from letta.schemas.tool_rule import ToolRule
@@ -84,6 +85,20 @@ def _process_tags(agent: AgentModel, tags: List[str], replace=True):
         agent.tags.extend([tag for tag in new_tags if tag.tag not in existing_tags])
 
 
+def _process_identity(agent: AgentModel, identifier_key: str, identity: Identity):
+    """
+    Handles identity for an agent.
+
+    Args:
+        agent: The AgentModel instance.
+        identifier_key: The identifier key of the identity to set or update.
+        identity: The Identity object to set or update.
+    """
+    agent.identifier_key = identifier_key
+    agent.identity = identity
+    agent.identity_id = identity.id
+
+
 def derive_system_message(agent_type: AgentType, system: Optional[str] = None):
     if system is None:
         # TODO: don't hardcode

letta/services/identity_manager.py

@@ -0,0 +1,140 @@
+from typing import List, Optional
+
+from fastapi import HTTPException
+from sqlalchemy.orm import Session
+
+from letta.orm.agent import Agent as AgentModel
+from letta.orm.identity import Identity as IdentityModel
+from letta.schemas.identity import Identity as PydanticIdentity
+from letta.schemas.identity import IdentityCreate, IdentityType, IdentityUpdate
+from letta.schemas.user import User as PydanticUser
+from letta.utils import enforce_types
+
+
+class IdentityManager:
+
+    def __init__(self):
+        from letta.server.db import db_context
+
+        self.session_maker = db_context
+
+    @enforce_types
+    def list_identities(
+        self,
+        name: Optional[str] = None,
+        project_id: Optional[str] = None,
+        identity_type: Optional[IdentityType] = None,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
+        limit: Optional[int] = 50,
+        actor: PydanticUser = None,
+    ) -> list[PydanticIdentity]:
+        with self.session_maker() as session:
+            filters = {"organization_id": actor.organization_id}
+            if project_id:
+                filters["project_id"] = project_id
+            if identity_type:
+                filters["identity_type"] = identity_type
+            identities = IdentityModel.list(
+                db_session=session,
+                query_text=name,
+                before=before,
+                after=after,
+                limit=limit,
+                **filters,
+            )
+            return [identity.to_pydantic() for identity in identities]
+
+    @enforce_types
+    def get_identity_from_identifier_key(self, identifier_key: str) -> PydanticIdentity:
+        with self.session_maker() as session:
+            identity = IdentityModel.read(db_session=session, identifier_key=identifier_key)
+            return identity.to_pydantic()
+
+    @enforce_types
+    def create_identity(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
+        with self.session_maker() as session:
+            new_identity = IdentityModel(**identity.model_dump(exclude={"agent_ids"}, exclude_unset=True))
+            new_identity.organization_id = actor.organization_id
+            self._process_agent_relationship(session=session, identity=new_identity, agent_ids=identity.agent_ids, allow_partial=False)
+            new_identity.create(session, actor=actor)
+            return new_identity.to_pydantic()
+
+    @enforce_types
+    def upsert_identity(self, identity: IdentityCreate, actor: PydanticUser) -> PydanticIdentity:
+        with self.session_maker() as session:
+            existing_identity = IdentityModel.read(
+                db_session=session,
+                identifier_key=identity.identifier_key,
+                project_id=identity.project_id,
+                organization_id=actor.organization_id,
+            )
+
+        if existing_identity is None:
+            return self.create_identity(identity=identity, actor=actor)
+        else:
+            if existing_identity.identifier_key != identity.identifier_key:
+                raise HTTPException(status_code=400, detail="Identifier key is an immutable field")
+            if existing_identity.project_id != identity.project_id:
+                raise HTTPException(status_code=400, detail="Project id is an immutable field")
+            identity_update = IdentityUpdate(name=identity.name, identity_type=identity.identity_type, agent_ids=identity.agent_ids)
+            return self.update_identity_by_key(identity.identifier_key, identity_update, actor, replace=True)
+
+    @enforce_types
+    def update_identity_by_key(
+        self, identifier_key: str, identity: IdentityUpdate, actor: PydanticUser, replace: bool = False
+    ) -> PydanticIdentity:
+        with self.session_maker() as session:
+            try:
+                existing_identity = IdentityModel.read(db_session=session, identifier_key=identifier_key)
+            except NoResultFound:
+                raise HTTPException(status_code=404, detail="Identity not found")
+            if existing_identity.organization_id != actor.organization_id:
+                raise HTTPException(status_code=403, detail="Forbidden")
+
+            existing_identity.name = identity.name if identity.name is not None else existing_identity.name
+            existing_identity.identity_type = (
+                identity.identity_type if identity.identity_type is not None else existing_identity.identity_type
+            )
+            self._process_agent_relationship(
+                session=session, identity=existing_identity, agent_ids=identity.agent_ids, allow_partial=False, replace=replace
+            )
+            existing_identity.update(session, actor=actor)
+            return existing_identity.to_pydantic()
+
+    @enforce_types
+    def delete_identity_by_key(self, identifier_key: str, actor: PydanticUser) -> None:
+        with self.session_maker() as session:
+            identity = IdentityModel.read(db_session=session, identifier_key=identifier_key)
+            if identity is None:
+                raise HTTPException(status_code=404, detail="Identity not found")
+            if identity.organization_id != actor.organization_id:
+                raise HTTPException(status_code=403, detail="Forbidden")
+            session.delete(identity)
+            session.commit()
+
+    def _process_agent_relationship(
+        self, session: Session, identity: IdentityModel, agent_ids: List[str], allow_partial=False, replace=True
+    ):
+        current_relationship = getattr(identity, "agents", [])
+        if not agent_ids:
+            if replace:
+                setattr(identity, "agents", [])
+            return
+
+        # Retrieve models for the provided IDs
+        found_items = session.query(AgentModel).filter(AgentModel.id.in_(agent_ids)).all()
+
+        # Validate all items are found if allow_partial is False
+        if not allow_partial and len(found_items) != len(agent_ids):
+            missing = set(agent_ids) - {item.id for item in found_items}
+            raise NoResultFound(f"Items not found in agents: {missing}")
+
+        if replace:
+            # Replace the relationship
+            setattr(identity, "agents", found_items)
+        else:
+            # Extend the relationship (only add new items)
+            current_ids = {item.id for item in current_relationship}
+            new_items = [item for item in found_items if item.id not in current_ids]
+            current_relationship.extend(new_items)

letta/services/job_manager.py

@@ -29,7 +29,7 @@ class JobManager:
 
     def __init__(self):
         # Fetching the db_context similarly as in OrganizationManager
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/message_manager.py

@@ -16,7 +16,7 @@ class MessageManager:
     """Manager class to handle business logic related to Messages."""
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/organization_manager.py

@@ -16,7 +16,7 @@ class OrganizationManager:
         # TODO: Please refactor this out
         # I am currently working on a ORM refactor and would like to make a more minimal set of changes
         # - Matt
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/passage_manager.py

@@ -16,7 +16,7 @@ class PassageManager:
     """Manager class to handle business logic related to Passages."""
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/provider_manager.py

@@ -10,7 +10,7 @@ from letta.utils import enforce_types
 class ProviderManager:
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/sandbox_config_manager.py

@@ -20,7 +20,7 @@ class SandboxConfigManager:
     """Manager class to handle business logic related to SandboxConfig and SandboxEnvironmentVariable."""
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/source_manager.py

@@ -15,7 +15,7 @@ class SourceManager:
     """Manager class to handle business logic related to Sources."""
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/step_manager.py

@@ -17,7 +17,7 @@ from letta.utils import enforce_types
 class StepManager:
 
     def __init__(self):
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/tool_manager.py

@@ -31,7 +31,7 @@ class ToolManager:
 
     def __init__(self):
         # Fetching the db_context similarly as in OrganizationManager
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/services/user_manager.py

@@ -17,7 +17,7 @@ class UserManager:
 
     def __init__(self):
         # Fetching the db_context similarly as in OrganizationManager
-        from letta.server.server import db_context
+        from letta.server.db import db_context
 
         self.session_maker = db_context
 

letta/settings.py

@@ -60,6 +60,9 @@ class ModelSettings(BaseSettings):
     openai_api_key: Optional[str] = None
     openai_api_base: str = "https://api.openai.com/v1"
 
+    # deepseek
+    deepseek_api_key: Optional[str] = None
+
     # groq
     groq_api_key: Optional[str] = None
 

letta/tracing.py

@@ -0,0 +1,205 @@
+import asyncio
+import inspect
+import sys
+import time
+from functools import wraps
+from typing import Any, Dict, Optional
+
+from fastapi import Request
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.requests import RequestsInstrumentor
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import Span, Status, StatusCode
+
+# Get a tracer instance - will be no-op until setup_tracing is called
+tracer = trace.get_tracer(__name__)
+
+# Track if tracing has been initialized
+_is_tracing_initialized = False
+
+
+def is_pytest_environment():
+    """Check if we're running in pytest"""
+    return "pytest" in sys.modules
+
+
+def trace_method(name=None):
+    """Decorator to add tracing to a method"""
+
+    def decorator(func):
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Skip tracing if not initialized
+            if not _is_tracing_initialized:
+                return await func(*args, **kwargs)
+
+            span_name = name or func.__name__
+            with tracer.start_as_current_span(span_name) as span:
+                span.set_attribute("code.namespace", inspect.getmodule(func).__name__)
+                span.set_attribute("code.function", func.__name__)
+
+                if len(args) > 0 and hasattr(args[0], "__class__"):
+                    span.set_attribute("code.class", args[0].__class__.__name__)
+
+                request = _extract_request_info(args, span)
+                if request and len(request) > 0:
+                    span.set_attribute("agent.id", kwargs.get("agent_id"))
+                    span.set_attribute("actor.id", request.get("http.user_id"))
+
+                try:
+                    result = await func(*args, **kwargs)
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR))
+                    span.record_exception(e)
+                    raise
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            # Skip tracing if not initialized
+            if not _is_tracing_initialized:
+                return func(*args, **kwargs)
+
+            span_name = name or func.__name__
+            with tracer.start_as_current_span(span_name) as span:
+                span.set_attribute("code.namespace", inspect.getmodule(func).__name__)
+                span.set_attribute("code.function", func.__name__)
+
+                if len(args) > 0 and hasattr(args[0], "__class__"):
+                    span.set_attribute("code.class", args[0].__class__.__name__)
+
+                request = _extract_request_info(args, span)
+                if request and len(request) > 0:
+                    span.set_attribute("agent.id", kwargs.get("agent_id"))
+                    span.set_attribute("actor.id", request.get("http.user_id"))
+
+                try:
+                    result = func(*args, **kwargs)
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+                except Exception as e:
+                    span.set_status(Status(StatusCode.ERROR))
+                    span.record_exception(e)
+                    raise
+
+        return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
+
+    return decorator
+
+
+def log_attributes(attributes: Dict[str, Any]) -> None:
+    """
+    Log multiple attributes to the current active span.
+
+    Args:
+        attributes: Dictionary of attribute key-value pairs
+    """
+    current_span = trace.get_current_span()
+    if current_span:
+        current_span.set_attributes(attributes)
+
+
+def log_event(name: str, attributes: Optional[Dict[str, Any]] = None, timestamp: Optional[int] = None) -> None:
+    """
+    Log an event to the current active span.
+
+    Args:
+        name: Name of the event
+        attributes: Optional dictionary of event attributes
+        timestamp: Optional timestamp in nanoseconds
+    """
+    current_span = trace.get_current_span()
+    if current_span:
+        if timestamp is None:
+            timestamp = int(time.perf_counter_ns())
+
+        current_span.add_event(name=name, attributes=attributes, timestamp=timestamp)
+
+
+def get_trace_id() -> str:
+    current_span = trace.get_current_span()
+    if current_span:
+        return format(current_span.get_span_context().trace_id, "032x")
+    else:
+        return ""
+
+
+def request_hook(span: Span, _request_context: Optional[Dict] = None, response: Optional[Any] = None):
+    """Hook to update span based on response status code"""
+    if response is not None:
+        if hasattr(response, "status_code"):
+            span.set_attribute("http.status_code", response.status_code)
+            if response.status_code >= 400:
+                span.set_status(Status(StatusCode.ERROR))
+            elif 200 <= response.status_code < 300:
+                span.set_status(Status(StatusCode.OK))
+
+
+def setup_tracing(endpoint: str, service_name: str = "memgpt-server") -> None:
+    """
+    Sets up OpenTelemetry tracing with OTLP exporter for specific endpoints
+
+    Args:
+        endpoint: OTLP endpoint URL
+        service_name: Name of the service for tracing
+    """
+    global _is_tracing_initialized
+
+    # Skip tracing in pytest environment
+    if is_pytest_environment():
+        print("ℹ️ Skipping tracing setup in pytest environment")
+        return
+
+    # Create a Resource to identify our service
+    resource = Resource.create({"service.name": service_name, "service.namespace": "default", "deployment.environment": "production"})
+
+    # Initialize the TracerProvider with the resource
+    provider = TracerProvider(resource=resource)
+
+    # Only set up OTLP export if endpoint is provided
+    if endpoint:
+        otlp_exporter = OTLPSpanExporter(endpoint=endpoint)
+        processor = BatchSpanProcessor(otlp_exporter)
+        provider.add_span_processor(processor)
+        _is_tracing_initialized = True
+    else:
+        print("⚠️ Warning: Tracing endpoint not provided, tracing will be disabled")
+
+    # Set the global TracerProvider
+    trace.set_tracer_provider(provider)
+
+    # Initialize automatic instrumentation for the requests library with response hook
+    if _is_tracing_initialized:
+        RequestsInstrumentor().instrument(response_hook=request_hook)
+
+
+def _extract_request_info(args: tuple, span: Span) -> Dict[str, Any]:
+    """
+    Safely extracts request information from function arguments.
+    Works with both FastAPI route handlers and inner functions.
+    """
+    attributes = {}
+
+    # Look for FastAPI Request object in args
+    request = next((arg for arg in args if isinstance(arg, Request)), None)
+
+    if request:
+        attributes.update(
+            {
+                "http.route": request.url.path,
+                "http.method": request.method,
+                "http.scheme": request.url.scheme,
+                "http.target": str(request.url.path),
+                "http.url": str(request.url),
+                "http.flavor": request.scope.get("http_version", ""),
+                "http.client_ip": request.client.host if request.client else None,
+                "http.user_id": request.headers.get("user_id"),
+            }
+        )
+
+    span.set_attributes(attributes)
+    return attributes

letta/utils.py

@@ -824,12 +824,16 @@ def parse_json(string) -> dict:
     result = None
     try:
         result = json_loads(string)
+        if not isinstance(result, dict):
+            raise ValueError(f"JSON from string input ({string}) is not a dictionary (type {type(result)}): {result}")
         return result
     except Exception as e:
         print(f"Error parsing json with json package: {e}")
 
     try:
         result = demjson.decode(string)
+        if not isinstance(result, dict):
+            raise ValueError(f"JSON from string input ({string}) is not a dictionary (type {type(result)}): {result}")
         return result
     except demjson.JSONDecodeError as e:
         print(f"Error parsing json with demjson package: {e}")