import json
from operator import attrgetter
from typing import List, Optional
from steamship import Block, MimeTypes, Tag
from steamship.agents.functional.output_parser import FunctionsBasedOutputParser
from steamship.agents.schema import Action, AgentContext, ChatAgent, ChatLLM, FinishAction, Tool
from steamship.data.tags.tag_constants import ChatTag, RoleTag, TagKind, TagValueKey
from steamship.data.tags.tag_utils import get_tag
[docs]
class FunctionsBasedAgent(ChatAgent):
"""Selects actions for AgentService based on OpenAI Function style LLM Prompting."""
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, tools: List[Tool], llm: ChatLLM, **kwargs):
super().__init__(
output_parser=FunctionsBasedOutputParser(tools=tools), llm=llm, tools=tools, **kwargs
)
[docs]
def default_system_message(self) -> Optional[str]:
return self.PROMPT
def _get_or_create_system_message(self, context: AgentContext) -> Block:
if context.chat_history.last_system_message:
return context.chat_history.last_system_message
return context.chat_history.append_system_message(
text=self.default_system_message(), mime_type=MimeTypes.TXT
)
[docs]
def build_chat_history_for_tool(self, context: AgentContext) -> List[Block]:
# system message should have already been created in context, but we double-check for safety
sys_msg = self._get_or_create_system_message(context)
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(self.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)
# get working history (completed actions)
messages.extend(self._function_calls_since_last_user_message(context))
return messages
[docs]
def next_action(self, context: AgentContext) -> Action:
# Build the Chat History that we'll provide as input to the action
messages = self.build_chat_history_for_tool(context)
# Run the default LLM on those messages
output_blocks = self.llm.chat(messages=messages, tools=self.tools)
future_action = self.output_parser.parse(output_blocks[0].text, context)
if not isinstance(future_action, FinishAction):
# record the LLM's function response in history
self._record_action_selection(future_action, context)
return future_action
def _function_calls_since_last_user_message(self, context: AgentContext) -> List[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 _to_openai_function_selection(self, action: Action) -> str:
"""NOTE: Temporary placeholder. Should be refactored"""
fc = {"name": action.tool}
args = {}
for block in action.input:
for t in block.tags:
if t.kind == TagKind.FUNCTION_ARG:
args[t.name] = block.as_llm_input(exclude_block_wrapper=True)
fc["arguments"] = json.dumps(args) # the arguments must be a string value NOT a dict
return json.dumps(fc)
def _record_action_selection(self, action: Action, context: AgentContext):
tags = [
Tag(
kind=TagKind.CHAT,
name=ChatTag.ROLE,
value={TagValueKey.STRING_VALUE: RoleTag.ASSISTANT},
),
Tag(kind=TagKind.FUNCTION_SELECTION, name=action.tool),
]
context.chat_history.file.append_block(
text=self._to_openai_function_selection(action), tags=tags, mime_type=MimeTypes.TXT
)
[docs]
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},
),
# need the following tag for backwards compatibility with older gpt-4 plugin
Tag(
kind="name",
name=action.tool,
),
]
# TODO(dougreid): I'm not convinced this is correct for tools that return multiple values.
# It _feels_ like these should be named and inlined as a single message in history, etc.
for block in action.output:
context.chat_history.file.append_block(
text=block.as_llm_input(exclude_block_wrapper=True),
tags=tags,
mime_type=block.mime_type,
)