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

letta/services/agent_manager.py

@@ -1,16 +1,26 @@
 from typing import Dict, List, Optional
+from datetime import datetime
+import numpy as np
 
-from letta.constants import BASE_MEMORY_TOOLS, BASE_TOOLS
+from sqlalchemy import select, union_all, literal, func, Select
+
+from letta.constants import BASE_MEMORY_TOOLS, BASE_TOOLS, MAX_EMBEDDING_DIM
+from letta.embeddings import embedding_model
+from letta.log import get_logger
 from letta.orm import Agent as AgentModel
 from letta.orm import Block as BlockModel
 from letta.orm import Source as SourceModel
 from letta.orm import Tool as ToolModel
+from letta.orm import AgentPassage, SourcePassage
+from letta.orm import SourcesAgents
 from letta.orm.errors import NoResultFound
+from letta.orm.sqlite_functions import adapt_array
 from letta.schemas.agent import AgentState as PydanticAgentState
 from letta.schemas.agent import AgentType, CreateAgent, UpdateAgent
 from letta.schemas.block import Block as PydanticBlock
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.llm_config import LLMConfig
+from letta.schemas.passage import Passage as PydanticPassage
 from letta.schemas.source import Source as PydanticSource
 from letta.schemas.tool_rule import ToolRule as PydanticToolRule
 from letta.schemas.user import User as PydanticUser
@@ -20,11 +30,13 @@ from letta.services.helpers.agent_manager_helper import (
     _process_tags,
     derive_system_message,
 )
-from letta.services.passage_manager import PassageManager
 from letta.services.source_manager import SourceManager
 from letta.services.tool_manager import ToolManager
+from letta.settings import settings
 from letta.utils import enforce_types
 
+logger = get_logger(__name__)
+
 
 # Agent Manager Class
 class AgentManager:
@@ -49,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:
@@ -226,13 +241,6 @@ class AgentManager:
         with self.session_maker() as session:
             # Retrieve the agent
             agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
-
-            # TODO: @mindy delete this piece when we have a proper passages/sources implementation
-            # TODO: This is done very hacky on purpose
-            # TODO: 1000 limit is also wack
-            passage_manager = PassageManager()
-            passage_manager.delete_passages(actor=actor, agent_id=agent_id, limit=1000)
-
             agent_state = agent.to_pydantic()
             agent.hard_delete(session)
             return agent_state
@@ -403,3 +411,330 @@ class AgentManager:
 
             agent.update(session, actor=actor)
             return agent.to_pydantic()
