agents/iterative_agent.py

"""
Iterative Agent - Multi-turn research agent with planning and refinement.

Strategy:
1. Generate a research plan with sub-topics
2. For each sub-topic, generate targeted queries and search
3. Synthesize findings from all sub-topics
4. Review synthesis and identify knowledge gaps
5. Run additional targeted searches to fill gaps
6. Produce the final report
"""

from __future__ import annotations

import json
from typing import Any, Optional

from benchmark.interface import BaseResearchAgent, ResearchOutput
from benchmark.websearch import BenchmarkWebSearchClient


class IterativeResearchAgent(BaseResearchAgent):
    """Multi-turn agent with planning, iterative search, and gap-filling."""

    @property
    def name(self) -> str:
        return "iterative-research"

    @property
    def description(self) -> str:
        return (
            "Multi-turn research agent: decomposes the topic into sub-questions, "
            "searches iteratively, identifies knowledge gaps, and fills them with "
            "additional targeted searches before final synthesis."
        )

    @property
    def author(self) -> str:
        return "DR-Bench Team"

    async def research(
        self,
        topic: str,
        llm: Any,
        websearch: BenchmarkWebSearchClient,
        *,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        max_iterations: int = 3,
        **kwargs,
    ) -> ResearchOutput:
        searches_made = []
        all_search_content: list[str] = []

        # Step 1: Research Planning
        plan_response = await llm.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a research planning assistant. Given a research topic, "
                        "break it down into 3-4 specific sub-questions that need to be "
                        "answered. For each, suggest 2 search queries.\n\n"
                        "Return JSON: [{\"question\": \"...\", \"queries\": [\"q1\", \"q2\"]}]"
                    ),
                },
                {"role": "user", "content": topic},
            ],
            temperature=0.3,
        )

        plan_text = plan_response.choices[0].message.content or "[]"
        try:
            plan = json.loads(plan_text)
            if not isinstance(plan, list):
                plan = [{"question": topic, "queries": [topic]}]
        except json.JSONDecodeError:
            plan = [{"question": topic, "queries": [topic]}]

        # Step 2: Iterative Search per sub-question
        for i, item in enumerate(plan[:4]):
            question = item.get("question", "")
            queries = item.get("queries", [question])[:3]

            results = await websearch.search(
                queries=queries, start_date=start_date, end_date=end_date,
            )
            searches_made.append({
                "queries": queries, "urls": results.get_all_urls(),
                "num_results": results.total_results,
                "sub_question": question, "iteration": 1,
            })

            content = results.get_all_content()
            if content:
                all_search_content.append(f"\n## Sub-question {i+1}: {question}\n{content}")

        # Step 3: Initial Synthesis
        combined = "\n".join(all_search_content)
        if len(combined) > 60000:
            combined = combined[:60000] + "\n... [truncated]"

        synth_resp = await llm.chat.completions.create(
            model=self.model_name,
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a business development research assistant. "
                        "Synthesize search results into a comprehensive response. "
                        "Be specific and cite real facts.\n\n"
                        "At the end, add '## KNOWLEDGE GAPS' listing 1-3 missing "
                        "pieces of information as search queries."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Research Request:\n{topic}\n\nFindings:\n{combined}",
                },
            ],
            temperature=0.4,
        )
        synthesis = synth_resp.choices[0].message.content or ""

        # Step 4: Gap-filling search
        if "KNOWLEDGE GAPS" in synthesis and max_iterations > 1:
            gap_section = synthesis.split("KNOWLEDGE GAPS")[-1]
            gap_queries = [
                line.strip().lstrip("- ").lstrip("0123456789. ")
                for line in gap_section.split("\n")
                if line.strip() and not line.strip().startswith("#") and len(line.strip()) > 10
            ][:3]

            if gap_queries:
                gap_results = await websearch.search(
                    queries=gap_queries, start_date=start_date, end_date=end_date,
                )
                searches_made.append({
                    "queries": gap_queries, "urls": gap_results.get_all_urls(),
                    "num_results": gap_results.total_results,
                    "iteration": 2, "type": "gap_filling",
                })

                gap_content = gap_results.get_all_content()
                if gap_content and len(gap_content) > 100:
                    final_resp = await llm.chat.completions.create(
                        model=self.model_name,
                        messages=[
                            {
                                "role": "system",
                                "content": (
                                    "Produce the final complete response incorporating "
                                    "all information. No knowledge gaps section."
                                ),
                            },
                            {
                                "role": "user",
                                "content": (
                                    f"Research Request:\n{topic}\n\n"
                                    f"Draft:\n{synthesis}\n\nAdditional:\n{gap_content}"
                                ),
                            },
                        ],
                        temperature=0.4,
                    )
                    synthesis = final_resp.choices[0].message.content or synthesis

        if "## KNOWLEDGE GAPS" in synthesis:
            synthesis = synthesis.split("## KNOWLEDGE GAPS")[0].strip()

        return ResearchOutput(
            report=synthesis,
            searches_made=searches_made,
            metadata={"plan": plan, "iterations": len(searches_made)},
        )