letta-nightly 0.6.2.dev20241210104242__py3-none-any.whl → 0.6.3.dev20241211050151__py3-none-any.whl

letta/services/passage_manager.py

@@ -0,0 +1,225 @@
+from typing import List, Optional, Dict, Tuple
+from letta.constants import MAX_EMBEDDING_DIM
+from datetime import datetime
+import numpy as np
+
+from letta.orm.errors import NoResultFound
+from letta.utils import enforce_types
+
+from letta.embeddings import embedding_model, parse_and_chunk_text
+from letta.schemas.embedding_config import EmbeddingConfig
+
+from letta.orm.passage import Passage as PassageModel
+from letta.orm.sqlalchemy_base import AccessType
+from letta.schemas.agent import AgentState
+from letta.schemas.passage import Passage as PydanticPassage
+from letta.schemas.user import User as PydanticUser
+
+class PassageManager:
+    """Manager class to handle business logic related to Passages."""
+
+    def __init__(self):
+        from letta.server.server import db_context
+        self.session_maker = db_context
+
+    @enforce_types
+    def get_passage_by_id(self, passage_id: str, actor: PydanticUser) -> Optional[PydanticPassage]:
+        """Fetch a passage by ID."""
+        with self.session_maker() as session:
+            try:
+                passage = PassageModel.read(db_session=session, identifier=passage_id, actor=actor)
+                return passage.to_pydantic()
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    def create_passage(self, pydantic_passage: PydanticPassage, actor: PydanticUser) -> PydanticPassage:
+        """Create a new passage."""
+        with self.session_maker() as session:
+            passage = PassageModel(**pydantic_passage.model_dump())
+            passage.create(session, actor=actor)
+            return passage.to_pydantic()
+
+    @enforce_types
+    def create_many_passages(self, passages: List[PydanticPassage], actor: PydanticUser) -> List[PydanticPassage]:
+        """Create multiple passages."""
+        return [self.create_passage(p, actor) for p in passages]
+
+    @enforce_types
+    def insert_passage(self, 
+        agent_state: AgentState,
+        agent_id: str,
+        text: str, 
+        actor: PydanticUser, 
+        return_ids: bool = False
+    ) -> List[PydanticPassage]:
+        """ Insert passage(s) into archival memory """
+
+        embedding_chunk_size = agent_state.embedding_config.embedding_chunk_size
+        embed_model = embedding_model(agent_state.embedding_config)
+
+        passages = []
+        
+        try:
+            # breakup string into passages
+            for text in parse_and_chunk_text(text, embedding_chunk_size):
+                embedding = embed_model.get_text_embedding(text)
+                if isinstance(embedding, dict):
+                    try:
+                        embedding = embedding["data"][0]["embedding"]
+                    except (KeyError, IndexError):
+                        # TODO as a fallback, see if we can find any lists in the payload
+                        raise TypeError(
+                            f"Got back an unexpected payload from text embedding function, type={type(embedding)}, value={embedding}"
+                        )
+                passage = self.create_passage(
+                    PydanticPassage(
+                        organization_id=actor.organization_id,
+                        agent_id=agent_id,
+                        text=text,
+                        embedding=embedding,
+                        embedding_config=agent_state.embedding_config
+                    ),
+                    actor=actor
+                )
+                passages.append(passage)
+
+            ids = [str(p.id) for p in passages]
+
+            if return_ids:
+                return ids
+            
+            return passages
+
+        except Exception as e:
+            raise e
+
+    @enforce_types
+    def update_passage_by_id(self, passage_id: str, passage: PydanticPassage, actor: PydanticUser, **kwargs) -> Optional[PydanticPassage]:
+        """Update a passage."""
+        if not passage_id:
+            raise ValueError("Passage ID must be provided.")
+
+        with self.session_maker() as session:
+            try:
+                # 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.")
+
+                # Update the database record with values from the provided record
+                update_data = passage.model_dump(exclude_unset=True, exclude_none=True)
+                for key, value in update_data.items():
+                    setattr(curr_passage, key, value)
+
+                # Commit changes
+                curr_passage.update(session, actor=actor)
+                return curr_passage.to_pydantic()
+            except NoResultFound:
+                return None
+
+    @enforce_types
+    def delete_passage_by_id(self, passage_id: str, actor: PydanticUser) -> bool:
+        """Delete a passage."""
+        if not passage_id:
+            raise ValueError("Passage ID must be provided.")
+
+        with self.session_maker() as session:
+            try:
+                passage = PassageModel.read(db_session=session, identifier=passage_id, actor=actor)
+                passage.hard_delete(session, actor=actor)
+            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,
+                      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,
+                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)
+
+    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
+                       ) -> 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)
+        
+        for passage in passages:
+            self.delete_passage_by_id(passage_id=passage.id, actor=actor)

letta/services/source_manager.py