+
+    # ======================================================================================================================
+    # Passage Management
+    # ======================================================================================================================
+    def _build_passage_query(
+        self,
+        actor: PydanticUser,
+        agent_id: Optional[str] = None,
+        file_id: Optional[str] = None,
+        query_text: Optional[str] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        cursor: Optional[str] = None,
+        source_id: Optional[str] = None,
+        embed_query: bool = False,
+        ascending: bool = True,
+        embedding_config: Optional[EmbeddingConfig] = None,
+        agent_only: bool = False,
+    ) -> Select:
+        """Helper function to build the base passage query with all filters applied.
+        
+        Returns the query before any limit or count operations are applied.
+        """
+        embedded_text = None
+        if embed_query:
+            assert embedding_config is not None, "embedding_config must be specified for vector search"
+            assert query_text is not None, "query_text must be specified for vector search"
+            embedded_text = embedding_model(embedding_config).get_text_embedding(query_text)
+            embedded_text = np.array(embedded_text)
+            embedded_text = np.pad(embedded_text, (0, MAX_EMBEDDING_DIM - embedded_text.shape[0]), mode="constant").tolist()
+
+        with self.session_maker() as session:
+            # Start with base query for source passages
+            source_passages = None
+            if not agent_only:  # Include source passages
+                if agent_id is not None:
+                    source_passages = (
+                        select(
+                            SourcePassage,
+                            literal(None).label('agent_id')
+                        )
+                        .join(SourcesAgents, SourcesAgents.source_id == SourcePassage.source_id)
+                        .where(SourcesAgents.agent_id == agent_id)
+                        .where(SourcePassage.organization_id == actor.organization_id)
+                    )
+                else:
+                    source_passages = (
+                        select(
+                            SourcePassage,
+                            literal(None).label('agent_id')
+                        )
+                        .where(SourcePassage.organization_id == actor.organization_id)
+                    )
+
+                if source_id:
+                    source_passages = source_passages.where(SourcePassage.source_id == source_id)
+                if file_id:
+                    source_passages = source_passages.where(SourcePassage.file_id == file_id)
+
+            # Add agent passages query
+            agent_passages = None
+            if agent_id is not None:
+                agent_passages = (
+                    select(
+                        AgentPassage.id,
+                        AgentPassage.text,
+                        AgentPassage.embedding_config,
+                        AgentPassage.metadata_,
+                        AgentPassage.embedding,
+                        AgentPassage.created_at,
+                        AgentPassage.updated_at,
+                        AgentPassage.is_deleted,
+                        AgentPassage._created_by_id,
+                        AgentPassage._last_updated_by_id,
+                        AgentPassage.organization_id,
+                        literal(None).label('file_id'),
+                        literal(None).label('source_id'),
+                        AgentPassage.agent_id
+                    )
+                    .where(AgentPassage.agent_id == agent_id)
+                    .where(AgentPassage.organization_id == actor.organization_id)
+                )
+
+            # Combine queries
+            if source_passages is not None and agent_passages is not None:
+                combined_query = union_all(source_passages, agent_passages).cte('combined_passages')
+            elif agent_passages is not None:
+                combined_query = agent_passages.cte('combined_passages')
+            elif source_passages is not None:
+                combined_query = source_passages.cte('combined_passages')
+            else:
+                raise ValueError("No passages found")
+
+            # Build main query from combined CTE
+            main_query = select(combined_query)
+
+            # Apply filters
+            if start_date:
+                main_query = main_query.where(combined_query.c.created_at >= start_date)
+            if end_date:
+                main_query = main_query.where(combined_query.c.created_at <= end_date)
+            if source_id:
+                main_query = main_query.where(combined_query.c.source_id == source_id)
+            if file_id:
+                main_query = main_query.where(combined_query.c.file_id == file_id)
+
+            # Vector search
+            if embedded_text:
+                if settings.letta_pg_uri_no_default:
+                    # PostgreSQL with pgvector
+                    main_query = main_query.order_by(
+                        combined_query.c.embedding.cosine_distance(embedded_text).asc()
+                    )
+                else:
+                    # SQLite with custom vector type
+                    query_embedding_binary = adapt_array(embedded_text)
+                    if ascending:
+                        main_query = main_query.order_by(
+                            func.cosine_distance(combined_query.c.embedding, query_embedding_binary).asc(),
+                            combined_query.c.created_at.asc(),
+                            combined_query.c.id.asc()
+                        )
+                    else:
+                        main_query = main_query.order_by(
+                            func.cosine_distance(combined_query.c.embedding, query_embedding_binary).asc(),
+                            combined_query.c.created_at.desc(),
+                            combined_query.c.id.asc()
+                        )
+            else:
+                if query_text:
+                    main_query = main_query.where(func.lower(combined_query.c.text).contains(func.lower(query_text)))
+
+            # Handle cursor-based pagination
+            if cursor:
+                cursor_query = select(combined_query.c.created_at).where(
+                    combined_query.c.id == cursor
+                ).scalar_subquery()
+                
+                if ascending:
+                    main_query = main_query.where(
+                        combined_query.c.created_at > cursor_query
+                    )
+                else:
+                    main_query = main_query.where(
+                        combined_query.c.created_at < cursor_query
+                    )
+
+            # Add ordering if not already ordered by similarity
+            if not embed_query:
+                if ascending:
+                    main_query = main_query.order_by(
+                        combined_query.c.created_at.asc(),
+                        combined_query.c.id.asc(),
+                    )
+                else:
+                    main_query = main_query.order_by(
+                        combined_query.c.created_at.desc(),
+                        combined_query.c.id.asc(),
+                    )
+
+            return main_query
+
+    @enforce_types
+    def list_passages(
+        self,
+        actor: PydanticUser,
+        agent_id: Optional[str] = None,
+        file_id: Optional[str] = None,
+        limit: Optional[int] = 50,
+        query_text: Optional[str] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        cursor: Optional[str] = None,
+        source_id: Optional[str] = None,
+        embed_query: bool = False,
+        ascending: bool = True,
+        embedding_config: Optional[EmbeddingConfig] = None,
+        agent_only: bool = False
+    ) -> List[PydanticPassage]:
+        """Lists all passages attached to an agent."""
+        with self.session_maker() as session:
+            main_query = self._build_passage_query(
+                actor=actor,
+                agent_id=agent_id,
+                file_id=file_id,
+                query_text=query_text,
+                start_date=start_date,
+                end_date=end_date,
+                cursor=cursor,
+                source_id=source_id,
+                embed_query=embed_query,
+                ascending=ascending,
+                embedding_config=embedding_config,
+                agent_only=agent_only,
+            )
+
+            # Add limit
+            if limit:
+                main_query = main_query.limit(limit)
+
+            # Execute query
+            results = list(session.execute(main_query))
+
+            passages = []
+            for row in results:
+                data = dict(row._mapping)
+                if data['agent_id'] is not None:
+                    # This is an AgentPassage - remove source fields
+                    data.pop('source_id', None)
+                    data.pop('file_id', None)
+                    passage = AgentPassage(**data)
+                else:
+                    # This is a SourcePassage - remove agent field
+                    data.pop('agent_id', None)
+                    passage = SourcePassage(**data)
+                passages.append(passage)
+            
+            return [p.to_pydantic() for p in passages]
+
+
+    @enforce_types
+    def passage_size(
+        self,
+        actor: PydanticUser,
+        agent_id: Optional[str] = None,
+        file_id: Optional[str] = None,
+        query_text: Optional[str] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
+        cursor: Optional[str] = None,
+        source_id: Optional[str] = None,
+        embed_query: bool = False,
+        ascending: bool = True,
+        embedding_config: Optional[EmbeddingConfig] = None,
+        agent_only: bool = False
+    ) -> int:
+        """Returns the count of passages matching the given criteria."""
+        with self.session_maker() as session:
+            main_query = self._build_passage_query(
+                actor=actor,
+                agent_id=agent_id,
+                file_id=file_id,
+                query_text=query_text,
+                start_date=start_date,
+                end_date=end_date,
+                cursor=cursor,
+                source_id=source_id,
+                embed_query=embed_query,
+                ascending=ascending,
+                embedding_config=embedding_config,
+                agent_only=agent_only,
+            )
+            
+            # Convert to count query
+            count_query = select(func.count()).select_from(main_query.subquery())
+            return session.scalar(count_query) or 0
+
+    # ======================================================================================================================
+    # Tool Management
+    # ======================================================================================================================
+    @enforce_types
+    def attach_tool(self, agent_id: str, tool_id: str, actor: PydanticUser) -> PydanticAgentState:
+        """
+        Attaches a tool to an agent.
+
+        Args:
+            agent_id: ID of the agent to attach the tool to.
+            tool_id: ID of the tool to attach.
+            actor: User performing the action.
+
+        Raises:
+            NoResultFound: If the agent or tool is not found.
+
+        Returns:
+            PydanticAgentState: The updated agent state.
+        """
+        with self.session_maker() as session:
+            # Verify the agent exists and user has permission to access it
+            agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+
+            # Use the _process_relationship helper to attach the tool
+            _process_relationship(
+                session=session,
+                agent=agent,
+                relationship_name="tools",
+                model_class=ToolModel,
+                item_ids=[tool_id],
+                allow_partial=False,  # Ensure the tool exists
+                replace=False,  # Extend the existing tools
+            )
+
+            # Commit and refresh the agent
+            agent.update(session, actor=actor)
+            return agent.to_pydantic()
+
+    @enforce_types
+    def detach_tool(self, agent_id: str, tool_id: str, actor: PydanticUser) -> PydanticAgentState:
+        """
+        Detaches a tool from an agent.
+
+        Args:
+            agent_id: ID of the agent to detach the tool from.
+            tool_id: ID of the tool to detach.
+            actor: User performing the action.
+
+        Raises:
+            NoResultFound: If the agent or tool is not found.
+
+        Returns:
+            PydanticAgentState: The updated agent state.
+        """
+        with self.session_maker() as session:
+            # Verify the agent exists and user has permission to access it
+            agent = AgentModel.read(db_session=session, identifier=agent_id, actor=actor)
+
+            # Filter out the tool to be detached
+            remaining_tools = [tool for tool in agent.tools if tool.id != tool_id]
+
+            if len(remaining_tools) == len(agent.tools):  # Tool ID was not in the relationship
+                logger.warning(f"Attempted to remove unattached tool id={tool_id} from agent id={agent_id} by actor={actor}")
+
+            # Update the tools relationship
+            agent.tools = remaining_tools
+
+            # Commit and refresh the agent
+            agent.update(session, actor=actor)
+            return agent.to_pydantic()

