letta-nightly 0.5.0.dev20241015104156__py3-none-any.whl → 0.5.0.dev20241017104103__py3-none-any.whl

letta/server/server.py

@@ -16,6 +16,7 @@ import letta.system as system
 from letta.agent import Agent, save_agent
 from letta.agent_store.db import attach_base
 from letta.agent_store.storage import StorageConnector, TableType
+from letta.client.utils import derive_function_name_regex
 from letta.credentials import LettaCredentials
 from letta.data_sources.connectors import DataConnector, load_data
 
@@ -72,9 +73,13 @@ from letta.schemas.file import FileMetadata
 from letta.schemas.job import Job
 from letta.schemas.letta_message import LettaMessage
 from letta.schemas.llm_config import LLMConfig
-from letta.schemas.memory import ArchivalMemorySummary, Memory, RecallMemorySummary
+from letta.schemas.memory import (
+    ArchivalMemorySummary,
+    ContextWindowOverview,
+    Memory,
+    RecallMemorySummary,
+)
 from letta.schemas.message import Message, MessageCreate, MessageRole, UpdateMessage
-from letta.schemas.openai.chat_completion_response import UsageStatistics
 from letta.schemas.organization import Organization, OrganizationCreate
 from letta.schemas.passage import Passage
 from letta.schemas.source import Source, SourceCreate, SourceUpdate
@@ -411,6 +416,7 @@ class SyncServer(Server):
             raise ValueError(f"messages should be a Message or a list of Message, got {type(input_messages)}")
 
         logger.debug(f"Got input messages: {input_messages}")
+        letta_agent = None
         try:
 
             # Get the agent object (loaded in memory)
@@ -422,83 +428,14 @@ class SyncServer(Server):
             token_streaming = letta_agent.interface.streaming_mode if hasattr(letta_agent.interface, "streaming_mode") else False
 
             logger.debug(f"Starting agent step")
-            no_verify = True
-            next_input_message = input_messages
-            counter = 0
-            total_usage = UsageStatistics()
-            step_count = 0
-            while True:
-                step_response = letta_agent.step(
-                    messages=next_input_message,
-                    first_message=False,
-                    skip_verify=no_verify,
-                    return_dicts=False,
-                    stream=token_streaming,
-                    # timestamp=timestamp,
-                    ms=self.ms,
-                )
-                step_response.messages
-                heartbeat_request = step_response.heartbeat_request
-                function_failed = step_response.function_failed
-                token_warning = step_response.in_context_memory_warning
-                usage = step_response.usage
-
-                step_count += 1
-                total_usage += usage
-                counter += 1
-                letta_agent.interface.step_complete()
-
-                logger.debug("Saving agent state")
-                # save updated state
-                save_agent(letta_agent, self.ms)
-
-                # Chain stops
-                if not self.chaining:
-                    logger.debug("No chaining, stopping after one step")
-                    break
-                elif self.max_chaining_steps is not None and counter > self.max_chaining_steps:
-                    logger.debug(f"Hit max chaining steps, stopping after {counter} steps")
-                    break
-                # Chain handlers
-                elif token_warning:
-                    assert letta_agent.agent_state.user_id is not None
-                    next_input_message = Message.dict_to_message(
-                        agent_id=letta_agent.agent_state.id,
-                        user_id=letta_agent.agent_state.user_id,
-                        model=letta_agent.model,
-                        openai_message_dict={
-                            "role": "user",  # TODO: change to system?
-                            "content": system.get_token_limit_warning(),
-                        },
-                    )
-                    continue  # always chain
-                elif function_failed:
-                    assert letta_agent.agent_state.user_id is not None
-                    next_input_message = Message.dict_to_message(
-                        agent_id=letta_agent.agent_state.id,
-                        user_id=letta_agent.agent_state.user_id,
-                        model=letta_agent.model,
-                        openai_message_dict={
-                            "role": "user",  # TODO: change to system?
-                            "content": system.get_heartbeat(constants.FUNC_FAILED_HEARTBEAT_MESSAGE),
-                        },
-                    )
-                    continue  # always chain
-                elif heartbeat_request:
-                    assert letta_agent.agent_state.user_id is not None
-                    next_input_message = Message.dict_to_message(
-                        agent_id=letta_agent.agent_state.id,
-                        user_id=letta_agent.agent_state.user_id,
-                        model=letta_agent.model,
-                        openai_message_dict={
-                            "role": "user",  # TODO: change to system?
-                            "content": system.get_heartbeat(constants.REQ_HEARTBEAT_MESSAGE),
-                        },
-                    )
-                    continue  # always chain
-                # Letta no-op / yield
-                else:
-                    break
+            usage_stats = letta_agent.step(
+                messages=input_messages,
+                chaining=self.chaining,
+                max_chaining_steps=self.max_chaining_steps,
+                stream=token_streaming,
+                ms=self.ms,
+                skip_verify=True,
+            )
 
         except Exception as e:
             logger.error(f"Error in server._step: {e}")
