backend/llm.py

# backend/llm.py
"""
Pluggable LLM narration layer.

Backend selection (env `LLM_BACKEND`):
  - "ollama" (default): fully local via Ollama. Documents never leave the
    machine. Slower on CPU, but private.
  - "groq": cloud, fast. Opt-in only — requires GROQ_API_KEY.

Deliberate design: there is NO silent fallback from local to cloud. If a
user chose the private backend and it's down, sending their filing to a
cloud API anyway would be a privacy violation, not a convenience. Instead
the caller gets an honest Mode-C message. Metric lookup, red flags, and
recommendations don't pass through here at all — they keep working with
no LLM whatsoever.
"""

import os
import re

import requests
from dotenv import load_dotenv

load_dotenv()

LLM_BACKEND = os.getenv("LLM_BACKEND", "ollama").strip().lower()

OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
# Default is deliberately small: narration only needs to phrase numbers the
# extractor already found, and a 1.5B q4 model runs in ~2GB RAM — it works
# even on a loaded 16GB laptop where a 4B model's KV cache gets OOM-killed.
# Swap via OLLAMA_MODEL when more memory is free.
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "qwen2.5:1.5b")
# Context cap keeps Ollama's KV-cache memory budget sane; our prompts are
# top-3 chunks + a metrics header, well under 4096 tokens.
OLLAMA_NUM_CTX = int(os.getenv("OLLAMA_NUM_CTX", "4096"))
# Local CPU generation is slow — a 400-token answer can take minutes.
# Generous by design; the frontend shows latency.
OLLAMA_TIMEOUT = int(os.getenv("OLLAMA_TIMEOUT", "300"))

GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")


class LLMUnavailable(Exception):
    """Raised when the selected backend can't produce a response."""


MODE_C_MESSAGE = (
    "LLM narration is unavailable (backend '{backend}': {reason}). "
    "Metric lookups, red flags, and recommendations still work without it. "
    "Start Ollama for local narration, or set LLM_BACKEND=groq with a "
    "GROQ_API_KEY for cloud narration."
)


# ── backends ─────────────────────────────────────────────────


def _generate_ollama(prompt: str, max_tokens: int) -> str:
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0,
            "num_predict": max_tokens,
            "num_ctx": OLLAMA_NUM_CTX,
        },
    }
    # qwen3/deepseek-r1 style models emit a chain-of-thought block that
    # burns the CPU token budget before the actual answer; Ollama >= 0.9
    # accepts `think: false` for them. Only sent for models known to
    # support it — some versions reject the flag on non-thinking models.
    if OLLAMA_MODEL.split(":")[0] in ("qwen3", "deepseek-r1"):
        payload["think"] = False

    try:
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json=payload,
            timeout=OLLAMA_TIMEOUT,
        )
        resp.raise_for_status()
        text = resp.json().get("response", "").strip()
    except requests.RequestException as e:
        raise LLMUnavailable(f"Ollama not reachable at {OLLAMA_URL} ({e.__class__.__name__})")

    # belt-and-braces: strip any thinking block that got through anyway
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

    if not text:
        raise LLMUnavailable(f"Ollama model {OLLAMA_MODEL} returned an empty response")
    return text


_groq_client = None


def _generate_groq(prompt: str, max_tokens: int) -> str:
    global _groq_client

    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        raise LLMUnavailable("GROQ_API_KEY is not set")

    if _groq_client is None:
        from groq import Groq
        _groq_client = Groq(api_key=api_key)

    try:
        r = _groq_client.chat.completions.create(
            model=GROQ_MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
            temperature=0,
        )
        return r.choices[0].message.content.strip()
    except Exception as e:
        raise LLMUnavailable(f"Groq call failed: {e}")


_BACKENDS = {
    "ollama": _generate_ollama,
    "groq": _generate_groq,
}


def generate(prompt: str, max_tokens: int = 400) -> str:
    """Route a prompt to the selected backend. Returns the Mode-C message
    instead of raising, so every endpoint stays a 200 with an honest
    explanation rather than a 500."""
    backend = _BACKENDS.get(LLM_BACKEND)
    if backend is None:
        return MODE_C_MESSAGE.format(
            backend=LLM_BACKEND, reason="unknown backend name"
        )
    try:
        return backend(prompt, max_tokens)
    except LLMUnavailable as e:
        return MODE_C_MESSAGE.format(backend=LLM_BACKEND, reason=e)


# ── public narration API (signatures unchanged for main.py) ──


def ask(question: str, context: str, max_tokens: int = 400) -> str:
    prompt = f"""You are a professional financial analyst.

Use ONLY the context below to answer the question.

Rules:
- Use only facts explicitly present in the context.
- Do not use external knowledge.
- Do not infer or speculate.
- If the answer is not present, reply: "Not available in filing."

Context:
{context}

Question:
{question}

Answer:"""

    return generate(prompt, max_tokens)


def generate_report(company: str, year: str, metrics: dict, context: str) -> str:
    def fmt(val):
        if val is None:
            return "N/A"
        if val >= 1_000_000_000:
            return f"${val/1_000_000_000:.1f}B"
        if val >= 1_000_000:
            return f"${val/1_000_000:.1f}M"
        return f"${val:.2f}"

    metrics_str = "\n".join(
        [
            f"- {k.replace('_', ' ').title()}: {fmt(v)}"
            for k, v in metrics.items()
        ]
    )

    prompt = f"""You are a senior Wall Street analyst. Generate a financial report.

Company: {company}
Year: {year}

Extracted Metrics:
{metrics_str}

Source Document Context:
{context[:3000]}

Write a report with these sections:

1. Executive Summary
2. Financial Performance
3. Bull Case
4. Bear Case
5. Risk Score (1-10)

Rules:
- Use ONLY facts explicitly present in the metrics or context.
- Do NOT infer, assume, speculate, or use external knowledge.
- Every claim must be traceable to the provided context.
- If evidence for a section is missing, write: "Not available in filing."
- Do not mention products, initiatives, risks, strategies, or events unless explicitly stated in the context.
"""

    return generate(prompt, max_tokens=800)


def compare_companies(
    companies: list,
    metric_data: dict,
    question: str = None
) -> str:

    lines = []

    for company, years in metric_data.items():
        for year, val in years.items():
            if val is None:
                continue

            if val >= 1_000_000_000:
                value = f"${val/1_000_000_000:.1f}B"
            else:
                value = f"${val/1_000_000:.1f}M"

            lines.append(f"{company} ({year}): {value}")

    data_str = "\n".join(lines)

    q = question or "Compare these companies based only on the provided data."

    prompt = f"""You are a financial analyst.

Use ONLY the information below.

Data:
{data_str}

Question:
{q}

Rules:
- Use only the provided data.
- Do not use external knowledge.
- Do not speculate.
- If data is insufficient, say so.
- Keep the answer concise.
"""

    return generate(prompt, max_tokens=400)


if __name__ == "__main__":
    context = "Apple reported revenue of $394.3 billion in FY2022, up 8% YoY."
    print(f"Backend: {LLM_BACKEND} (model: {OLLAMA_MODEL if LLM_BACKEND == 'ollama' else GROQ_MODEL})")
    print(ask("What was Apple's revenue?", context))