letta/services/passage_manager.py

@@ -1,12 +1,13 @@
-from datetime import datetime
 from typing import List, Optional
-
+from datetime import datetime
 import numpy as np
 
+from sqlalchemy import select, union_all, literal
+
 from letta.constants import MAX_EMBEDDING_DIM
 from letta.embeddings import embedding_model, parse_and_chunk_text
 from letta.orm.errors import NoResultFound
-from letta.orm.passage import Passage as PassageModel
+from letta.orm.passage import AgentPassage, SourcePassage
 from letta.schemas.agent import AgentState
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.passage import Passage as PydanticPassage
@@ -14,6 +15,7 @@ from letta.schemas.user import User as PydanticUser
 from letta.utils import enforce_types
 
 
+
 class PassageManager:
     """Manager class to handle business logic related to Passages."""
 
@@ -26,14 +28,51 @@ class PassageManager:
     def get_passage_by_id(self, passage_id: str, actor: PydanticUser) -> Optional[PydanticPassage]:
         """Fetch a passage by ID."""
         with self.session_maker() as session:
-            passage = PassageModel.read(db_session=session, identifier=passage_id, actor=actor)
-            return passage.to_pydantic()
+            # Try source passages first
+            try:
+                passage = SourcePassage.read(db_session=session, identifier=passage_id, actor=actor)
+                return passage.to_pydantic()
+            except NoResultFound:
+                # Try archival passages
+                try:
+                    passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                    return passage.to_pydantic()
+                except NoResultFound:
+                    raise NoResultFound(f"Passage with id {passage_id} not found in database.")
 
     @enforce_types
     def create_passage(self, pydantic_passage: PydanticPassage, actor: PydanticUser) -> PydanticPassage:
-        """Create a new passage."""
+        """Create a new passage in the appropriate table based on whether it has agent_id or source_id."""
+        # Common fields for both passage types
+        data = pydantic_passage.model_dump()
+        common_fields = {
+            "id": data.get("id"),
+            "text": data["text"],
+            "embedding": data["embedding"],
+            "embedding_config": data["embedding_config"],
+            "organization_id": data["organization_id"],
+            "metadata_": data.get("metadata_", {}),
+            "is_deleted": data.get("is_deleted", False),
+            "created_at": data.get("created_at", datetime.utcnow()),
+        }
+
+        if "agent_id" in data and data["agent_id"]:
+            assert not data.get("source_id"), "Passage cannot have both agent_id and source_id"
+            agent_fields = {
+                "agent_id": data["agent_id"],
+            }
+            passage = AgentPassage(**common_fields, **agent_fields)
+        elif "source_id" in data and data["source_id"]:
+            assert not data.get("agent_id"), "Passage cannot have both agent_id and source_id"
+            source_fields = {
+                "source_id": data["source_id"],
+                "file_id": data.get("file_id"),
+            }
+            passage = SourcePassage(**common_fields, **source_fields)
+        else:
+            raise ValueError("Passage must have either agent_id or source_id")
+
         with self.session_maker() as session:
-            passage = PassageModel(**pydantic_passage.model_dump())
             passage.create(session, actor=actor)
             return passage.to_pydantic()
 
@@ -93,14 +132,23 @@ class PassageManager:
             raise ValueError("Passage ID must be provided.")
 
         with self.session_maker() as session:
-            # Fetch existing message from database
-            curr_passage = PassageModel.read(
-                db_session=session,
-                identifier=passage_id,
-                actor=actor,
-            )
-            if not curr_passage:
-                raise ValueError(f"Passage with id {passage_id} does not exist.")
+            # Try source passages first
+            try:
+                curr_passage = SourcePassage.read(
+                    db_session=session,
+                    identifier=passage_id,
+                    actor=actor,
+                )
+            except NoResultFound:
+                # Try agent passages
+                try:
+                    curr_passage = AgentPassage.read(
+                        db_session=session,
+                        identifier=passage_id,
+                        actor=actor,
+                    )
+                except NoResultFound:
+                    raise ValueError(f"Passage with id {passage_id} does not exist.")
 
             # Update the database record with values from the provided record
             update_data = passage.model_dump(exclude_unset=True, exclude_none=True)