@@ -506,9 +443,10 @@ class SyncServer(Server):
             raise
         finally:
             logger.debug("Calling step_yield()")
-            letta_agent.interface.step_yield()
+            if letta_agent:
+                letta_agent.interface.step_yield()
 
-        return LettaUsageStatistics(**total_usage.model_dump(), step_count=step_count)
+        return usage_stats
 
     def _command(self, user_id: str, agent_id: str, command: str) -> LettaUsageStatistics:
         """Process a CLI command"""
@@ -794,7 +732,7 @@ class SyncServer(Server):
                 message_objects.append(message)
 
         else:
-            raise ValueError(f"All messages must be of type Message or MessageCreate, got {type(messages)}")
+            raise ValueError(f"All messages must be of type Message or MessageCreate, got {[type(message) for message in messages]}")
 
         # Run the agent state forward
         return self._step(user_id=user_id, agent_id=agent_id, input_messages=message_objects)
@@ -1028,6 +966,80 @@ class SyncServer(Server):
         # TODO: probably reload the agent somehow?
         return letta_agent.agent_state
 
+    def add_tool_to_agent(
+        self,
+        agent_id: str,
+        tool_id: str,
+        user_id: str,
+    ):
+        """Update the agents core memory block, return the new state"""
+        if self.ms.get_user(user_id=user_id) is None:
+            raise ValueError(f"User user_id={user_id} does not exist")
+        if self.ms.get_agent(agent_id=agent_id) is None:
+            raise ValueError(f"Agent agent_id={agent_id} does not exist")
+
+        # Get the agent object (loaded in memory)
+        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+
+        # Get all the tool objects from the request
+        tool_objs = []
+        tool_obj = self.ms.get_tool(tool_id=tool_id, user_id=user_id)
+        assert tool_obj, f"Tool with id={tool_id} does not exist"
+        tool_objs.append(tool_obj)
+
+        for tool in letta_agent.tools:
+            tool_obj = self.ms.get_tool(tool_id=tool.id, user_id=user_id)
+            assert tool_obj, f"Tool with id={tool.id} does not exist"
+
+            # If it's not the already added tool
+            if tool_obj.id != tool_id:
+                tool_objs.append(tool_obj)
+
+        # replace the list of tool names ("ids") inside the agent state
+        letta_agent.agent_state.tools = [tool.name for tool in tool_objs]
+
+        # then attempt to link the tools modules
+        letta_agent.link_tools(tool_objs)
+
+        # save the agent
+        save_agent(letta_agent, self.ms)
+        return letta_agent.agent_state
+
+    def remove_tool_from_agent(
+        self,
+        agent_id: str,
+        tool_id: str,
+        user_id: str,
+    ):
+        """Update the agents core memory block, return the new state"""
+        if self.ms.get_user(user_id=user_id) is None:
+            raise ValueError(f"User user_id={user_id} does not exist")
+        if self.ms.get_agent(agent_id=agent_id) is None:
+            raise ValueError(f"Agent agent_id={agent_id} does not exist")
+
+        # Get the agent object (loaded in memory)
+        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+
+        # Get all the tool_objs
+        tool_objs = []
+        for tool in letta_agent.tools:
+            tool_obj = self.ms.get_tool(tool_id=tool.id, user_id=user_id)
+            assert tool_obj, f"Tool with id={tool.id} does not exist"
+
+            # If it's not the tool we want to remove
+            if tool_obj.id != tool_id:
+                tool_objs.append(tool_obj)
+
+        # replace the list of tool names ("ids") inside the agent state
+        letta_agent.agent_state.tools = [tool.name for tool in tool_objs]
+
+        # then attempt to link the tools modules
+        letta_agent.link_tools(tool_objs)
+
+        # save the agent
+        save_agent(letta_agent, self.ms)
+        return letta_agent.agent_state
+
     def _agent_state_to_config(self, agent_state: AgentState) -> dict:
         """Convert AgentState to a dict for a JSON response"""
         assert agent_state is not None
@@ -1683,6 +1695,9 @@ class SyncServer(Server):
 
         return job
 
