Add a few standardizing methods to AgentService (#513)

A handful of user support conversations over the weeks have related to:

1. How do I fetch "my agent" from the service, and
2. How do I modify / build the AgentContext in "situation Y"

This PR:

1. Adds a `get_default_agent` method to AgentService that acts as an
official way to get the agent (before it just had a special name, but
the code makes it unclear if it was `self.agent` or `self._agent`
2. Adds a `build_default_context` method which implements the standard
context that's been copy-pasted every time we write a new endpoint. Now
these endpoints can all rely on this `build_default_context` method, and
developers can be told that overriding that method will -- for example
-- alter the mechanics of their Telegram bot

Author: eob

src/steamship/agents/examples/telegram_bot.py

@@ -76,7 +76,7 @@ def __init__(self, **kwargs):
                 client=self.client,
                 config=TelegramTransportConfig(bot_token=self.config.bot_token),
                 agent_service=self,
-                agent=self._agent,
+                agent=self.get_default_agent(),
             )
         )
 

src/steamship/agents/mixins/transports/slack.py

@@ -9,7 +9,7 @@
 from steamship import Block, Steamship
 from steamship.agents.llms import OpenAI
 from steamship.agents.mixins.transports.transport import Transport
-from steamship.agents.schema import Agent, AgentContext, EmitFunc, Metadata
+from steamship.agents.schema import Agent, EmitFunc, Metadata
 from steamship.agents.service.agent_service import AgentService
 from steamship.agents.utils import with_llm
 from steamship.invocable import Config, InvocableResponse, InvocationContext, get, post
@@ -370,12 +370,7 @@ def _respond_to_block(self, incoming_message: Block):
         """Respond to a single inbound message from Slack, posting the response back to Slack."""
         try:
             chat_id = incoming_message.chat_id
