From 331387efa4c87193600562efc5d423b2ed27fc73 Mon Sep 17 00:00:00 2001
From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com>
Date: Mon, 8 Dec 2025 13:07:52 +0530
Subject: [PATCH] feat(search): add better context handling

---
 src/lib/agents/search/index.ts | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/lib/agents/search/index.ts b/src/lib/agents/search/index.ts
index bc6ff5b..91bed9a 100644
--- a/src/lib/agents/search/index.ts
+++ b/src/lib/agents/search/index.ts
@@ -55,19 +55,20 @@ class SearchAgent {
     });
 
     const finalContext =
-      searchResults?.findings
-        .filter((f) => f.type === 'search_results')
-        .flatMap((f) => f.results)
-        .map((f) => `${f.metadata.title}: ${f.content}`)
+      searchResults?.searchFindings
+        .map(
+          (f, index) =>
+            `<result index=${index} title=${f.metadata.title}>${f.content}</result>`,
+        )
         .join('\n') || '';
 
     const widgetContext = widgetOutputs
       .map((o) => {
-        return `${o.type}: ${o.llmContext}`;
+        return `<result>${o.llmContext}</result>`;
       })
       .join('\n-------------\n');
 
-    const finalContextWithWidgets = `<search_results note="These are the search results and you can cite these">${finalContext}</search_results>\n<widgets_result noteForAssistant="Its output is already showed to the user, you can use this information to answer the query but do not CITE this as a souce">${widgetContext}</widgets_result>`;
+    const finalContextWithWidgets = `<search_results note="These are the search results and assistant can cite these">\n${finalContext}\n</search_results>\n<widgets_result noteForAssistant="Its output is already showed to the user, assistant can use this information to answer the query but do not CITE this as a souce">\n${widgetContext}\n</widgets_result>`;
 
     const writerPrompt = getWriterPrompt(finalContextWithWidgets);
     const answerStream = input.config.llm.streamText({