feat: migrate from BFS to DFS for working with research_plan

.gitignore

@@ -9,7 +9,7 @@ backend/*.pyo
 backend/.venv/
 backend/.env*
 backend/downloads/*
-backend/output.json
+backend/*.log.*
 backend/.ruff_cache/
 
 # Next.js ignore files

backend/app.py

@@ -86,7 +86,6 @@ async def start_research(sid, data):
         data = json.loads(data) if type(data) is not dict else data
         topic = data.get("topic").strip()
         max_depth: int = data.get("max_depth")
-        max_breadth: int = data.get("max_breadth")
         num_sites_per_query: int = data.get("num_sites_per_query")
 
         knet, _ = await session_manager.get_or_create_session(sid)
@@ -94,14 +93,10 @@ async def start_research(sid, data):
         session_id = sid
         logger.info(f"Starting research for client {session_id}.\nTopic '{topic}'")
 
-        async def progress_callback(status):
-            await sio.emit(
-                "status",
-                {"message": status["message"], "progress": status["progress"]},
-                room=session_id,
-            )
+        async def progress_callback(status: dict):
+            await sio.emit("status", status, room=session_id)
 
-        research_results = await knet.conduct_research(topic, progress_callback, max_depth, max_breadth, num_sites_per_query)
+        research_results = await knet.conduct_research(topic, progress_callback, max_depth, num_sites_per_query)
         logger.info(f"Research completed for topic: {topic}")
         await sio.emit("research_complete", research_results, room=session_id)
 

backend/knet.py

@@ -19,27 +19,50 @@ load_dotenv()
 
 class Prompt:
     def __init__(self) -> None:
-        self.continue_branch = dedent("""Given the current research context and findings, should we explore this branch deeper?
+        self.research_plan = dedent("""You are an expert AI Deep Research agent, part of a Multiagent system.
+
+        User query:
+        "{topic}".
+
+        ---
+        Generate few very high level steps on which other agents can do info collection runs. Provide only data collection steps, no data identification, summarization, manipulation, selection, etc.
+        Return a string array of steps.""")
+
+        self.site_summary = dedent("""Extract specific verbatim key information from the following content that is related to the topic "{query}". No small talk.
+        Findings:
+        {findings}""")
+
+        self.continue_branch = dedent("""Given the current state of research, decide whether to continue exploring the current branch or not.
+        Global Research Plan:
+        {research_plan}
 
         Current Topic: {query}
-        Current Depth: {depth}
-        Path from Root: {path}
-        Key Findings:
-        {findings}
+        Searched Queries:
+        {past_queries}
+
+        Findings under current topic:
+        {ctx_manager}
 
         Consider:
-        1. Relevance to main topic
-        2. Potential for new insights
-        3. Depth vs breadth tradeoff
-        4. Information saturation
+        - Information saturation
+        - Information duplication
+        - Coverage of current topic
+        - Potential for new insights
 
-        Return only: decision: true/false""")
+        Return only decision: true/false""")
 
-        self.search_query = dedent("""Based on the following findings about "{topic}", suggest new research directions.
-        Findings:
+        self.search_query = dedent("""Based on the following findings on topic {vertical}, suggest new research directions.
+        Global Research Plan:
+        {research_plan}
+
+        Searched queries:
+        {past_queries}
+
+        Findings under current topic:
         {ctx_manager}
 
-        Suggest up to {n} specific google search queries that would help data which:
+        Suggest up to {n} specific google search queries that:
+        - Covers what has not been covered yet
         - Builds upon these findings
         - Explores different aspects
         - Goes deeper into important details
@@ -47,34 +70,67 @@ class Prompt:
         Return as JSON array of objects with properties:
         - query (string)""")
 
+        self.report_outline = dedent("""Generate a comprehensive report outline on the user query based on the following research findings:
+        User query:
+        {topic}
+
+        Findings:
+        {ctx_manager}
+
+        The outline should include:
+        - Title
+        - List of h2 headings""")
+
+        self.report_fillin = dedent("""Fill in the content for the following report outline on the user query based on the following research findings:
+        User query:
+        {topic}
+
+        Findings:
+        {ctx_manager}
+
+        Report Outline:
+        {report_outline}
+
+        Current slot to fill in: (h2 heading)
+        {slot}
+        """)
+
 
 class Schema:
     def __init__(self) -> None:
+        self.research_plan = genai.types.Schema(
+            type=genai.types.Type.OBJECT,
+            required=["steps"],
+            properties={"steps": genai.types.Schema(type=genai.types.Type.ARRAY, items=genai.types.Schema(type=genai.types.Type.STRING))},
+        )
+
         self.continue_branch = genai.types.Schema(
             type=genai.types.Type.OBJECT,
             required=["decision"],
-            properties={
-                "decision": genai.types.Schema(type=genai.types.Type.BOOLEAN),
-            },
+            properties={"decision": genai.types.Schema(type=genai.types.Type.BOOLEAN)},
         )
 
         self.search_query = genai.types.Schema(
             type=genai.types.Type.OBJECT,
             required=["branches"],
+            properties={"branches": genai.types.Schema(type=genai.types.Type.ARRAY, items=genai.types.Schema(type=genai.types.Type.STRING))},
+        )
+
+        self.report_outline = genai.types.Schema(
+            type=genai.types.Type.OBJECT,
+            required=["title", "headings"],
             properties={
-                "branches": genai.types.Schema(
-                    type=genai.types.Type.ARRAY,
-                    items=genai.types.Schema(
-                        type=genai.types.Type.OBJECT,
-                        required=["query"],
-                        properties={
-                            "query": genai.types.Schema(type=genai.types.Type.STRING),
-                        },
-                    ),
-                )
+                "title": genai.types.Schema(type=genai.types.Type.STRING),
+                "headings": genai.types.Schema(type=genai.types.Type.ARRAY, items=genai.types.Schema(type=genai.types.Type.STRING)),
             },
         )
 
+        self.report_fillin = genai.types.Schema(
+            type=genai.types.Type.OBJECT,
+            required=["content"],
+            properties={"content": genai.types.Schema(type=genai.types.Type.STRING)},
+        )
+
 
 class ResearchProgress:
     def __init__(self, callback):
@@ -82,90 +138,104 @@ class ResearchProgress:
         self.callback = callback
 
     async def update(self, progress: int, message: str):
-        self.progress += progress
-        if self.progress > 100:
-            self.progress = 100
-        if self.callback:
-            await self.callback({"progress": self.progress, "message": message})
+        self.progress = min(100, self.progress + progress)  # max 100
+        await self.callback({"progress": self.progress, "message": message})
+
+    async def setter(self, progress: int, message: str):
+        self.progress = min(100, progress)  # max 100
+        await self.callback({"progress": self.progress, "message": message})
 
 
 class KNet:
-    def __init__(self, scraper_instance: CrawlForAIScraper, max_depth: int = 1, max_breadth: int = 1, num_sites_per_query: int = 5):
+    def __init__(self, scraper_instance: CrawlForAIScraper, max_depth: int = 1, num_sites_per_query: int = 5):
         self.api_key = os.getenv("GOOGLE_API_KEY")
         assert self.api_key, "Google API key is required"
         self.scraper = scraper_instance
         self.logger = logging.getLogger(__name__)
         self.prompt = Prompt()
         self.schema = Schema()
+        self.progress = None
 
         # Init Google GenAI client
         self.genai_client = genai.Client(api_key=self.api_key)
 
         # Parameters
         self.max_depth = max_depth
-        self.max_breadth = max_breadth
         self.num_sites_per_query = num_sites_per_query
 
         # Global State
+        self.research_plan: list[str] = []
+        self.idx_research_plan: int = 0
         self.ctx_researcher: list[str] = []
         self.ctx_manager: list[str] = []
         self.token_count: int = 0
 
-    async def conduct_research(self, topic: str, progress_callback, max_depth: int, max_breadth: int, num_sites_per_query: int) -> dict:
+    async def conduct_research(self, topic: str, progress_callback, max_depth: int, num_sites_per_query: int) -> dict:
         # Local Runtime State
-        progress = ResearchProgress(progress_callback)
+        self.progress = ResearchProgress(progress_callback)
         self.max_depth = max_depth
-        self.max_breadth = max_breadth
         self.num_sites_per_query = num_sites_per_query
 
         # Reset global state
+        self.research_plan = []
+        self.idx_research_plan = 0
         self.ctx_researcher = []
         self.ctx_manager = []
         self.token_count = 0
 
         try:
+            # Generate research plan
+            await self.progress.update(0, "Generating research plan...")
+            self.research_plan = self.generate_content(self.prompt.research_plan.format(topic=topic), schema=self.schema.research_plan)["steps"]
+            self.logger.info(f"Research plan:\n{json.dumps(self.research_plan, indent=2)}")
+
             # Generate initial search query
             query = self.generate_content(
-                self.prompt.search_query.format(topic=topic, ctx_manager=json.dumps(self.ctx_manager, indent=2), n=1),
-                schema=self.schema.search_query,
-            )
-            root_node = ResearchNode(query.get("branches")[0]["query"])
+                self.prompt.search_query.format(vertical=self.research_plan[self.idx_research_plan]), schema=self.schema.search_query
+            )["branches"][0]
+
+            # Initialize research tree
+            root_node = ResearchNode(query)
             to_explore = deque([(root_node, 0)])  # (node, depth) pairs
             explored_queries = set()  # {string, string, ...}
 
-            await progress.update(5, "Starting research...")
+            await self.progress.update(0, "Starting research...")
 
-            while to_explore:
+            # Iterate on research plan
+            for self.idx_research_plan, _ in enumerate(self.research_plan):
                 current_node, current_depth = to_explore.popleft()
-
-                if current_node.query in explored_queries or current_depth > self.max_depth:
-                    continue
-
-                self.logger.info(f"Exploring: {current_node.query} (Depth: {current_depth})")
-                await progress.update(5, f"Exploring: {current_node.query}")
-
-                # Search and scrape
-                current_node.data = await self.scraper.search_and_scrape(
-                    current_node.query, self.num_sites_per_query
-                )  # node -> data = [{url:...}, {url:...}, ...]
-                self.ctx_researcher.append(json.dumps(current_node.data, indent=2))
-                explored_queries.add(current_node.query)
-
-                # Only branch if we have data and haven't reached max depth
-                if self._should_continue_branch(current_node, topic):
-                    if current_node.data and current_depth < self.max_depth:
-                        new_branches = self._gen_queries(current_node, topic)
-                        for branch in new_branches:
-                            to_explore.append((branch, current_depth + 1))
+                await self.progress.update(100 / (len(self.research_plan) + 1), f"{self.research_plan[self.idx_research_plan]}")
+
+                while to_explore:
+                    current_node, current_depth = to_explore.popleft()
+                    if current_depth > self.max_depth:
+                        continue
+
+                    self.logger.info(f"Exploring: {current_node.query} (depth: {current_depth})")
+                    await self.progress.update(0, f"s_{current_node.query}")
+
+                    # Search and scrape
+                    current_node.data = await self.scraper.search_and_scrape(
+                        current_node.query, self.num_sites_per_query
+                    )  # node -> data = [{url:...}, {url:...}, ...]
+                    self.ctx_researcher.append(json.dumps(current_node.data, indent=2))
+                    explored_queries.add(current_node.query)
+
+                    # Only branch if we have data and haven't reached max depth
+                    if self._should_continue_branch(current_node, topic):
+                        if current_node.data and current_depth < self.max_depth:
+                            new_branches = self._gen_queries(current_node, topic)
+                            for branch in new_branches:
+                                to_explore.appendleft((branch, current_depth + 1))
 
             # Generate final report
-            await progress.update(30, "Generating comprehensive report...")
-            final_report = self._generate_final_report(root_node)
+            await self.progress.update(100 / (len(self.research_plan) + 1), "Generating final report...")
+            final_report = self._generate_final_report(root_node, topic)
 
             self.logger.info(f"Research completed. Explored {len(explored_queries)} queries across {root_node.max_depth()} levels")
-            await progress.update(100, "Research complete!")
+            await self.progress.update(100, "Research complete!")
 
-            with open("output.json", "a", encoding="utf-8") as f:
+            with open("output.log.json", "w", encoding="utf-8") as f:
                 json.dump(final_report, f, indent=2)
             return final_report
 
@@ -173,15 +243,30 @@ class KNet:
             self.logger.error("Research failed", exc_info=True)
             raise
 
-    def _generate_final_report(self, root_node: ResearchNode, retry_count: int = 1) -> Dict[str, Any]:
+    def _generate_final_report(self, root_node: ResearchNode, topic: str, retry_count: int = 1) -> Dict[str, Any]:
         try:
-            findings = "\n".join(self.ctx_manager)
-            with open("output.json", "w", encoding="utf-8") as f:
+            self.progress.setter(0, "Generating report...")
+            findings = "\n\n------\n\n".join(self.ctx_manager)
+            with open("ctx_manager.log.txt", "w", encoding="utf-8") as f:
                 f.write(findings)
-            prompt = f"""Generate a comprehensive report on the topic "{root_node.query}" based on the following research findings:
-            {findings}
-            """
-            response = self.generate_content(prompt)
+
+            # Generate report outline
+            outline = self.generate_content(self.prompt.report_outline.format(topic=topic, ctx_manager=findings), schema=self.schema.report_outline)
+            self.logger.info(f"Report outline:\n{json.dumps(outline, indent=2)}")
+            report = []
+            # Fill in report outline
+            for i, heading in enumerate(outline["headings"]):
+                self.progress.update(100 / (len(outline["headings"] + 1)), "Generating report...")
+                content = self.generate_content(
+                    self.prompt.report_fillin.format(
+                        topic=topic,
+                        ctx_manager=findings,
+                        report_outline=["[done] " + outline["title"]] + [f"[done] {h}" for _, h in enumerate(outline["headings"]) if i < _],
+                        slot=heading,
+                    ),
+                    schema=self.schema.report_fillin,
+                )["content"]
+                report.append({"heading": heading, "content": content})
 
             # Collate multimedia content
             media_content = {"images": [], "videos": [], "links": [], "references": []}
@@ -203,7 +288,6 @@ class KNet:
             def build_tree_structure(node: ResearchNode) -> Dict:
                 if not node:
                     return {}
-
                 sources = [d["url"] for d in node.data if d.get("url")]
                 return {
                     "query": node.query,
@@ -215,7 +299,7 @@ class KNet:
             return {
                 "topic": root_node.query,
                 "timestamp": datetime.now().isoformat(),
-                "content": response,
+                "content": report,
                 "media": media_content,
                 "research_tree": build_tree_structure(root_node),
                 "metadata": {
@@ -229,7 +313,7 @@ class KNet:
         except Exception as e:
             if e in ["GEMINI_RECITATION", "NO_RESPONSE"] and retry_count < 3:
                 self.logger.error(f"Retrying final report:C:{retry_count / 3}", exc_info=True)
-                self._generate_final_report(root_node, retry_count + 1)
+                return self._generate_final_report(root_node, retry_count + 1)
             self.logger.error("Error generating final report", exc_info=True)
             raise
 
@@ -239,11 +323,13 @@ class KNet:
                 return []
 
             prompt = self.prompt.search_query.format(
-                topic=topic,
-                ctx_manager=json.dumps(self.ctx_manager, indent=2),
-                n=self.max_breadth,
+                vertical=self.research_plan[self.idx_research_plan],
+                research_plan="\n".join([f"[done] {step}" for i, step in enumerate(self.research_plan) if i < self.idx_research_plan]),
+                past_queries="\n".join([f"[done] {query}" for query in node.get_path_to_root()[1:]]),
+                ctx_manager="\n\n---\n\n".join(self.ctx_manager),
+                n=1,
             )
-            response = self.generate_content(prompt, schema=self.schema.search_query)
+            response = self.generate_content(prompt, schema=self.schema.search_query, temp=1.5)
             self.logger.info(f"Spawn branches '{node.query}':\n{json.dumps(response['branches'], indent=2)}")
 
             # Add children to current node
@@ -252,7 +338,7 @@ class KNet:
             #       |-> child
             new_nodes = []
             for branch in response.get("branches", []):
-                child_node = node.add_child(branch["query"])
+                child_node = node.add_child(branch)
                 new_nodes.append(child_node)
 
             self.logger.info(f"Spawned {len(new_nodes)} new branch(es)")
@@ -261,7 +347,7 @@ class KNet:
         except Exception as e:
             if e in ["GEMINI_RECITATION", "NO_RESPONSE"] and retry_count < 3:
                 self.logger.error(f"Retrying _gen_queries | C:{retry_count / 3}", exc_info=True)
-                self._gen_queries(node, topic, retry_count + 1)
+                return self._gen_queries(node, topic, retry_count + 1)
             self.logger.error("_gen_queries failed", exc_info=True)
             raise
 
@@ -273,15 +359,15 @@ class KNet:
             # Generate summary of key findings into the manager's context
             if node.data:
                 findings = ("\n" + "-" * 10 + "Next data" + "-" * 10 + "\n").join([json.dumps(d, indent=2) for d in node.data])
-                response = self.generate_content(f"Extract key findings from the following data related to the topic '{topic}':\n{findings}")
+                response = self.generate_content(self.prompt.site_summary.format(query=node.query, findings=findings), temp=0.2)
                 self.ctx_manager.append(response)
 
             # Research manager takes decision to proceed or not
             prompt = self.prompt.continue_branch.format(
+                research_plan="\n".join([f"[done] {step}" for i, step in enumerate(self.research_plan) if i < self.idx_research_plan]),
                 query=node.query,
-                depth=node.depth,
-                path=" -> ".join(node.get_path_to_root()),
-                findings="\n".join(self.ctx_manager),
+                past_queries="\n".join([f"[done] {query}" for query in node.get_path_to_root()[1:]]),
+                ctx_manager="\n\n---\n\n".join(self.ctx_manager),
             )
             response = self.generate_content(prompt, schema=self.schema.continue_branch)
             self.logger.info(f"Branch decision '{node.query}': {response['decision']}")
@@ -291,11 +377,11 @@ class KNet:
         except Exception as e:
             if e in ["GEMINI_RECITATION", "NO_RESPONSE"] and retry_count < 3:
                 self.logger.error(f"Retrying branch decision:C:{retry_count / 3}", exc_info=True)
-                self._should_continue_branch(node, topic, retry_count + 1)
+                return self._should_continue_branch(node, topic, retry_count + 1)
             self.logger.error("Branch decision failed:", exc_info=True)
             raise
 
-    def generate_content(self, prompt: str, schema: Dict[str, Any] = {}) -> Dict[str, Any] | str:
+    def generate_content(self, prompt: str, schema: Dict[str, Any] = {}, temp: float = 0.9) -> Dict[str, Any] | str:
         safe = [
             types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=types.HarmBlockThreshold.BLOCK_NONE),
             types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=types.HarmBlockThreshold.BLOCK_NONE),
@@ -305,10 +391,10 @@ class KNet:
         ]
         if schema:
             generate_content_config = types.GenerateContentConfig(
-                temperature=0.9, response_mime_type="application/json", safety_settings=safe, response_schema=schema
+                temperature=temp, response_mime_type="application/json", safety_settings=safe, response_schema=schema
             )
         else:
-            generate_content_config = types.GenerateContentConfig(temperature=0.9, response_mime_type="text/plain", safety_settings=safe)
+            generate_content_config = types.GenerateContentConfig(temperature=temp, response_mime_type="text/plain", safety_settings=safe)
 
         try:
             response = self.genai_client.models.generate_content(model="gemini-2.0-flash", contents=prompt, config=generate_content_config)

backend/research_node.py

@@ -3,9 +3,7 @@ from typing import Any, Dict, List, Optional
 
 
 class ResearchNode:
-    def __init__(
-        self, query: str, parent: Optional["ResearchNode"] = None, depth: int = 0
-    ):
+    def __init__(self, query: str, parent: Optional["ResearchNode"] = None, depth: int = 0):
         self.query = query
         self.parent = parent
         self.depth = depth
@@ -18,6 +16,10 @@ class ResearchNode:
         return child
 
     def get_path_to_root(self) -> List[str]:
+        """
+        Returns the path from this node to the root node.
+        List[str]: [root.query, ..., self.query]
+        """
         path = [self.query]
         current = self
         while current.parent:
@@ -33,9 +35,7 @@ class ResearchNode:
     def total_children(self) -> int:
         if not self.children:
             return 0
-        return len(self.children) + sum(
-            [child.total_children() for child in self.children]
-        )
+        return len(self.children) + sum([child.total_children() for child in self.children])
 
     def get_all_data(self) -> List[Dict[str, Any]]:
         data = copy.deepcopy(self.data)