FabledScribe/src/fabledassistant/services/research.py

"""Web research pipeline: sub-queries → SearXNG → fetch → synthesize → note."""

import asyncio
import json
import logging
import re

import httpx

from fabledassistant.config import Config
from fabledassistant.services.llm import fetch_url_content, generate_completion, stream_chat
from fabledassistant.services.notes import create_note
from fabledassistant.models.note import Note

logger = logging.getLogger(__name__)

SEARXNG_QUERIES = 5      # sub-queries to generate
RESULTS_PER_QUERY = 3    # results fetched from SearXNG per query
PAGES_PER_QUERY = 3      # pages actually read per sub-query (top N results)
MAX_SYNTHESIS_SOURCES = 12  # deduplicated sources passed to synthesis LLM
CHARS_PER_SOURCE = 2000     # content chars per source sent to synthesis


async def run_research_pipeline(
    topic: str,
    user_id: int,
    model: str,
    buf=None,
    project_id: int | None = None,
) -> Note:
    """Full research pipeline: search → fetch → synthesize → create note.

    Emits status events via buf.append_event throughout (when buf is provided).
    Returns the created Note.
    """
    # Step 1: Generate sub-queries
    if buf is not None:
        buf.append_event("status", {"status": "Generating search queries..."})
    queries = await _generate_sub_queries(topic, model)
    logger.info("Research: generated %d sub-queries for topic '%s'", len(queries), topic)

    # Step 2: Search all queries in parallel (200 ms stagger to avoid hammering SearXNG)
    async def _search_with_stagger(i: int, query: str) -> tuple[str, list[dict]]:
        if i > 0:
            await asyncio.sleep(0.2 * i)
        if buf is not None:
            buf.append_event("status", {"status": f"Searching: {query}..."})
        results = await _search_searxng(query)
        logger.info("Research: query '%s' → %d results", query, len(results))
        return query, results

    search_results = await asyncio.gather(
        *[_search_with_stagger(i, q) for i, q in enumerate(queries)]
    )

    # Deduplicate URLs across all queries
    seen_urls: set[str] = set()
    url_tasks: list[tuple[str, dict, str]] = []  # (url, result_dict, query)
    for query, results in search_results:
        for result in results[:PAGES_PER_QUERY]:
            url = result.get("url", "")
            if url and url not in seen_urls:
                seen_urls.add(url)
                url_tasks.append((url, result, query))

    # Fetch all unique URLs in parallel
    async def _fetch_source(url: str, result: dict, query: str) -> dict:
        title = result.get("title", url)
        if buf is not None:
            buf.append_event("status", {"status": f"Reading: {title[:60]}..."})
        content = await fetch_url_content(url)
        return {
            "url": url,
            "title": title,
            "query": query,
            "snippet": result.get("snippet", ""),
            "content": content,
        }

    all_sources: list[dict] = list(await asyncio.gather(
        *[_fetch_source(url, result, query) for url, result, query in url_tasks]
    ))

    if not all_sources:
        raise ValueError(f"No results found for '{topic}'")

    # Step 3: Filter failed fetches
    good_sources = [
        s for s in all_sources
        if not s["content"].startswith("[Failed to fetch")
    ]

    if not good_sources:
        raise ValueError(f"Could not read any sources for '{topic}'")

    # Limit to top N sources for synthesis (already deduplicated by URL)
    synthesis_sources = good_sources[:MAX_SYNTHESIS_SOURCES]
    logger.info(
        "Research: %d/%d sources successfully fetched, using %d for synthesis",
        len(good_sources), len(all_sources), len(synthesis_sources),
    )

    # Step 4: Synthesize (streams tokens into chat as the note is being written)
    if buf is not None:
        buf.append_event("status", {"status": f"Synthesizing report from {len(synthesis_sources)} sources..."})
    title, body = await _synthesize_note(topic, synthesis_sources, model, buf)

    # Step 5: Create note
    if buf is not None:
        buf.append_event("status", {"status": "Saving note..."})
    note = await create_note(
        user_id=user_id,
        title=title,
        body=body,
        tags=["research"],
        project_id=project_id,
    )
    logger.info("Research: created note id=%d title='%s'", note.id, note.title)
    return note


async def _generate_sub_queries(topic: str, model: str) -> list[str]:
    """Ask the model for focused search queries for the topic."""
    messages = [
        {
            "role": "system",
            "content": (
                f"You are a research assistant. Given a research topic, generate exactly {SEARXNG_QUERIES} "
                "focused web search queries that together would provide comprehensive coverage of the topic. "
                "Vary the angle of each query: include overview, implementation details, best practices, "
                "common problems, and real-world examples. "
                "Respond with ONLY a JSON array of strings, no other text. "
                'Example: ["query one", "query two", "query three"]'
            ),
        },
        {"role": "user", "content": f"Topic: {topic}"},
    ]
    try:
        raw = await generate_completion(messages, model, max_tokens=200)
        raw = raw.strip()
        raw = re.sub(r"^```(?:json)?\s*", "", raw)
        raw = re.sub(r"\s*```$", "", raw)
        idx = raw.find("[")
        if idx >= 0:
            parsed, _ = json.JSONDecoder().raw_decode(raw[idx:])
            if isinstance(parsed, list) and parsed:
                queries = [str(q).strip() for q in parsed if str(q).strip()]
                if queries:
                    return queries[:SEARXNG_QUERIES]
    except Exception:
        logger.warning("Sub-query generation failed, falling back to topic", exc_info=True)
    return [topic]


