Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Agents Refactor 2/2 (client side): Use capabilities for LLMs, agent impl. #583

Merged
merged 10 commits into from
Oct 24, 2023
50 changes: 50 additions & 0 deletions src/steamship/agents/basic_chat.py
Original file line number Diff line number Diff line change
@@ -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)
137 changes: 137 additions & 0 deletions src/steamship/agents/functions_based.py
Original file line number Diff line number Diff line change
@@ -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):
GitOnUp marked this conversation as resolved.
Show resolved Hide resolved
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.
Comment on lines +132 to +133
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Callout for removal

]
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)
58 changes: 58 additions & 0 deletions src/steamship/agents/tools/text_generation/custom_llm_prompt.py
Original file line number Diff line number Diff line change
@@ -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())
51 changes: 50 additions & 1 deletion src/steamship/agents/utils.py
Original file line number Diff line number Diff line change
@@ -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"

Expand All @@ -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
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pulled generally applicable parts out of functions based agent.

) -> 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
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm also noticing while doing plugin implementations that we tend to not dedupe between this and "default_system_prompt" in configs, which seems like it could confuse some models.

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
4 changes: 4 additions & 0 deletions src/steamship/data/block.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down
Loading
Loading