+    def delete_file_from_source(self, source_id: str, file_id: str, user_id: Optional[str]) -> Optional[FileMetadata]:
+        return self.ms.delete_file_from_source(source_id=source_id, file_id=file_id, user_id=user_id)
+
     def load_data(
         self,
         user_id: str,
@@ -1811,6 +1826,15 @@ class SyncServer(Server):
         """Get tool by ID."""
         return self.ms.get_tool(tool_id=tool_id)
 
+    def tool_with_name_and_user_id_exists(self, tool: Tool, user_id: Optional[str] = None) -> bool:
+        """Check if tool exists"""
+        tool = self.ms.get_tool_with_name_and_user_id(tool_name=tool.name, user_id=user_id)
+
+        if tool is None:
+            return False
+        else:
+            return True
+
     def get_tool_id(self, name: str, user_id: str) -> Optional[str]:
         """Get tool ID from name and user_id."""
         tool = self.ms.get_tool(tool_name=name, user_id=user_id)
@@ -1818,16 +1842,27 @@ class SyncServer(Server):
             return None
         return tool.id
 
-    def update_tool(
-        self,
-        request: ToolUpdate,
-    ) -> Tool:
+    def update_tool(self, request: ToolUpdate, user_id: Optional[str] = None) -> Tool:
         """Update an existing tool"""
-        existing_tool = self.ms.get_tool(tool_id=request.id)
-        if not existing_tool:
-            raise ValueError(f"Tool does not exist")
+        if request.name:
+            existing_tool = self.ms.get_tool_with_name_and_user_id(tool_name=request.name, user_id=user_id)
+            if existing_tool is None:
+                raise ValueError(f"Tool with name={request.name}, user_id={user_id} does not exist")
+        else:
+            existing_tool = self.ms.get_tool(tool_id=request.id)
+            if existing_tool is None:
+                raise ValueError(f"Tool with id={request.id} does not exist")
+
+        # Preserve the original tool id
+        # As we can override the tool id as well
+        # This is probably bad design if this is exposed to users...
+        original_id = existing_tool.id
 
         # override updated fields
+        if request.id:
+            existing_tool.id = request.id
+        if request.description:
+            existing_tool.description = request.description
         if request.source_code:
             existing_tool.source_code = request.source_code
         if request.source_type:
@@ -1836,10 +1871,15 @@ class SyncServer(Server):
             existing_tool.tags = request.tags
         if request.json_schema:
             existing_tool.json_schema = request.json_schema
+
+        # If name is explicitly provided here, overide the tool name
         if request.name:
             existing_tool.name = request.name
+        # Otherwise, if there's no name, and there's source code, we try to derive the name
+        elif request.source_code:
+            existing_tool.name = derive_function_name_regex(request.source_code)
 
-        self.ms.update_tool(existing_tool)
+        self.ms.update_tool(original_id, existing_tool)
         return self.ms.get_tool(tool_id=request.id)
 
     def create_tool(self, request: ToolCreate, user_id: Optional[str] = None, update: bool = True) -> Tool:  # TODO: add other fields
@@ -1866,7 +1906,7 @@ class SyncServer(Server):
 
             # TODO: not sure if this always works
             func = env[functions[-1]]
-            json_schema = generate_schema(func)
+            json_schema = generate_schema(func, terminal=request.terminal)
         else:
             # provided by client
             json_schema = request.json_schema
@@ -1877,15 +1917,23 @@ class SyncServer(Server):
             assert request.name, f"Tool name must be provided in json_schema {json_schema}. This should never happen."
 
         # check if already exists:
-        existing_tool = self.ms.get_tool(tool_name=request.name, user_id=user_id)
+        existing_tool = self.ms.get_tool(tool_id=request.id, tool_name=request.name, user_id=user_id)
         if existing_tool:
             if update:
-                updated_tool = self.update_tool(ToolUpdate(id=existing_tool.id, **vars(request)))
+                # id is an optional field, so we will fill it with the existing tool id
+                if not request.id:
+                    request.id = existing_tool.id
+                updated_tool = self.update_tool(ToolUpdate(**vars(request)), user_id)
                 assert updated_tool is not None, f"Failed to update tool {request.name}"
                 return updated_tool
             else:
                 raise ValueError(f"Tool {request.name} already exists and update=False")
 
+        # check for description
+        description = None
+        if request.description:
+            description = request.description
+
         tool = Tool(
             name=request.name,
             source_code=request.source_code,
@@ -1893,9 +1941,14 @@ class SyncServer(Server):
             tags=request.tags,
             json_schema=json_schema,
             user_id=user_id,
+            description=description,
         )
+
+        if request.id:
+            tool.id = request.id
+
         self.ms.create_tool(tool)
-        created_tool = self.ms.get_tool(tool_name=request.name, user_id=user_id)
+        created_tool = self.ms.get_tool(tool_id=tool.id, user_id=user_id)
         return created_tool
 
     def delete_tool(self, tool_id: str):
@@ -2087,3 +2140,13 @@ class SyncServer(Server):
 
     def add_embedding_model(self, request: EmbeddingConfig) -> EmbeddingConfig:
         """Add a new embedding model"""
+
+    def get_agent_context_window(
+        self,
+        user_id: str,
+        agent_id: str,
+    ) -> ContextWindowOverview:
+
+        # Get the current message
+        letta_agent = self._get_or_load_agent(agent_id=agent_id)
+        return letta_agent.get_context_window()