research_workflow.py

from __future__ import annotations

import asyncio
import json
import logging
import os
from functools import lru_cache
from typing import Any

from dotenv import load_dotenv
from duckduckgo_search import DDGS
from openai import AsyncOpenAI
from pydantic import BaseModel, field_validator

try:
    from tavily import AsyncTavilyClient

    TAVILY_AVAILABLE = True
except ImportError:
    AsyncTavilyClient = None
    TAVILY_AVAILABLE = False

load_dotenv()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

DOMAIN_DEPTH_MAP = {
    "finance": 5,
    "law": 5,
    "healthcare": 5,
    "governance": 4,
    "economics": 4,
    "ai": 4,
    "academia": 4,
    "science": 4,
}
DEFAULT_MAX_SUBTOPICS = 3
DEFAULT_MAX_SEARCH_RESULTS = 4


class Subtopics(BaseModel):
    subtopics: list[str]


class ResearchResults(BaseModel):
    findings: str | list[Any] | dict[str, Any]

    @field_validator("findings", mode="after")
    @classmethod
    def normalize_findings(cls, value: str | list[Any] | dict[str, Any]) -> str:
        return normalize_model_text(value)


class OptimizationDecision(BaseModel):
    justification: str
    needs_more_research: bool


class SummaryReport(BaseModel):
    report: str | list[Any] | dict[str, Any]

    def as_text(self) -> str:
        return normalize_report_text(self.report)