@@ -64,7 +64,7 @@ class SourceManager:
             return source.to_pydantic()
 
     @enforce_types
-    def list_sources(self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50) -> List[PydanticSource]:
+    def list_sources(self, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50, **kwargs) -> List[PydanticSource]:
         """List all sources with optional pagination."""
         with self.session_maker() as session:
             sources = SourceModel.list(
@@ -72,6 +72,7 @@ class SourceManager:
                 cursor=cursor,
                 limit=limit,
                 organization_id=actor.organization_id,
+                **kwargs,
             )
             return [source.to_pydantic() for source in sources]
 

letta/services/tool_execution_sandbox.py

@@ -127,11 +127,12 @@ class ToolExecutionSandbox:
 
         # Save the old stdout
         old_stdout = sys.stdout
+        old_stderr = sys.stderr
         try:
             if local_configs.use_venv:
                 return self.run_local_dir_sandbox_venv(sbx_config, env, temp_file_path)
             else:
-                return self.run_local_dir_sandbox_runpy(sbx_config, env_vars, temp_file_path, old_stdout)
+                return self.run_local_dir_sandbox_runpy(sbx_config, env_vars, temp_file_path, old_stdout, old_stderr)
         except Exception as e:
             logger.error(f"Executing tool {self.tool_name} has an unexpected error: {e}")
             logger.error(f"Logging out tool {self.tool_name} auto-generated code for debugging: \n\n{code}")
@@ -139,6 +140,7 @@ class ToolExecutionSandbox:
         finally:
             # Clean up the temp file and restore stdout
             sys.stdout = old_stdout
+            sys.stderr = old_stderr
             os.remove(temp_file_path)
 
     def run_local_dir_sandbox_venv(self, sbx_config: SandboxConfig, env: Dict[str, str], temp_file_path: str) -> SandboxRunResult:
@@ -201,7 +203,11 @@ class ToolExecutionSandbox:
             func_result, stdout = self.parse_out_function_results_markers(result.stdout)
             func_return, agent_state = self.parse_best_effort(func_result)
             return SandboxRunResult(
-                func_return=func_return, agent_state=agent_state, stdout=[stdout], sandbox_config_fingerprint=sbx_config.fingerprint()
+                func_return=func_return, 
+                agent_state=agent_state,
+                stdout=[stdout],
+                stderr=[result.stderr],
+                sandbox_config_fingerprint=sbx_config.fingerprint(),
             )
         except subprocess.TimeoutExpired:
             raise TimeoutError(f"Executing tool {self.tool_name} has timed out.")
@@ -213,11 +219,13 @@ class ToolExecutionSandbox:
             raise e
 
     def run_local_dir_sandbox_runpy(
-        self, sbx_config: SandboxConfig, env_vars: Dict[str, str], temp_file_path: str, old_stdout: TextIO
+        self, sbx_config: SandboxConfig, env_vars: Dict[str, str], temp_file_path: str, old_stdout: TextIO, old_stderr: TextIO
     ) -> SandboxRunResult:
-        # Redirect stdout to capture script output
+        # Redirect stdout and stderr to capture script output
         captured_stdout = io.StringIO()
+        captured_stderr = io.StringIO()
         sys.stdout = captured_stdout
+        sys.stderr = captured_stderr
 
         # Execute the temp file
         with self.temporary_env_vars(env_vars):
@@ -227,14 +235,17 @@ class ToolExecutionSandbox:
         func_result = result.get(self.LOCAL_SANDBOX_RESULT_VAR_NAME)
         func_return, agent_state = self.parse_best_effort(func_result)
 
-        # Restore stdout and collect captured output
+        # Restore stdout and stderr and collect captured output
         sys.stdout = old_stdout
+        sys.stderr = old_stderr
         stdout_output = captured_stdout.getvalue()
+        stderr_output = captured_stderr.getvalue()
 
         return SandboxRunResult(
             func_return=func_return,
             agent_state=agent_state,
             stdout=[stdout_output],
+            stderr=[stderr_output],
             sandbox_config_fingerprint=sbx_config.fingerprint(),
         )
 
@@ -294,7 +305,8 @@ class ToolExecutionSandbox:
             return SandboxRunResult(
                 func_return=func_return,
                 agent_state=agent_state,
-                stdout=execution.logs.stdout + execution.logs.stderr,
+                stdout=execution.logs.stdout,
+                stderr=execution.logs.stderr,
                 sandbox_config_fingerprint=sbx_config.fingerprint(),
             )
 

letta/settings.py

@@ -17,6 +17,8 @@ class ToolSettings(BaseSettings):
 
 class ModelSettings(BaseSettings):
 
+    model_config = SettingsConfigDict(env_file='.env')
+
     # env_prefix='my_prefix_'
 
     # when we use /completions APIs (instead of /chat/completions), we need to specify a model wrapper