pass in user message only on first turn

letta-ai · Oct 10, 2024 · f76a149 · f76a149
1 parent f122bf2
commit f76a149
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/letta/o1_agent.py b/letta/o1_agent.py
@@ -21,7 +21,7 @@ def send_thinking_message(self: Agent, message: str) -> Optional[str]:
     Returns:
         Optional[str]: None is always returned as this function does not produce a response.
     """
-    self.interface.assistant_message(message, msg_obj=self._messages[-1])
+    self.interface.internal_monologue(message, msg_obj=self._messages[-1])
     return None
 
 
@@ -35,7 +35,7 @@ def send_final_message(self: Agent, message: str) -> Optional[str]:
     Returns:
         Optional[str]: None is always returned as this function does not produce a response.
     """
-    self.interface.assistant_message(message, msg_obj=self._messages[-1])
+    self.interface.internal_monologue(message, msg_obj=self._messages[-1])
     return None
 
 
@@ -45,7 +45,7 @@ def __init__(
         interface: AgentInterface,
         agent_state: AgentState,
         tools: List[Tool] = [],
-        max_thinking_steps: int = 5,
+        max_thinking_steps: int = 10,
         first_message_verify_mono: bool = False,
     ):
         super().__init__(interface, agent_state, tools)
@@ -75,13 +75,29 @@ def step(
             tools=self.tools,
             first_message_verify_mono=self.first_message_verify_mono,
         )
+        print("O1 AGENT", self.agent_state.system)
+        print("O1 AGENT", self.agent_state)
+
+        response = thinking_agent.step(
+            user_message,
+            first_message,
+            first_message_retry_limit,
+            skip_verify,
+            return_dicts,
+            recreate_message_timestamp,
+            stream,
+            timestamp,
+            inner_thoughts_in_kwargs,
+            ms,
+        )
+
+        print("THINKING AGENT", thinking_agent.agent_state.system)
         for _ in range(self.max_thinking_steps):
             # assert isinstance(self.agent_state.memory, Memory), f"Memory object is not of type Memory: {type(self.agent.agent_state.memory)}"
 
             # assert isinstance(thinking_agent.agent_state.memory, Memory), f"Memory object is not of type Memory"
-            # print("THINKING AGENT", thinking_agent.agent_state.system)
             response = thinking_agent.step(
-                user_message,
+                None,
                 first_message,
                 first_message_retry_limit,
                 skip_verify,

diff --git a/tests/test_o1_agent.py b/tests/test_o1_agent.py
@@ -36,14 +36,15 @@ def test_o1_agent():
     agent_state = client.create_agent(
         agent_type=AgentType.o1_agent,
         tools=[thinking_tool.name, final_tool.name],
-        llm_config=LLMConfig.default_config("gpt-4"),
+        llm_config=LLMConfig.default_config("gpt-4o-mini"),
         embedding_config=EmbeddingConfig.default_config("text-embedding-ada-002"),
         memory=ChatMemory(human=get_human_text(DEFAULT_HUMAN), persona=get_persona_text("o1_persona")),
     )
     agent = client.get_agent(agent_id=agent_state.id)
     assert agent is not None
 
-    response = client.user_message(agent_id=agent_state.id, message="How many Rs are there in strawberry?")
+    # response = client.user_message(agent_id=agent_state.id, message="How many Rs are there in strawberry?")
+    response = client.user_message(agent_id=agent_state.id, message="9.9 or 9.11, which is a larger number?")
     assert response is not None
     assert len(response.messages) > 3
     print("\n\n".join([str(i) for i in response.messages]))