From 686b59be497e6ee94a67d3763f876451410d09a5 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Wed, 11 Dec 2024 14:00:49 +0100
Subject: [PATCH] feat(agents): improve parsing in Streamlit agent

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 src/agents/experimental/streamlit/agent.ts |  2 +-
 src/internals/helpers/string.ts            | 47 ++++++++++++++++++++--
 tests/e2e/agents/streamlit.test.ts         |  4 +-
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/agents/experimental/streamlit/agent.ts b/src/agents/experimental/streamlit/agent.ts
index 13dc8aa9..c847bc05 100644
--- a/src/agents/experimental/streamlit/agent.ts
+++ b/src/agents/experimental/streamlit/agent.ts
@@ -170,7 +170,7 @@ export class StreamlitAgent extends BaseAgent<StreamlitRunInput, StreamlitRunOut
     for (let i = 0; i < raw.length; ) {
       const text = raw.substring(i);
 
-      const code = findFirstPair(text, ["```python-app\n", "```\n"]);
+      const code = findFirstPair(text, ["```python-app\n", "\n```"], { allowOverlap: true });
       if (!code) {
         blocks.push({ start: i, end: i + text.length, content: text, name: "text" });
         break;
diff --git a/src/internals/helpers/string.ts b/src/internals/helpers/string.ts
index 932dd5fe..22e8af00 100644
--- a/src/internals/helpers/string.ts
+++ b/src/internals/helpers/string.ts
@@ -17,7 +17,7 @@
 import { ValueOf } from "@/internals/types.js";
 import * as R from "remeda";
 import { ValueError } from "@/errors.js";
-import { unique } from "remeda";
+import { isString, unique } from "remeda";
 
 export function* splitString(
   text: string,
@@ -149,7 +149,25 @@ export function halveString(
   }
 }
 
-export function findFirstPair(text: string, pair: [string, string]) {
+export function countSharedStartEndLetters(a: string, b: string): number {
+  if (!isString(a) || !isString(b)) {
+    throw new ValueError("Provided values must be all strings.");
+  }
+
+  const minLength = Math.min(a.length, b.length);
+  for (let i = 0; i < minLength; i++) {
+    if (a.at((i + 1) * -1) !== b.at(i)) {
+      return i;
+    }
+  }
+  return minLength;
+}
+
+export function findFirstPair(
+  text: string,
+  pair: [string, string],
+  options: { allowOverlap?: boolean } = {},
+) {
   const [opening, closing] = pair || [];
   if (!pair || !opening || !closing) {
     throw new ValueError(`The "pair" parameter is required and must be non-empty!`);
@@ -157,6 +175,7 @@ export function findFirstPair(text: string, pair: [string, string]) {
 
   let balance = 0;
   let startIndex = -1;
+  const pairOverlap = options.allowOverlap ? countSharedStartEndLetters(opening, closing) : 0;
 
   const isSame = opening === closing;
   for (let index = 0; index < text.length; index++) {
@@ -165,19 +184,41 @@ export function findFirstPair(text: string, pair: [string, string]) {
         startIndex = index;
       }
       balance++;
+      if (!options.allowOverlap) {
+        index += opening.length - 1;
+      }
     } else if (text.substring(index, index + closing.length) === closing) {
       if (balance > 0) {
         balance--;
         if (balance === 0) {
+          const inner = {
+            start: startIndex + opening.length,
+            get end() {
+              let innerEnd = index;
+              const innerSize = innerEnd - this.start;
+
+              if (innerSize < 0) {
+                innerEnd = this.start;
+              } else {
+                innerEnd += pairOverlap;
+              }
+
+              return innerEnd;
+            },
+          };
+
           return {
             start: startIndex,
             end: index + closing.length,
             pair,
-            inner: text.substring(startIndex + opening.length, index),
+            inner: text.substring(inner.start, inner.end),
             outer: text.substring(startIndex, index + closing.length),
           };
         }
       }
+      if (!options.allowOverlap) {
+        index += closing.length - 1;
+      }
     }
   }
 
diff --git a/tests/e2e/agents/streamlit.test.ts b/tests/e2e/agents/streamlit.test.ts
index 2f82b103..288c060b 100644
--- a/tests/e2e/agents/streamlit.test.ts
+++ b/tests/e2e/agents/streamlit.test.ts
@@ -43,7 +43,9 @@ describe("Streamlit Agent", () => {
     let response: Awaited<ReturnType<typeof agent.run>>;
     try {
       response = await agent
-        .run({ prompt: `Generate me a minimalistic "Hello World" app.` })
+        .run({
+          prompt: `Generate a minimalistic "Hello World" app and provide some short explanation.`,
+        })
         .observe((emitter) => {
           emitter.registerCallbacks({
             newToken: callbacks.create("newToken", {