Skip to content

Commit

Permalink
LLM Ouptut Parsing back into Blocks (#387)
Browse files Browse the repository at this point in the history
This allows Telegram (and other transports) to appropriately send them
  • Loading branch information
dkolas authored May 24, 2023
1 parent 3cb597a commit d8c1fed
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 8 deletions.
6 changes: 4 additions & 2 deletions src/steamship/agents/examples/my_assistant.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@ def prompt(self, prompt: str) -> str:

def sync_emit(blocks: List[Block], meta: Metadata):
nonlocal output
block_text = "\n".join([b.text for b in blocks if b.is_text()])
block_text = "\n".join(
[b.text if b.is_text() else f"({b.mime_type}: {b.id})" for b in blocks]
)
output += block_text

context.emit_funcs.append(sync_emit)
Expand All @@ -44,4 +46,4 @@ def sync_emit(blocks: List[Block], meta: Metadata):


if __name__ == "__main__":
AgentREPL(MyAssistant, "prompt").run()
AgentREPL(MyAssistant, "prompt", agent_package_config={}).run()
27 changes: 24 additions & 3 deletions src/steamship/agents/react/output_parser.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import re
from typing import Dict, Optional
from typing import Dict, List, Optional

from steamship import Block
from steamship import Block, Steamship
from steamship.agents.schema import Action, AgentContext, FinishAction, OutputParser, Tool


Expand All @@ -19,7 +19,9 @@ def parse(self, text: str, context: AgentContext) -> Action:
raise RuntimeError(f"Could not parse LLM output: `{text}`")

if "AI:" in text:
return FinishAction(output=[Block(text=text.split("AI:")[-1].strip())], context=context)
return FinishAction(
output=ReACTOutputParser._blocks_from_text(context.client, text), context=context
)

regex = r"Action: (.*?)[\n]*Action Input: (.*)"
match = re.search(regex, text)
Expand All @@ -35,3 +37,22 @@ def parse(self, text: str, context: AgentContext) -> Action:
input=[Block(text=action_input)],
context=context,
)

@staticmethod
def _blocks_from_text(client: Steamship, text: str) -> List[Block]:
last_response = text.split("AI:")[-1].strip()

block_id_regex = r"(?:\[Block)?\(?([A-F0-9]{8}\-[A-F0-9]{4}\-[A-F0-9]{4}\-[A-F0-9]{4}\-[A-F0-9]{12})\)?\]?"
remaining_text = last_response
result_blocks: List[Block] = []
while remaining_text is not None and len(remaining_text) > 0:
match = re.search(block_id_regex, remaining_text)
if match:
result_blocks.append(Block(text=remaining_text[0 : match.start()]))
result_blocks.append(Block.get(client, _id=match.group(1)))
remaining_text = remaining_text[match.end() :]
else:
result_blocks.append(Block(text=remaining_text))
remaining_text = ""

return result_blocks
8 changes: 5 additions & 3 deletions src/steamship/experimental/package_starters/telegram_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from pydantic import Field

from steamship import Block
from steamship.agents.schema import AgentContext, EmitFunc, Metadata
from steamship.agents.schema import Agent, AgentContext, EmitFunc, Metadata
from steamship.experimental.package_starters.web_agent import SteamshipWidgetAgentService
from steamship.experimental.package_starters.web_bot import response_for_exception
from steamship.experimental.transports import TelegramTransport
Expand All @@ -20,18 +20,20 @@ class TelegramBotConfig(Config):
class TelegramAgentService(SteamshipWidgetAgentService, ABC):
config: TelegramBotConfig
telegram_transport: TelegramTransport
incoming_message_agent: Agent

@classmethod
def config_cls(cls) -> Type[Config]:
"""Return the Configuration class."""
return TelegramBotConfig

def __init__(self, **kwargs):
def __init__(self, incoming_message_agent: Agent, **kwargs):
super().__init__(**kwargs)
self.api_root = f"https://api.telegram.org/bot{self.config.bot_token}"
self.telegram_transport = TelegramTransport(
bot_token=self.config.bot_token, client=self.client
)
self.incoming_message_agent = incoming_message_agent

def instance_init(self):
"""This instance init method is called automatically when an instance of this package is created. It registers the URL of the instance as the Telegram webhook for messages."""
Expand Down Expand Up @@ -62,7 +64,7 @@ def respond(self, **kwargs) -> InvocableResponse[str]:
if len(context.emit_funcs) == 0:
context.emit_funcs.append(self.build_emit_func(chat_id=chat_id))

response = self.run_agent(context)
response = self.run_agent(self.incoming_message_agent, context)
if response is not None:
self.telegram_transport.send(response, metadata={})
else:
Expand Down
Empty file.
37 changes: 37 additions & 0 deletions tests/steamship_tests/agents/test_react_parse_output.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import pytest

from steamship import Block, File, Steamship
from steamship.agents.react import ReACTOutputParser


@pytest.mark.usefixtures("client")
def test_parse_output(client: Steamship):

file = File.create(client, blocks=[Block(text="test"), Block(text="Another test")])
block_id = file.blocks[0].id
block_2_id = file.blocks[1].id

example1 = f"some text [Block({block_id})] some more text"

parsed_blocks1 = ReACTOutputParser._blocks_from_text(client, example1)
assert len(parsed_blocks1) == 3
assert parsed_blocks1[0].text == "some text "
assert parsed_blocks1[1].id == block_id
assert parsed_blocks1[2].text == " some more text"

example2 = f"some text [Block({block_id})] some more text {block_2_id} even more text"
parsed_blocks2 = ReACTOutputParser._blocks_from_text(client, example2)
assert len(parsed_blocks2) == 5
assert parsed_blocks2[0].text == "some text "
assert parsed_blocks2[1].id == block_id
assert parsed_blocks2[2].text == " some more text "
assert parsed_blocks2[3].id == block_2_id
assert parsed_blocks2[4].text == " even more text"

example3 = f"some text [Block({block_id})] some more text {block_2_id}"
parsed_blocks3 = ReACTOutputParser._blocks_from_text(client, example3)
assert len(parsed_blocks3) == 4
assert parsed_blocks3[0].text == "some text "
assert parsed_blocks3[1].id == block_id
assert parsed_blocks3[2].text == " some more text "
assert parsed_blocks3[3].id == block_2_id

0 comments on commit d8c1fed

Please sign in to comment.