bugfix: Bug in which context_id wasn't stable across REPL messages (#480)

Co-authored-by: Douglas Reid <douglas-reid@users.noreply.github.com>

Author: eob

.gitignore

@@ -108,4 +108,6 @@ fabric.properties
 # Generated website
 docs/_build
 venv
+
+# Logs directory
 logs/

src/steamship/agents/examples/example_assistant.py

@@ -1,5 +1,3 @@
-import uuid
-
 from steamship.agents.functional import FunctionsBasedAgent
 from steamship.agents.llms.openai import ChatOpenAI
 from steamship.agents.schema.message_selectors import MessageWindowMessageSelector
@@ -31,4 +29,4 @@ def __init__(self, **kwargs):
     # AgentREPL provides a mechanism for local execution of an AgentService method.
     # This is used for simplified debugging as agents and tools are developed and
     # added.
-    AgentREPL(MyFunctionsBasedAssistant, agent_package_config={}).run(context_id=uuid.uuid4())
+    AgentREPL(MyFunctionsBasedAssistant, agent_package_config={}).run()

src/steamship/agents/examples/example_react_assistant.py

@@ -1,5 +1,3 @@
-import uuid
-
 from steamship.agents.llms.openai import OpenAI
 from steamship.agents.react import ReACTAgent
 from steamship.agents.schema.message_selectors import MessageWindowMessageSelector
@@ -33,4 +31,4 @@ def __init__(self, **kwargs):
     # AgentREPL provides a mechanism for local execution of an AgentService method.
     # This is used for simplified debugging as agents and tools are developed and
     # added.
-    AgentREPL(MyAssistant, agent_package_config={}).run(context_id=uuid.uuid4())
+    AgentREPL(MyAssistant, agent_package_config={}).run()

src/steamship/agents/service/agent_service.py

@@ -88,16 +88,21 @@ def run_agent(self, agent: Agent, context: AgentContext):
             func(action.output, context.metadata)
 
     @post("prompt")
-    def prompt(self, prompt: Optional[str] = None, **kwargs) -> List[Block]:
+    def prompt(
+        self, prompt: Optional[str] = None, context_id: Optional[str] = None, **kwargs
+    ) -> List[Block]:
         """Run an agent with the provided text as the input."""
         prompt = prompt or kwargs.get("question")
 
         # AgentContexts serve to allow the AgentService to run agents
         # with appropriate information about the desired tasking.
-        # Here, we create a new context on each prompt, and append the
-        # prompt to the message history stored in the context.
-        context_id = uuid.uuid4()
-        context = AgentContext.get_or_create(self.client, {"id": f"{context_id}"})
+        if context_id is None:
+            context_id = uuid.uuid4()
+            logging.warning(
+                f"No context_id was provided; generated {context_id}. This likely means no conversational history will be present."
+            )
+
+        context = AgentContext.get_or_create(self.client, context_keys={"id": f"{context_id}"})
         context.chat_history.append_user_message(prompt)
 
         # Add a default LLM

src/steamship/utils/repl.py

@@ -1,6 +1,7 @@
 import abc
 import contextlib
 import logging
+import uuid
 from abc import ABC
 from json import JSONDecodeError
 from typing import Any, Dict, List, Optional, Type, Union, cast
@@ -142,6 +143,7 @@ def run(self):
 class AgentREPL(SteamshipREPL):
     agent_class: Type[AgentService]
     agent_instance: Optional[AgentService]
+    context_id: Optional[str] = None
     client = Steamship
     config = None
 
@@ -151,6 +153,7 @@ def __init__(
         method: Optional[str] = None,
         agent_package_config: Optional[Dict[str, Any]] = None,
         client: Optional[Steamship] = None,
+        context_id: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -159,6 +162,7 @@ def __init__(
         self.client = client or Steamship()
         self.config = agent_package_config
         self.agent_instance = None
+        self.context_id = context_id or uuid.uuid4()
 
     def run_with_client(self, client: Steamship, **kwargs):
         try:
@@ -178,7 +182,7 @@ def colored(text: str, color: str, **kwargs):
 
         while True:
             input_text = input(colored(text="Input: ", color="blue"))  # noqa: F821
-            output = responder(input_text)
+            output = responder(prompt=input_text, context_id=self.context_id)
             self.print_object_or_objects(output)
 
     def run(self, **kwargs):
@@ -192,11 +196,19 @@ class HttpREPL(SteamshipREPL):
     prompt_url: Optional[AgentService]
     client = Steamship
     config = None
+    context_id: Optional[str] = None
 
-    def __init__(self, prompt_url: str, client: Optional[Steamship] = None, **kwargs):
+    def __init__(
+        self,
+        prompt_url: str,
+        context_id: Optional[str] = None,
+        client: Optional[Steamship] = None,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         self.prompt_url = prompt_url
         self.client = client or Steamship()
+        self.context_id = context_id or uuid.uuid4()
 
     def run_with_client(self, client: Steamship, **kwargs):  # noqa: C901
         try:
@@ -212,7 +224,7 @@ def colored(text: str, color: str):
             input_text = input(colored(text="Input: ", color="blue"))  # noqa: F821
             resp = requests.post(
                 self.prompt_url,
-                json={"prompt": input_text},
+                json={"prompt": input_text, "context_id": self.context_id},
                 headers={
                     "Content-Type": "application/json",
                     "Authorization": f"Bearer {self.client.config.api_key.get_secret_value()}",