synth/llm.py

"""Thin LLM client used to generate realistic file content for synthetic tasks.

Uses LiteLLM with the DeepSeek key from .mcp_env by default. If no key is set or
``--no-llm`` is passed, callers fall back to a built-in offline content generator
so the pipeline still works without network access.
"""

import json
import os
import random
import re
from pathlib import Path
from typing import List, Optional

from dotenv import load_dotenv

# Load .mcp_env at the repo root so DEEPSEEK_API_KEY etc. are available.
_REPO_ROOT = Path(__file__).resolve().parent.parent
load_dotenv(_REPO_ROOT / ".mcp_env")

_OFFLINE_WORDS = (
    "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi "
    "omicron pi rho sigma tau upsilon phi chi psi omega quartz nimbus harbor "
    "ledger compass orchard meadow lantern cobalt ember willow ".split()
)


class LLMClient:
    """Minimal completion wrapper around LiteLLM with an offline fallback."""

    def __init__(
        self,
        model: str = "openai/deepseek-v4-pro",
        api_key_var: str = "DEEPSEEK_API_KEY",
        enabled: bool = True,
        seed: int = 0,
        base_url: Optional[str] = None,
    ):
        self.model = model
        self.api_key = os.getenv(api_key_var)
        # dimcode OpenAI-compatible endpoint; falls back to env if not passed.
        self.base_url = base_url or os.getenv("DEEPSEEK_BASE_URL")
        self.enabled = enabled and bool(self.api_key)
        self._rng = random.Random(seed)

    # -- public API ---------------------------------------------------------

    def complete(self, prompt: str, temperature: float = 1.0,
                 max_tokens: int = 300) -> Optional[str]:
        """One-shot completion. Returns the text, or None if disabled/failed.

        Used for paraphrasing task questions (never for computing answers)."""
        if not self.enabled:
            return None
        try:
            import litellm
            litellm.suppress_debug_info = True
            resp = litellm.completion(
                model=self.model, api_key=self.api_key, base_url=self.base_url,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature, max_tokens=max_tokens,
            )
            return (resp.choices[0].message.content or "").strip()
        except Exception as e:  # pragma: no cover - network/parse issues
            print(f"[llm] complete() failed, falling back ({e})")
            return None

    def gen_snippets(self, theme: str, n: int, max_words: int = 30) -> List[dict]:
        """Return ``n`` dicts of {"filename", "content"} on a theme.

        Filenames are slugs without extension (the caller assigns extensions and
        adjusts sizes). Falls back to offline lorem text on any failure.
        """
        if self.enabled:
            try:
                return self._gen_snippets_llm(theme, n, max_words)
            except Exception as e:  # pragma: no cover - network/parse issues
                print(f"[llm] falling back to offline content ({e})")
        return self._gen_snippets_offline(theme, n, max_words)

    # -- LLM path -----------------------------------------------------------

    def _gen_snippets_llm(self, theme: str, n: int, max_words: int) -> List[dict]:
        import litellm

        litellm.suppress_debug_info = True
        prompt = (
            f"Generate {n} short, realistic text snippets for files that might appear "
            f"on a desktop/project about: {theme}.\n"
            f"Return ONLY a JSON array of objects with keys 'filename' (a short "
            f"lowercase slug, no extension, e.g. 'meeting_notes', unique across the "
            f"array) and 'content' (1-2 sentences, at most {max_words} words, plain "
            f"text, no markdown headers). No commentary, no code fences."
        )
        resp = litellm.completion(
            model=self.model,
            api_key=self.api_key,
            base_url=self.base_url,
            messages=[{"role": "user", "content": prompt}],
            temperature=1.0,
            max_tokens=2000,
        )
        text = resp.choices[0].message.content
        items = self._parse_json_array(text)
        out: List[dict] = []
        seen = set()
        for it in items:
            name = _slug(str(it.get("filename", "")).strip())
            body = str(it.get("content", "")).strip()
            if not name or not body:
                continue
            name = _dedupe(name, seen)
            out.append({"filename": name, "content": body})
        # Top up if the model returned too few.
        if len(out) < n:
            out.extend(self._gen_snippets_offline(theme, n - len(out), max_words, seen))
        return out[:n]

    @staticmethod
    def _parse_json_array(text: str) -> List[dict]:
        text = text.strip()
        # Strip code fences if present.
        text = re.sub(r"^```(?:json)?|```$", "", text, flags=re.MULTILINE).strip()
        start, end = text.find("["), text.rfind("]")
        if start != -1 and end != -1:
            text = text[start : end + 1]
        data = json.loads(text)
        if not isinstance(data, list):
            raise ValueError("expected a JSON array")
        return data

    # -- offline path -------------------------------------------------------

    def _gen_snippets_offline(
        self, theme: str, n: int, max_words: int, seen: Optional[set] = None
    ) -> List[dict]:
        seen = seen if seen is not None else set()
        out = []
        for _ in range(n):
            wc = self._rng.randint(6, max_words)
            words = [self._rng.choice(_OFFLINE_WORDS) for _ in range(wc)]
            name = _dedupe(_slug("_".join(words[:2]) or "note"), seen)
            body = (theme + ": " + " ".join(words)).strip().capitalize() + "."
            out.append({"filename": name, "content": body})
        return out


def _slug(s: str) -> str:
    s = re.sub(r"[^a-zA-Z0-9]+", "_", s).strip("_").lower()
    return s or "file"


def _dedupe(name: str, seen: set) -> str:
    base, i = name, 1
    while name in seen:
        i += 1
        name = f"{base}_{i}"
    seen.add(name)
    return name