async def _search_searxng(query: str) -> list[dict]:
    """Search SearXNG and return top results as [{url, title, snippet}]."""
    url = Config.SEARXNG_URL.rstrip("/") + "/search"
    params = {"q": query, "format": "json", "categories": "general"}
    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(url, params=params)
                if resp.status_code == 429:
                    retry_after = int(resp.headers.get("Retry-After", "5"))
                    wait = min(retry_after, 10) * (attempt + 1)
                    logger.warning(
                        "SearXNG 429 for query '%s' (attempt %d/3), waiting %ds",
                        query, attempt + 1, wait,
                    )
                    await asyncio.sleep(wait)
                    continue
                resp.raise_for_status()
                data = resp.json()
                results = data.get("results", [])
                out = []
                for r in results[:RESULTS_PER_QUERY]:
                    out.append({
                        "url": r.get("url", ""),
                        "title": r.get("title", ""),
                        "snippet": r.get("content", ""),
                    })
                return out
        except httpx.HTTPStatusError:
            logger.warning("SearXNG search failed for query '%s'", query, exc_info=True)
            return []
        except Exception:
            logger.warning("SearXNG search failed for query '%s'", query, exc_info=True)
            return []
    logger.warning("SearXNG search gave up after 3 attempts for query '%s'", query)
    return []


async def _search_searxng_images(query: str) -> list[dict]:
    """Search SearXNG image category and return [{img_src, page_url, title, source_domain}]."""
    url = Config.SEARXNG_URL.rstrip("/") + "/search"
    params = {"q": query, "format": "json", "categories": "images"}
    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(url, params=params)
                if resp.status_code == 429:
                    retry_after = int(resp.headers.get("Retry-After", "5"))
                    wait = min(retry_after, 10) * (attempt + 1)
                    logger.warning(
                        "SearXNG image 429 for '%s' (attempt %d/3), waiting %ds",
                        query, attempt + 1, wait,
                    )
                    await asyncio.sleep(wait)
                    continue
                resp.raise_for_status()
                data = resp.json()
                out = []
                for r in data.get("results", []):
                    img_src = r.get("img_src") or r.get("thumbnail_src", "")
                    if not img_src:
                        continue
                    try:
                        from urllib.parse import urlparse
                        source_domain = urlparse(r.get("url", "")).netloc or ""
                    except Exception:
                        source_domain = ""
                    out.append({
                        "img_src": img_src,
                        "page_url": r.get("url", ""),
                        "title": r.get("title", ""),
                        "source_domain": source_domain,
                    })
                return out
        except httpx.HTTPStatusError:
            logger.warning("SearXNG image search failed for '%s'", query, exc_info=True)
            return []
        except Exception:
            logger.warning("SearXNG image search failed for '%s'", query, exc_info=True)
            return []
    logger.warning("SearXNG image search gave up after 3 attempts for '%s'", query)
    return []


async def _synthesize_note(
    topic: str,
    sources: list[dict],
    model: str,
    buf=None,
) -> tuple[str, str]:
    """Synthesize a comprehensive markdown research document from fetched sources.

    Returns (title, body_markdown).
    When buf is provided, tokens are streamed into the chat buffer in real time
    so the user can see the note being written. Uses an extended context window.
    """
    sources_text_parts = []
    for i, s in enumerate(sources, 1):
        content = (s.get("content") or s.get("snippet") or "")[:CHARS_PER_SOURCE]
        sources_text_parts.append(
            f"[Source {i}] {s['title']}\nURL: {s['url']}\nSearch query: {s['query']}\n\n{content}"
        )
    sources_block = "\n\n" + ("─" * 60) + "\n\n".join(sources_text_parts)

    messages = [
        {
            "role": "system",
            "content": (
                "You are a thorough researcher and writer. "
                "Your task is to write an exhaustive, well-structured document on the given topic — "
                "not a brief summary or intro paragraph.\n\n"
                "Requirements:\n"
                "- Write at least 2500 words of substantive content (excluding the Sources section)\n"
                "- Choose sections (##) that make sense for the topic — let the subject matter determine the structure. "
                "A technical topic might need implementation, configuration, and troubleshooting sections. "
                "A comparison topic might need dedicated sections per subject being compared plus a summary. "
                "A scientific topic might need background, mechanisms, research findings, and implications. "
                "Use your judgment — minimum 6 major sections.\n"
                "- Use ### for subsections where they add clarity\n"
                "- Write in detailed prose paragraphs — do not reduce sections to bullet-point lists\n"
                "- Include specific details, examples, data points, comparisons, and nuance from the sources\n"
                "- Do not pad with vague generalities — every paragraph should say something concrete\n"
                "- The first line must be the document title starting with '# '\n"
                "- End with a '## Sources' section listing every source as a markdown hyperlink\n\n"
                "The reader wants to finish this document with a thorough understanding of the topic, "
                "not just an overview."
            ),
        },
        {
            "role": "user",
            "content": (
                f"Write a comprehensive reference document on: {topic}\n\n"
                f"Sources ({len(sources)} pages fetched):\n{sources_block}"
            ),
        },
    ]

    if buf is not None:
        # Stream tokens into the chat buffer so the user sees the note being written
        raw_parts: list[str] = []
        async for token in stream_chat(
            messages, model, options={"num_ctx": 16384, "num_predict": 8192}
        ):
            raw_parts.append(token)
            buf.append_event("chunk", {"chunk": token})
            buf.content_so_far += token
        raw = "".join(raw_parts).strip()
    else:
        raw = await generate_completion(
            messages,
            model,
            max_tokens=8192,
            num_ctx=16384,
        )
        raw = raw.strip()

    # Extract title from first # heading
    lines = raw.splitlines()
    title = f"Research: {topic}"
    body_lines = lines
    if lines and lines[0].startswith("# "):
        title = lines[0][2:].strip()
        body_lines = lines[1:]

    body = "\n".join(body_lines).strip()
    return title, body