backend/research.py

"""
Research agent backend using DR-TULU model - model-driven deep research

DR-TULU drives the research loop - it decides when to search, what to search for,
and when it has enough information to answer.
"""
import json
import logging
import os
import re
import uuid
from typing import List, Dict, Optional, Tuple
import requests
import trafilatura

logger = logging.getLogger(__name__)


def search_web(query: str, api_key: str, num_results: int = 10) -> List[Dict[str, str]]:
    """Search the web using Serper API"""
    url = "https://google.serper.dev/search"
    payload = json.dumps({"q": query, "num": num_results})
    headers = {
        'X-API-KEY': api_key,
        'Content-Type': 'application/json'
    }

    try:
        response = requests.post(url, headers=headers, data=payload, timeout=10)

        if response.status_code != 200:
            return []

        data = response.json()
        results = []
        for item in data.get('organic', []):
            results.append({
                'title': item.get('title', ''),
                'url': item.get('link', ''),
                'snippet': item.get('snippet', '')
            })
        return results
    except Exception as e:
        logger.error(f"Search error: {e}")
        return []


def extract_content(url: str) -> Optional[str]:
    """Extract main content from a URL"""
    try:
        downloaded = trafilatura.fetch_url(url)
        if downloaded is None:
            return None
        text = trafilatura.extract(
            downloaded,
            include_comments=False,
            include_tables=False,
            no_fallback=False
        )
        return text
    except Exception as e:
        logger.error(f"Content extraction error for {url}: {e}")
        return None


def generate_snippet_id() -> str:
    """Generate unique snippet ID"""
    return f"S_{uuid.uuid4().hex[:8]}"


def generate_webpage_id() -> str:
    """Generate unique webpage ID"""
    return f"W_{uuid.uuid4().hex[:8]}"


def parse_tool_calls(text: str) -> List[Dict]:
    """
    Parse <call_tool name="...">query</call_tool> from model output.
    Returns list of {"name": str, "query": str, "params": dict}
    """
    pattern = r'<call_tool\s+name="([^"]+)"([^>]*)>([^<]+)</call_tool>'
    matches = re.findall(pattern, text)

    tool_calls = []
    for name, params_str, query in matches:
        # Parse optional params like limit="8" year="2021-2025"
        params = {}
        param_pattern = r'(\w+)="([^"]+)"'
        for param_name, param_value in re.findall(param_pattern, params_str):
            params[param_name] = param_value

        tool_calls.append({
            "name": name.strip(),
            "query": query.strip(),
            "params": params
        })

    return tool_calls


def parse_think_blocks(text: str) -> List[str]:
    """Extract <think>...</think> content"""
    pattern = r'<think>(.*?)</think>'
    return re.findall(pattern, text, re.DOTALL)


def parse_answer(text: str) -> Optional[str]:
    """Extract <answer>...</answer> content"""
    pattern = r'<answer>(.*?)</answer>'
    match = re.search(pattern, text, re.DOTALL)
    return match.group(1).strip() if match else None


