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 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>; 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", {