-
-            if not chat_id:
-                logging.error(f"No chat id on incoming block {incoming_message}")
-                return
-            # TODO: It feels like context is something the Agent should be providing.
-            context = AgentContext.get_or_create(self.client, context_keys={"chat_id": chat_id})
+            context = self.agent_service.build_default_context(context_id=chat_id)
 
             context.chat_history.append_user_message(
                 text=incoming_message.text, tags=incoming_message.tags

src/steamship/agents/mixins/transports/steamship_widget.py

@@ -2,11 +2,9 @@
 from typing import List, Optional
 
 from steamship import Block, Steamship, SteamshipError
-from steamship.agents.llms import OpenAI
 from steamship.agents.mixins.transports.transport import Transport
-from steamship.agents.schema import Agent, AgentContext, Metadata
+from steamship.agents.schema import Agent, Metadata
 from steamship.agents.service.agent_service import AgentService
-from steamship.agents.utils import with_llm
 from steamship.invocable import Config, InvocationContext, post
 
 API_BASE = "https://api.telegram.org/bot"
@@ -60,23 +58,14 @@ def _parse_inbound(self, payload: dict, context: Optional[dict] = None) -> Optio
     def answer(self, **payload) -> List[Block]:
         """Endpoint that implements the contract for Steamship embeddable chat widgets. This is a PUBLIC endpoint since these webhooks do not pass a token."""
         incoming_message = self.parse_inbound(payload)
-        context = AgentContext.get_or_create(
-            self.client, context_keys={"chat_id": incoming_message.chat_id}
-        )
+
+        context = self.agent_service.build_default_context(context_id=incoming_message.chat_id)
+
         context.chat_history.append_user_message(
             text=incoming_message.text, tags=incoming_message.tags
         )
         context.emit_funcs = [self.save_for_emit]
 
-        # Add an LLM to the context, using the Agent's if it exists.
-        llm = None
-        if hasattr(self.agent, "llm"):
-            llm = self.agent.llm
-        else:
-            llm = OpenAI(client=self.client)
-
-        context = with_llm(context=context, llm=llm)
-
         try:
             self.agent_service.run_agent(self.agent, context)
         except Exception as e:

src/steamship/agents/mixins/transports/telegram.py

@@ -6,11 +6,9 @@
 from pydantic import Field
 
 from steamship import Block, Steamship, SteamshipError
-from steamship.agents.llms import OpenAI
 from steamship.agents.mixins.transports.transport import Transport
-from steamship.agents.schema import Agent, AgentContext, EmitFunc, Metadata
+from steamship.agents.schema import Agent, EmitFunc, Metadata
 from steamship.agents.service.agent_service import AgentService
-from steamship.agents.utils import with_llm
 from steamship.invocable import Config, InvocableResponse, InvocationContext, post
 from steamship.utils.kv_store import KeyValueStore
 
@@ -211,21 +209,13 @@ def telegram_respond(self, **kwargs) -> InvocableResponse[str]:
         try:
             incoming_message = self.parse_inbound(message)
             if incoming_message is not None:
-                context = AgentContext.get_or_create(self.client, context_keys={"chat_id": chat_id})
+                context = self.agent_service.build_default_context(chat_id)
+
                 context.chat_history.append_user_message(
                     text=incoming_message.text, tags=incoming_message.tags
                 )
                 context.emit_funcs = [self.build_emit_func(chat_id=chat_id)]
 
-                # Add an LLM to the context, using the Agent's if it exists.
-                llm = None
-                if hasattr(self.agent, "llm"):
-                    llm = self.agent.llm
-                else:
-                    llm = OpenAI(client=self.client)
-
-                context = with_llm(context=context, llm=llm)
-
                 response = self.agent_service.run_agent(self.agent, context)
                 if response is not None:
                     self.send(response)

src/steamship/agents/service/agent_service.py

@@ -3,7 +3,7 @@
 from typing import List, Optional
 
 from steamship import Block, SteamshipError, Task
-from steamship.agents.llms.openai import ChatOpenAI
+from steamship.agents.llms.openai import ChatOpenAI, OpenAI
 from steamship.agents.logging import AgentLogging
 from steamship.agents.schema import Action, Agent, FinishAction
 from steamship.agents.schema.context import AgentContext, Metadata
@@ -165,12 +165,33 @@ def run_agent(self, agent: Agent, context: AgentContext):
             logging.info(f"Emitting via function: {func.__name__}")
             func(action.output, context.metadata)
 
-    @post("prompt")
-    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")
+    def get_default_agent(self, throw_if_missing: bool = True) -> Optional[Agent]:
+        """Return the default agent of this agent service.
+
+        This is a helper wrapper to safeguard naming conventions that have formed.
+        """
+        if hasattr(self, "agent"):
+            return self.agent
+        elif hasattr(self, "_agent"):
+            return self._agent
+        else:
+            if throw_if_missing:
+                raise SteamshipError(
+                    message="No Agent object found in the Agent Service. "
+                    "Please name it either self.agent or self._agent."
+                )
+            else:
+                return None
+
+    def build_default_context(self, context_id: Optional[str] = None, **kwargs) -> AgentContext:
+        """Build's the agent's default context.
+
+        The provides a single place to implement (or override) the default context that will be used by endpoints
+        that transports define. This allows an Agent developer to use, eg, the TelegramTransport but with a custom
+        type of memory or caching.
+
+        The returned context does not have any emit functions yet registered to it.
+        """
 
         # AgentContexts serve to allow the AgentService to run agents
         # with appropriate information about the desired tasking.
@@ -194,6 +215,28 @@ def prompt(
             use_llm_cache=use_llm_cache,
             use_action_cache=use_action_cache,
         )
+
+        # Add a default LLM to the context, using the Agent's if it exists.
+        llm = None
+        if agent := self.get_default_agent():
+            if hasattr(agent, "llm"):
+                llm = agent.llm
+        if llm is None:
+            llm = OpenAI(client=self.client)
+
+        context = with_llm(context=context, llm=llm)
+
+        return context
+
+    @post("prompt")
+    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")
+
+        context = self.build_default_context(context_id, **kwargs)
+
         context.chat_history.append_user_message(prompt)
 
         # Add a default LLM
@@ -215,18 +258,7 @@ def sync_emit(blocks: List[Block], meta: Metadata):
         context.emit_funcs.append(sync_emit)
 
         # Get the agent
-        agent: Optional[Agent] = None
-        if hasattr(self, "agent"):
-            agent = self.agent
-        elif hasattr(self, "_agent"):
-            agent = self._agent
-
-        if not agent:
-            raise SteamshipError(
-                message="No Agent object found in the Agent Service. "
-                "Please name it either self.agent or self._agent."
-            )
-
+        agent: Optional[Agent] = self.get_default_agent()
         self.run_agent(agent, context)
 
         # Now append the output blocks to the chat history