pipeline/search_terms.py

"""Step 1: LLM turns the topic/keywords/brief into optimized web-search queries."""
from __future__ import annotations

import json
import re
from typing import List

from huggingface_hub import InferenceClient

from . import config, llm

_SYSTEM = (
    "You are an SEO research assistant. Given a blog topic, its target keywords, and a "
    "brief, produce a small set of high-signal web search queries that will surface the "
    "most authoritative, information-rich sources to write the post from. Vary angle and "
    "specificity. Return ONLY a JSON array of query strings, nothing else."
)


def _fallback_terms(topic: str, primary: str, secondary: str) -> List[str]:
    base = [t for t in [topic, primary, secondary] if t]
    extra = [f"{topic} {primary}".strip(), f"{topic} guide", f"{topic} best practices"]
    seen, out = set(), []
    for t in base + extra:
        t = t.strip()
        if t and t.lower() not in seen:
            seen.add(t.lower())
            out.append(t)
    return out[: config.N_SEARCH_TERMS]


def generate_search_terms(
    client: InferenceClient,
    topic: str,
    primary_keyword: str,
    secondary_keyword: str,
    brief: str,
) -> List[str]:
    user = (
        f"Topic: {topic}\n"
        f"Primary keyword: {primary_keyword}\n"
        f"Secondary keyword: {secondary_keyword}\n"
        f"Brief: {brief}\n\n"
        f"Return {config.N_SEARCH_TERMS} search queries as a JSON array."
    )
    try:
        raw = llm.chat(
            client, config.MODEL_REASONING, _SYSTEM, user,
            max_tokens=400, temperature=0.4,
        )
        terms = _parse_terms(raw)
        if terms:
            return terms[: config.N_SEARCH_TERMS]
    except Exception:
        pass
    return _fallback_terms(topic, primary_keyword, secondary_keyword)


def _parse_terms(raw: str) -> List[str]:
    """Extract a list of strings from a (possibly fenced) LLM response."""
    match = re.search(r"\[.*\]", raw, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group(0))
            return [str(t).strip() for t in data if str(t).strip()]
        except Exception:
            pass
    # line-based fallback
    lines = [re.sub(r'^[\s\-\*\d\.\)"]+', "", ln).strip().strip('"') for ln in raw.splitlines()]
    return [ln for ln in lines if ln]