def normalize_model_text(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return "\n".join(format_model_item(item, index) for index, item in enumerate(value, 1))
    if isinstance(value, dict):
        return format_model_item(value)
    return str(value)


def normalize_report_text(value: Any) -> str:
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return "\n\n".join(format_model_item(item) for item in value)
    if isinstance(value, dict):
        sections = []
        for key, section_value in value.items():
            heading = str(key).replace("_", " ").title()
            body = normalize_model_text(section_value)
            sections.append(f"## {heading}\n{body}")
        return "\n\n".join(sections)
    return str(value)


def format_model_item(item: Any, index: int | None = None) -> str:
    prefix = f"{index}. " if index is not None else ""
    if isinstance(item, str):
        return f"{prefix}{item}"
    if not isinstance(item, dict):
        return f"{prefix}{item}"

    fact = item.get("fact") or item.get("finding") or item.get("insight") or item.get("summary")
    source = item.get("source") or item.get("url") or item.get("citation")
    title = item.get("title")

    parts = []
    if fact:
        parts.append(str(fact))
    else:
        parts.append(json.dumps(item, ensure_ascii=False))
    if title:
        parts.append(f"Title: {title}")
    if source:
        parts.append(f"Source: {source}")
    return prefix + " | ".join(parts)


@lru_cache(maxsize=1)
def get_deepseek_client() -> AsyncOpenAI:
    api_key = os.getenv("DEEPSEEK_API_KEY")
    if not api_key:
        raise RuntimeError("DEEPSEEK_API_KEY is not configured.")
    return AsyncOpenAI(api_key=api_key, base_url="https://api.deepseek.com/v1")


@lru_cache(maxsize=1)
def get_tavily_client():
    api_key = os.getenv("TAVILY_API_KEY")
    if not TAVILY_AVAILABLE or not api_key:
        return None
    return AsyncTavilyClient(api_key=api_key)


def get_research_depth(topic: str) -> tuple[int, int]:
    topic_lower = topic.lower()
    for domain, depth in DOMAIN_DEPTH_MAP.items():
        if domain in topic_lower:
            logger.info("Detected domain '%s' -> depth %s", domain, depth)
            return depth, depth
    return DEFAULT_MAX_SUBTOPICS, DEFAULT_MAX_SEARCH_RESULTS


async def call_deepseek_json(
    system_prompt: str,
    user_prompt: str,
    output_model: type[BaseModel],
    temperature: float = 0.3,
) -> BaseModel:
    response = await get_deepseek_client().chat.completions.create(
        model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        response_format={"type": "json_object"},
        temperature=temperature,
    )
    content = response.choices[0].message.content or "{}"
    try:
        data = json.loads(content)
    except json.JSONDecodeError as exc:
        logger.error("Invalid JSON from DeepSeek: %s", content[:300])
        raise ValueError("Model returned invalid JSON.") from exc
    return output_model.model_validate(data)


async def call_deepseek_text(
    messages: list[dict[str, str]],
    temperature: float = 0.4,
    max_tokens: int = 700,
) -> str:
    response = await get_deepseek_client().chat.completions.create(
        model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content or ""


async def run_chat_completion(user_query: str, memory: list[dict[str, str]]) -> str:
    system = (
        "You are the normal chat mode of a deep research assistant. Answer directly, "
        "briefly, and helpfully. Do not claim to have searched the web. If the user "
        "needs current sources or a detailed report, suggest asking for deep research."
    )
    recent_memory = memory[-8:]
    messages = [{"role": "system", "content": system}, *recent_memory, {"role": "user", "content": user_query}]
    try:
        return await call_deepseek_text(messages)
    except Exception as exc:
        logger.error("Chat completion failed: %s", exc)
        return (
            "I can help with quick chat or deep research. The chat model is not available "
            "right now, so please check the DeepSeek API key configuration."
        )


async def web_search_multi(query: str, max_results: int = 4) -> str:
    tasks = [web_search_duckduckgo(query, max_results)]
    if get_tavily_client():
        tasks.append(web_search_tavily(query, max_results))

    results = await asyncio.gather(*tasks, return_exceptions=True)
    combined = []
    for result in results:
        if isinstance(result, str) and result.strip():
            combined.append(result)
        elif isinstance(result, Exception):
            logger.warning("Search engine error: %s", result)
    return "\n\n---\n\n".join(combined) if combined else "No web results found."


async def web_search_duckduckgo(query: str, max_results: int) -> str:
    try:
        return await asyncio.to_thread(_duckduckgo_search_sync, query, max_results)
    except Exception as exc:
        logger.error("DuckDuckGo error: %s", exc)
        return f"[DuckDuckGo] Search failed: {exc}"


def _duckduckgo_search_sync(query: str, max_results: int) -> str:
    with DDGS() as ddgs:
        results = list(ddgs.text(query, max_results=max_results))
    if not results:
        return "[DuckDuckGo] No results found."

    snippets = []
    for index, result in enumerate(results, 1):
        title = result.get("title", "No title")
        body = result.get("body", "No content")
        href = result.get("href", "")
        snippets.append(f"{index}. {title}\n   {body}\n   Source: {href}")
    return "[DuckDuckGo Results]\n" + "\n\n".join(snippets)


async def web_search_tavily(query: str, max_results: int) -> str:
    tavily_client = get_tavily_client()
    if not tavily_client:
        return "[Tavily] Not configured."
    try:
        response = await tavily_client.search(
            query=query,
            max_results=max_results,
            search_depth="basic",
            include_answer=True,
            include_raw_content=False,
        )
    except Exception as exc:
        logger.error("Tavily error: %s", exc)
        return f"[Tavily] Search failed: {exc}"

    if not response.get("results"):
        return "[Tavily] No results found."

    snippets = []
    for index, result in enumerate(response["results"][:max_results], 1):
        title = result.get("title", "No title")
        content = result.get("content", "No content")
        url = result.get("url", "")
        snippets.append(f"{index}. {title}\n   {content}\n   Source: {url}")

    answer = response.get("answer", "")
    result_text = "[Tavily Results]\n" + "\n\n".join(snippets)
    return f"Tavily AI Summary: {answer}\n\n{result_text}" if answer else result_text


async def split_topic(user_query: str, num_subtopics: int) -> Subtopics:
    system = (
        f"Break the user's research request into exactly {num_subtopics} concrete, "
        "searchable subtopics. Return JSON with a 'subtopics' list."
    )
    return await call_deepseek_json(system, user_query, Subtopics)


async def research_subtopic(subtopic: str, max_search_results: int) -> ResearchResults:
    search_results = await web_search_multi(subtopic, max_results=max_search_results)
    system = (
        "Given web search results for a subtopic, extract relevant facts, data, "
        "source-backed claims, and insights. Return JSON with a 'findings' field."
    )
    user_prompt = f"Subtopic: {subtopic}\n\nWeb search results:\n{search_results}"
    return await call_deepseek_json(system, user_prompt, ResearchResults)


async def optimize_research(findings: str) -> OptimizationDecision:
    system = (
        "Decide whether the findings are sufficient to answer the original research "
        "request. Return JSON with 'justification' and 'needs_more_research'."
    )
    return await call_deepseek_json(system, f"Findings so far:\n{findings}", OptimizationDecision)


async def synthesize_report(findings: str) -> SummaryReport:
    system = (
        "Combine the provided research findings into a structured final report. "
        "Include an introduction, key findings, caveats, and conclusion. Return JSON "
        "with a 'report' field."
    )
    return await call_deepseek_json(system, findings, SummaryReport)


async def research_workflow(user_query: str) -> str:
    num_subtopics, max_search_results = get_research_depth(user_query)
    logger.info("Research depth: %s subtopics, %s results each", num_subtopics, max_search_results)

    subtopics_obj = await split_topic(user_query, num_subtopics)
    subtopics_list = subtopics_obj.subtopics[:num_subtopics]
    logger.info("Subtopics: %s", subtopics_list)

    research_tasks = [
        research_subtopic(subtopic, max_search_results)
        for subtopic in subtopics_list
    ]
    research_results = await asyncio.gather(*research_tasks)
    return "\n\n---\n\n".join(result.findings for result in research_results)


async def run_research_pipeline(user_query: str, memory: list[dict[str, str]]) -> str:
    try:
        context_text = "\n".join(
            f"{item['role']}: {item['content']}" for item in memory[-8:]
        )
        enhanced_query = f"Previous conversation:\n{context_text}\n\nCurrent query:\n{user_query}"

        research = await research_workflow(enhanced_query)
        optimizer_decision = await optimize_research(research)
        if optimizer_decision.needs_more_research:
            logger.info("Optimizer requested one additional research pass")
            research = await research_workflow(enhanced_query)

        final_report = await synthesize_report(research)
        report_text = final_report.as_text()

        memory.append({"role": "user", "content": user_query})
        memory.append({"role": "assistant", "content": report_text})
        return report_text
    except Exception as exc:
        logger.error("Pipeline error: %s", exc)
        return f"Error: {exc}"