LLM Ouptut Parsing back into Blocks (#387)

This allows Telegram (and other transports) to appropriately send them

Author: dkolas

src/steamship/agents/examples/my_assistant.py

@@ -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)
@@ -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()

src/steamship/agents/react/output_parser.py

@@ -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
 
 
@@ -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)
@@ -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

src/steamship/experimental/package_starters/telegram_agent.py

@@ -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
@@ -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."""
@@ -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:

tests/steamship_tests/agents/__init__.py

@@ -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