Output parsing for block IDs

Author: dkolas

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

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