def format_search_results(results: List[Dict], query: str) -> str:
    """
    Format search results as DR-TULU tool output.
    """
    if not results:
        return "<tool_output>No results found.</tool_output>"

    snippets = []
    for r in results:
        snippet_id = generate_snippet_id()
        # Escape XML special chars in content
        title = r.get("title", "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        snippet_text = r.get("snippet", "").replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        url = r.get("url", "")

        snippets.append(
            f'<snippet id="{snippet_id}" url="{url}" title="{title}">\n'
            f'{snippet_text}\n'
            f'</snippet>'
        )

    return f"<tool_output>\n" + "\n".join(snippets) + "\n</tool_output>"


def format_webpage_content(url: str, title: str, content: str) -> str:
    """
    Format extracted webpage as DR-TULU tool output.
    """
    if not content:
        return f"<tool_output>Could not extract content from {url}</tool_output>"

    webpage_id = generate_webpage_id()
    # Truncate very long content
    if len(content) > 8000:
        content = content[:8000] + "\n[Content truncated...]"

    # Escape XML special chars
    content = content.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    title = title.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

    return (
        f"<tool_output>\n"
        f'<webpage id="{webpage_id}" url="{url}" title="{title}">\n'
        f'{content}\n'
        f'</webpage>\n'
        f"</tool_output>"
    )


def execute_tool(
    tool_name: str,
    query: str,
    params: dict,
    serper_key: str
) -> Tuple[str, List[Dict]]:
    """
    Execute a tool and return (formatted_output, raw_results).
    """
    if tool_name == "google_search":
        num_results = int(params.get("limit", 10))
        results = search_web(query, serper_key, num_results=num_results)
        formatted = format_search_results(results, query)
        return formatted, results

    elif tool_name == "browse_webpage":
        # query is the URL for browse_webpage
        url = query
        content = extract_content(url)
        title = url  # Could extract from content if needed
        formatted = format_webpage_content(url, title, content or "")
        return formatted, [{"url": url, "content": content, "title": title}]

    else:
        return f"<tool_output>Unknown tool: {tool_name}</tool_output>", []


def get_dr_tulu_system_prompt() -> str:
    """Return the DR-TULU system prompt"""
    return '''You are a research assistant who answers questions through iterative reasoning and research.

## Process
- Use <think></think> tags to show your reasoning at any point.
- Use <call_tool name="...">query</call_tool> when you need information (see tools below).
- You can alternate between thinking and searching multiple times.
- Only provide <answer></answer> tags when you have enough information for a complete response.
- Support every non-trivial claim with retrieved evidence. Wrap the exact claim span in <cite id="ID1,ID2">...</cite>, where id are snippet IDs from searched results.

## Calling Tools (<call_tool name="...">query</call_tool>)

1. google_search
- Purpose: general web search.
- Input via: <call_tool name="google_search">your query</call_tool>
- Output: web search snippets.

2. browse_webpage
- Purpose: open a specific URL and extract readable page text.
- Input via: <call_tool name="browse_webpage">https://example.com/article</call_tool>
- Output: webpage content.

## Tool Output
- After you issue a tool call, we will execute it and return results wrapped in <tool_output> tags.
- For web search: <tool_output><snippet id=UNIQUE_ID url="..." title="...">content</snippet>...</tool_output>
- For web browsing: <tool_output><webpage id=UNIQUE_ID url="..." title="...">content</webpage></tool_output>

## Answer and Citation Format
- Once you collect all necessary information, generate the final answer with <answer></answer> tags.
- In your answer, wrap supported text in <cite id="SNIPPET_ID">...</cite> using exact IDs from returned snippets.
- Write comprehensive, well-structured answers with clear sections when appropriate.
'''


def stream_research(
    client,
    model: str,
    question: str,
    serper_key: str,
    max_iterations: int = 5,
    max_websites: int = 50,
    system_prompt: str = "",
    sub_agent_model: Optional[str] = None,
    parallel_workers: int = 8,
    max_tool_calls: int = 20,
    abort_event=None,
    **kwargs
):
    """
    Stream deep research results using DR-TULU.

    The model drives the research loop - it decides when to search,
    what to search for, and when it has enough information to answer.

    Yields same event types as the original research.py for API compatibility.
    """

    # Build system prompt
    dr_tulu_system = get_dr_tulu_system_prompt()
    if system_prompt:
        dr_tulu_system += f"\n\n{system_prompt}"

    messages = [
        {"role": "system", "content": dr_tulu_system},
        {"role": "user", "content": question}
    ]

    yield {"type": "status", "message": f"Starting DR-TULU research: {question}"}

    tool_call_count = 0
    findings = []  # Track sources for compatibility
    all_queries = []  # Track queries for compatibility
    iteration = 0
    max_iterations_without_progress = 3
    iterations_without_tool_calls = 0

    while tool_call_count < max_tool_calls:
        # Check abort before each iteration
        if abort_event and abort_event.is_set():
            yield {"type": "aborted"}
            yield {"type": "done"}
            return

        iteration += 1

        # Call DR-TULU
        yield {"type": "status", "message": "Thinking..."}

        try:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                max_tokens=4096,
                temperature=0.7,
            )
            assistant_message = response.choices[0].message.content
        except Exception as e:
            yield {"type": "error", "content": f"Model error: {str(e)}"}
            yield {"type": "done"}
            return

        # Parse thinking blocks and yield as status
        think_blocks = parse_think_blocks(assistant_message)
        for thought in think_blocks:
            # Truncate long thoughts for status display
            thought_preview = thought[:300].strip()
            if len(thought) > 300:
                thought_preview += "..."
            yield {"type": "status", "message": f"Reasoning: {thought_preview}"}

        # Check for final answer
        answer = parse_answer(assistant_message)
        if answer:
            yield {"type": "status", "message": "Research complete! Generating result..."}

            # Wrap in <result> tags for compatibility with command center
            result_content = answer
            if "<result>" not in answer:
                result_content = f"<result>\n{answer}\n</result>"

            yield {"type": "result_preview", "content": answer, "figures": {}}
            yield {"type": "result", "content": answer, "figures": {}}
            yield {"type": "done"}
            return

        # Parse and execute tool calls
        tool_calls = parse_tool_calls(assistant_message)

        if not tool_calls:
            # No tool calls and no answer - model might be stuck
            iterations_without_tool_calls += 1

            if iterations_without_tool_calls >= max_iterations_without_progress:
                # Force final answer
                yield {"type": "status", "message": "No more searches needed, generating final answer..."}
                messages.append({"role": "assistant", "content": assistant_message})
                messages.append({"role": "user", "content": "Please provide your final answer now using <answer></answer> tags."})
                continue

            # Append message and continue to prompt for more
            messages.append({"role": "assistant", "content": assistant_message})
            messages.append({"role": "user", "content": "Please continue your research or provide your answer using <answer></answer> tags."})
            continue

        # Reset counter since we have tool calls
        iterations_without_tool_calls = 0

        # Track queries for compatibility - build map of query -> global index
        new_queries = [tc["query"] for tc in tool_calls if tc["name"] == "google_search"]
        query_start_idx = len(all_queries)  # Starting global index for this batch

        if new_queries:
            all_queries.extend(new_queries)
            yield {
                "type": "queries",
                "queries": new_queries,
                "iteration": iteration
            }

        # Track stats per query for this batch
        query_stats = {}  # local index -> {relevant, irrelevant, error}

        # Execute tools and collect results
        tool_outputs = []
        search_idx = 0  # Track which search query we're on within this batch

        for i, tc in enumerate(tool_calls):
            # Check abort between tool executions
            if abort_event and abort_event.is_set():
                yield {"type": "aborted"}
                yield {"type": "done"}
                return

            tool_call_count += 1

            if tc["name"] == "google_search":
                # Calculate global query index
                global_query_idx = query_start_idx + search_idx

                yield {
                    "type": "status",
                    "message": f"Searching: {tc['query'][:50]}..."
                }

                # Initialize stats for this query
                if global_query_idx not in query_stats:
                    query_stats[global_query_idx] = {"relevant": 0, "irrelevant": 0, "error": 0}
            else:
                global_query_idx = None  # browse_webpage doesn't have a query index
                yield {
                    "type": "status",
                    "message": f"Browsing: {tc['query'][:50]}..."
                }

            formatted_output, raw_results = execute_tool(
                tc["name"],
                tc["query"],
                tc["params"],
                serper_key
            )
            tool_outputs.append(formatted_output)

            # Yield source events for compatibility
            if tc["name"] == "google_search":
                for j, result in enumerate(raw_results):
                    findings.append({
                        "source": result.get("url", ""),
                        "title": result.get("title", ""),
                        "analysis": result.get("snippet", "")
                    })

                    # All search results are considered relevant (DR-TULU decides what to use)
                    query_stats[global_query_idx]["relevant"] += 1

                    yield {
                        "type": "source",
                        "query_index": global_query_idx,
                        "query_text": tc["query"],
                        "title": result.get("title", ""),
                        "url": result.get("url", ""),
                        "analysis": result.get("snippet", ""),
                        "finding_count": len(findings),
                        "is_relevant": True,  # DR-TULU decides relevance
                        "is_error": False,
                        "error_message": ""
                    }

                # Emit query_stats after processing all results for this query
                yield {
                    "type": "query_stats",
                    "query_index": global_query_idx,
                    "relevant_count": query_stats[global_query_idx]["relevant"],
                    "irrelevant_count": query_stats[global_query_idx]["irrelevant"],
                    "error_count": query_stats[global_query_idx]["error"]
                }

                search_idx += 1  # Move to next search query

            elif tc["name"] == "browse_webpage":
                # browse_webpage results don't belong to a search query
                # We can associate them with the last search query or skip
                for result in raw_results:
                    content = result.get("content", "")
                    is_error = not content
                    findings.append({
                        "source": tc["query"],
                        "title": result.get("title", tc["query"]),
                        "analysis": content[:500] if content else "Failed to extract"
                    })

                    # For browse_webpage, use a pseudo-index or skip query association
                    # Since it's browsing a specific URL, we'll emit it without query grouping
                    yield {
                        "type": "source",
                        "query_index": -1,  # Special index for browse results
                        "query_text": f"Browse: {tc['query'][:50]}",
                        "title": result.get("title", tc["query"]),
                        "url": tc["query"],
                        "analysis": content[:500] if content else "Failed to extract content",
                        "finding_count": len(findings),
                        "is_relevant": not is_error,
                        "is_error": is_error,
                        "error_message": "Content extraction failed" if is_error else ""
                    }

            if tool_call_count >= max_tool_calls:
                break

        # Append assistant message and tool results to conversation
        messages.append({"role": "assistant", "content": assistant_message})

        # Combine all tool outputs into one user message
        combined_output = "\n\n".join(tool_outputs)
        messages.append({"role": "user", "content": combined_output})

    # Max tool calls reached - ask for final answer
    yield {"type": "status", "message": "Maximum searches reached, generating final answer..."}

    messages.append({
        "role": "user",
        "content": "You have reached the maximum number of tool calls. Please provide your final answer now using <answer></answer> tags based on the information gathered."
    })

    try:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=4096,
            temperature=0.7,
        )
        final_message = response.choices[0].message.content
        answer = parse_answer(final_message) or final_message

        yield {"type": "result_preview", "content": answer, "figures": {}}
        yield {"type": "result", "content": answer, "figures": {}}
    except Exception as e:
        yield {"type": "error", "content": f"Failed to generate final answer: {str(e)}"}

    yield {"type": "done"}