Agents Refactor 2/2 (client side): Use capabilities for LLMs, agent i…

…mpl. (#583)

Implement SteamshipLLM and Plugin Capabilities via Agents.

- FunctionsBasedAgent in steamship.agents.functions_based, mostly copied
from the existing one, but using plugin capabilities.
- Example Tool for CustomLLMPrompt takes Steamship LLM, which assumes
you're providing a PluginInstance. That said, the PluginInstance might
be more direct.
- Factored chat history building out into shared code.
- This change needs plugins to adapt to new behavior, coming in another PR.

---------

Co-authored-by: Ted Benson <edward.benson@gmail.com>

src/steamship/agents/basic_chat.py

@@ -0,0 +1,50 @@
+from steamship import SteamshipError
+from steamship.agents.llms.steamship_llm import SteamshipLLM
+from steamship.agents.schema import Action, Agent, AgentContext, FinishAction
+from steamship.agents.utils import build_chat_history
+from steamship.data.tags.tag_constants import RoleTag
+from steamship.plugin.capabilities import ConversationSupport, RequestLevel, SystemPromptSupport
+
+DEFAULT_PROMPT = """You are a helpful AI assistant.
+You chat with users are eager to engage on a variety of topics. You answer their questions, provide help thinking
+through challenges they have, and engage as a trusted assistant."""
+
+
+class _BasicChatAgent(Agent):
+    """BasicChatAgent implements a conversational agent with no function calling or tool reasoning.
+
+    This class is under active development.
+
+    This Agent class is useful in a number of situations:
+    1) You are looking for an agent with nothing more than open-ended chat, possibly with personality and backstory.
+    2) You are using an LLM that does not support Function-calling natively
+    3) You are using an LLM that, in practice, performs poorly with ReACT-style prompts
+    In these cases, an agent whose only loop does not include an attempt to reason about tool invocation is ideal.
+    A secondary side effect of eliminating tooling is that the less processing must occur before a response begins.
+    """
+
+    PROMPT = DEFAULT_PROMPT
+
+    def __init__(self, llm: SteamshipLLM, **kwargs):
+        # Throw if the user has provided tools as a way to ensure the user understands the limitation of this agent.
+        if "tools" in kwargs:
+            raise SteamshipError(
+                "BasicChatAgent does not support tools. For tool-based agents, please use a base class such as FunctionsBasedAgent."
+            )
+
+        super().__init__(llm=llm, tools=[], **kwargs)
+        self.capabilities = [
+            SystemPromptSupport(),
+            ConversationSupport(request_level=RequestLevel.BEST_EFFORT),
+        ]
+        self.llm = llm
+
+    def next_action(self, context: AgentContext) -> Action:
+        # Build the Chat History that we'll provide as input to the action
+        messages = build_chat_history(self.PROMPT, self.message_selector, context)
+
+        output_blocks = self.llm.generate(messages=messages, capabilities=self.capabilities)
+
+        for block in output_blocks:
+            block.set_chat_role(RoleTag.ASSISTANT)
+        return FinishAction(output=output_blocks, context=context)

src/steamship/agents/functions_based.py

@@ -0,0 +1,137 @@
+from typing import Iterable, List, Optional
+
+from steamship import Block, MimeTypes, SteamshipError, Tag
+from steamship.agents.llms.steamship_llm import SteamshipLLM
+from steamship.agents.schema import Action, Agent, AgentContext, FinishAction, Tool
+from steamship.agents.utils import build_chat_history
+from steamship.data.tags.tag_constants import ChatTag, RoleTag, TagKind, TagValueKey
+from steamship.data.tags.tag_utils import get_tag
+from steamship.plugin.capabilities import (
+    ConversationSupport,
+    FunctionCallingSupport,
+    SystemPromptSupport,
+)
+
+
+class _FunctionsBasedAgent(Agent):
+    """Selects actions for AgentService based on a set of Tools.
+
+    This class is part of active development and not ready for usage yet.
+    """
+
+    PROMPT = """You are a helpful AI assistant.
+
+NOTE: Some functions return images, video, and audio files. These multimedia files will be represented in messages as
+UUIDs for Steamship Blocks. When responding directly to a user, you SHOULD print the Steamship Blocks for the images,
+video, or audio as follows: `Block(UUID for the block)`.
+
+Example response for a request that generated an image:
+Here is the image you requested: Block(288A2CA1-4753-4298-9716-53C1E42B726B).
+
+Only use the functions you have been provided with."""
+
+    def __init__(self, llm: SteamshipLLM, tools: List[Tool], **kwargs):
+        super().__init__(tools=tools, **kwargs)
+        self.llm = llm
+        self.capabilities = [
+            SystemPromptSupport(),
+            ConversationSupport(),
+            FunctionCallingSupport(tools=tools),
+        ]
+        self.tools_map = {tool.name: tool for tool in tools}
+
+    def default_system_message(self) -> Optional[str]:
+        return self.PROMPT
+
+    def next_action(self, context: AgentContext) -> Action:
+        # Build the Chat History that we'll provide as input to the action
+        messages = build_chat_history(self.default_system_message(), self.message_selector, context)
+        # get working history (completed actions)
+        messages.extend(self._function_calls_since_last_user_message(context))
+
+        # Run the default LLM on those messages
+        output_blocks = self.llm.generate(messages=messages, capabilities=self.capabilities)
+
+        for block in output_blocks:
+            if block.mime_type == MimeTypes.STEAMSHIP_PLUGIN_FUNCTION_CALL_INVOCATION:
+                invocation = FunctionCallingSupport.FunctionCallInvocation.from_block(block)
+                tool = self.tools_map.get(invocation.tool_name)
+                if tool is None:
+                    raise SteamshipError(
+                        f"LLM attempted to invoke tool {invocation.tool_name}, but {self.__class__.__name__} does not have a tool with that name."
+                    )
+                # TODO Block parse for input.  text/uuid is the default argument that we currently pass in for Tools.
+                #  As part of a refactor to allow for other parameters, this would need to change.
+                input_blocks = []
+                if text := invocation.args.get("text"):
+                    input_blocks.append(
+                        Block(
+                            text=text,
+                            tags=[Tag(kind=TagKind.FUNCTION_ARG, name="text")],
+                            mime_type=MimeTypes.TXT,
+                        )
+                    )
+                if uuid_arg := invocation.args.get("uuid"):
+                    existing_block = Block.get(context.client, _id=uuid_arg)
+                    tag = Tag.create(
+                        existing_block.client,
+                        file_id=existing_block.file_id,
+                        block_id=existing_block.id,
+                        kind=TagKind.FUNCTION_ARG,
+                        name="uuid",
+                    )
+                    existing_block.tags.append(tag)
+                    input_blocks.append(existing_block)
+                future_action = Action(tool=tool.name, input=input_blocks, output=None)
+                break
+        else:
+            future_action = FinishAction()
+            invocation = None
+        if not isinstance(future_action, FinishAction):
+            # record the LLM's function response in history
+            assert invocation
+            self._record_function_invocation(invocation, context)
+        return future_action
+
+    def _function_calls_since_last_user_message(self, context: AgentContext) -> Iterable[Block]:
+        function_calls = []
+        for block in context.chat_history.messages[::-1]:  # is this too inefficient at scale?
+            if block.chat_role == RoleTag.USER:
+                return reversed(function_calls)
+            if get_tag(block.tags, kind=TagKind.ROLE, name=RoleTag.FUNCTION):
+                function_calls.append(block)
+            elif get_tag(block.tags, kind=TagKind.FUNCTION_SELECTION):
+                function_calls.append(block)
+        return reversed(function_calls)
+
+    def _record_function_invocation(
+        self, invocation: FunctionCallingSupport.FunctionCallInvocation, context: AgentContext
+    ):
+        tags = [
+            Tag(
+                kind=TagKind.CHAT,
+                name=ChatTag.ROLE,
+                value={TagValueKey.STRING_VALUE: RoleTag.ASSISTANT},
+            ),
+            Tag(kind=TagKind.FUNCTION_SELECTION, name=invocation.tool_name),
+        ]
+        invocation.create_block(context.client, context.chat_history.file.id, tags=tags)
+
+    def record_action_run(self, action: Action, context: AgentContext):
+        super().record_action_run(action, context)
+
+        if isinstance(action, FinishAction):
+            return
+
+        tags = [
+            Tag(
+                kind=TagKind.ROLE,
+                name=RoleTag.FUNCTION,
+                value={TagValueKey.STRING_VALUE: action.tool},
+            ),
+            # TODO (PR): we're asserting capabilities support in next_action so the "name" tag is no longer needed for
+            #  backcompat as we won't be able to run against older versions anyway.
+        ]
+        output = [block.as_llm_input(exclude_block_wrapper=False) for block in action.output]
+        result = FunctionCallingSupport.FunctionCallResult(tool_name=action.tool, result=output)
+        result.create_block(context.client, context.chat_history.file.id, tags=tags)

src/steamship/agents/tools/text_generation/custom_llm_prompt.py

@@ -0,0 +1,58 @@
+from typing import Any, List, Union
+
+from steamship import Block, Steamship, Task
+from steamship.agents.llms.steamship_llm import SteamshipLLM
+from steamship.agents.schema import Tool
+from steamship.agents.schema.tool import AgentContext
+from steamship.utils.repl import ToolREPL
+
+DEFAULT_PROMPT = """You are a pirate who is stranded at sea.
+You love chatting with the people you meet while drifting about.
+SOMEONE ON A YACHT SAYS:
+{input}
+YOU REPLY:"""
+
+
+class _CustomLLMPromptTool(Tool):
+    """Example tool to illustrate how to respond with a different LLM than the one running your main agent.
+
+    This tool is under active development.
+
+    This can be useful if you are building an agent which:
+    - Wants to use tools, e.g. to generate images, which in practice requires GPT for you Agent Reasoning, but
+    - Also wants to use non-OpenAI LLM for some responses
+
+    In this case, one option is to direct the requests that deserve the non-GPT response to an instance of this
+    tool, which sets the `is_final` bit to `True`, causing its output to be considered the final Agent output.
+    Note that this is just a simple prompt-completion demo. If you want to build a Chat Agent with a non-GPT tool,
+    the best option is to see the steamship.agents.basic_chat.BasicChatAgent class, which incorporates chat history.
+    """
+
+    llm: SteamshipLLM
+    rewrite_prompt: str = DEFAULT_PROMPT
+    name: str = "CustomLLMTool"
+    human_description: str = "Applies a custom LLM to the user input."
+    agent_description: str = "Used to generate the answer for any user input. Always use every time tool or function selection is being performed."
+    is_final: bool = True
+
+    def run(self, tool_input: List[Block], context: AgentContext) -> Union[List[Block], Task[Any]]:
+        """Applies a custom LLM tool to the provided text block.
+
+        :param tool_input: A list of blocks to be rewritten if they contain text. Each block will be considered a separate input.
+        :param context: The active AgentContext.
+        :return: a list of Blocks whose content has been rewritten.  Synchronously produced (for now).
+        """
+        blocks = []
+        for block in tool_input:
+            if not block.is_text():
+                continue
+            prompt = self.rewrite_prompt.format(input=block.text)
+            output_blocks = self.llm.generate([Block(text=prompt)], assert_capabilities=False)
+            blocks.extend(output_blocks)
+
+        return blocks
+
+
+if __name__ == "__main__":
+    with Steamship.temporary_workspace() as client:
+        ToolREPL(_CustomLLMPromptTool()).run_with_client(client=client, context=AgentContext())

src/steamship/agents/utils.py

@@ -1,7 +1,10 @@
-from typing import Optional
+from operator import attrgetter
+from typing import List, Optional
 
+from steamship import Block, MimeTypes
 from steamship.agents.schema import AgentContext
 from steamship.agents.schema.llm import LLM
+from steamship.agents.schema.message_selectors import MessageSelector
 
 _LLM_KEY = "llm"
 
@@ -18,3 +21,49 @@ def with_llm(llm: LLM, context: Optional[AgentContext] = None) -> AgentContext:
 def get_llm(context: AgentContext, default: Optional[LLM] = None) -> Optional[LLM]:
     """Retrieves the LLM from the provided AgentContext (if it exists)."""
     return context.metadata.get(_LLM_KEY, default)
+
+
+def build_chat_history(
+    default_system_message: str, message_selector: MessageSelector, context: AgentContext
+) -> List[Block]:
+    # system message should have already been created in context, but we double-check for safety
+    if context.chat_history.last_system_message:
+        sys_msg = context.chat_history.last_system_message
+    else:
+        sys_msg = context.chat_history.append_system_message(
+            text=default_system_message, mime_type=MimeTypes.TXT
+        )
+    messages: List[Block] = [sys_msg]
+
+    messages_from_memory = []
+    # get prior conversations
+    if context.chat_history.is_searchable():
+        messages_from_memory.extend(
+            context.chat_history.search(context.chat_history.last_user_message.text, k=3)
+            .wait()
+            .to_ranked_blocks()
+        )
+        # TODO(dougreid): we need a way to threshold message inclusion, especially for small contexts
+
+    # get most recent context
+    messages_from_memory.extend(context.chat_history.select_messages(message_selector))
+
+    messages_from_memory.sort(key=attrgetter("index_in_file"))
+
+    # de-dupe the messages from memory
+    ids = [
+        sys_msg.id,
+        context.chat_history.last_user_message.id,
+    ]  # filter out last user message, it is appended afterwards
+    for msg in messages_from_memory:
+        if msg.id not in ids:
+            messages.append(msg)
+            ids.append(msg.id)
+
+    # TODO(dougreid): sort by dates? we SHOULD ensure ordering, given semantic search
+
+    # put the user prompt in the appropriate message location
+    # this should happen BEFORE any agent/assistant messages related to tool selection
+    messages.append(context.chat_history.last_user_message)
+
+    return messages

src/steamship/data/block.py

@@ -386,6 +386,10 @@ def abort_stream(self):
         )
 
 
+def is_block_id(value: str) -> bool:
+    return value.startswith("Block(") and value.endswith(")")
+
+
 class BlockQueryResponse(Response):
     blocks: List[Block]