@@ -113,104 +161,32 @@ class PassageManager:
 
     @enforce_types
     def delete_passage_by_id(self, passage_id: str, actor: PydanticUser) -> bool:
-        """Delete a passage."""
+        """Delete a passage from either source or archival passages."""
         if not passage_id:
             raise ValueError("Passage ID must be provided.")
 
         with self.session_maker() as session:
+            # Try source passages first
             try:
-                passage = PassageModel.read(db_session=session, identifier=passage_id, actor=actor)
+                passage = SourcePassage.read(db_session=session, identifier=passage_id, actor=actor)
                 passage.hard_delete(session, actor=actor)
+                return True
             except NoResultFound:
-                raise ValueError(f"Passage with id {passage_id} not found.")
-
-    @enforce_types
-    def list_passages(
-        self,
-        actor: PydanticUser,
-        agent_id: Optional[str] = None,
-        file_id: Optional[str] = None,
-        cursor: Optional[str] = None,
-        limit: Optional[int] = 50,
-        query_text: Optional[str] = None,
-        start_date: Optional[datetime] = None,
-        end_date: Optional[datetime] = None,
-        ascending: bool = True,
-        source_id: Optional[str] = None,
-        embed_query: bool = False,
-        embedding_config: Optional[EmbeddingConfig] = None,
-    ) -> List[PydanticPassage]:
-        """List passages with pagination."""
-        with self.session_maker() as session:
-            filters = {"organization_id": actor.organization_id}
-            if agent_id:
-                filters["agent_id"] = agent_id
-            if file_id:
-                filters["file_id"] = file_id
-            if source_id:
-                filters["source_id"] = source_id
-
-            embedded_text = None
-            if embed_query:
-                assert embedding_config is not None
-
-                # Embed the text
-                embedded_text = embedding_model(embedding_config).get_text_embedding(query_text)
-
-                # Pad the embedding with zeros
-                embedded_text = np.array(embedded_text)
-                embedded_text = np.pad(embedded_text, (0, MAX_EMBEDDING_DIM - embedded_text.shape[0]), mode="constant").tolist()
-
-            results = PassageModel.list(
-                db_session=session,
-                cursor=cursor,
-                start_date=start_date,
-                end_date=end_date,
-                limit=limit,
-                ascending=ascending,
-                query_text=query_text if not embedded_text else None,
-                query_embedding=embedded_text,
-                **filters,
-            )
-            return [p.to_pydantic() for p in results]
-
-    @enforce_types
-    def size(self, actor: PydanticUser, agent_id: Optional[str] = None, **kwargs) -> int:
-        """Get the total count of messages with optional filters.
-
-        Args:
-            actor   : The user requesting the count
-            agent_id: The agent ID
-        """
-        with self.session_maker() as session:
-            return PassageModel.size(db_session=session, actor=actor, agent_id=agent_id, **kwargs)
+                # Try archival passages
+                try:
+                    passage = AgentPassage.read(db_session=session, identifier=passage_id, actor=actor)
+                    passage.hard_delete(session, actor=actor)
+                    return True
+                except NoResultFound:
+                    raise NoResultFound(f"Passage with id {passage_id} not found.")
 
     def delete_passages(
         self,
         actor: PydanticUser,
-        agent_id: Optional[str] = None,
-        file_id: Optional[str] = None,
-        start_date: Optional[datetime] = None,
-        end_date: Optional[datetime] = None,
-        limit: Optional[int] = 50,
-        cursor: Optional[str] = None,
-        query_text: Optional[str] = None,
-        source_id: Optional[str] = None,
+        passages: List[PydanticPassage],
     ) -> bool:
-
-        passages = self.list_passages(
-            actor=actor,
-            agent_id=agent_id,
-            file_id=file_id,
-            cursor=cursor,
-            limit=limit,
-            start_date=start_date,
-            end_date=end_date,
-            query_text=query_text,
-            source_id=source_id,
-        )
-
         # TODO: This is very inefficient
         # TODO: We should have a base `delete_all_matching_filters`-esque function
         for passage in passages:
             self.delete_passage_by_id(passage_id=passage.id, actor=actor)
+        return True