stonks-oracle/services/recommendation/thesis_llm.py

"""Optional LLM wording layer for thesis generation.

Takes a deterministic thesis string (built from trend data and eligibility
rules) and rewrites it into natural, analyst-quality prose using a local
Ollama model. The deterministic thesis is always preserved as the fallback
and audit reference.

This module is opt-in: callers must explicitly request LLM rewriting.
If the LLM call fails or is disabled, the original deterministic thesis
is returned unchanged.

Requirements: 7.1, 7.2
"""
from __future__ import annotations

import logging
import time

import httpx

from services.shared.config import OllamaConfig
from services.shared.schemas import TrendSummary

logger = logging.getLogger(__name__)

THESIS_PROMPT_VERSION = "thesis-rewrite-v1"

THESIS_SYSTEM_PROMPT = """\
You are a concise financial analyst. You rewrite structured trade thesis \
summaries into clear, professional prose suitable for an internal research note.

STRICT RULES:
1. Do NOT add any information that is not present in the input.
2. Do NOT fabricate numbers, dates, company names, or analyst opinions.
3. Keep the rewrite under 150 words.
4. Preserve all factual claims, risk notes, and evidence counts from the input.
5. Use a neutral, professional tone. Avoid hype or marketing language.
6. Return ONLY the rewritten thesis text. No JSON, no markdown, no commentary."""


def build_thesis_rewrite_prompt(
    deterministic_thesis: str,
    summary: TrendSummary,
) -> dict[str, str]:
    """Build system and user prompts for thesis rewriting.

    Provides the model with the deterministic thesis and key trend
    context so it can produce a natural-language version.
    """
    context_parts = [
        f"Ticker: {summary.entity_id}",
        f"Window: {summary.window.value}",
        f"Direction: {summary.trend_direction.value}",
        f"Strength: {summary.trend_strength:.2f}",
        f"Confidence: {summary.confidence:.2f}",
        f"Contradiction score: {summary.contradiction_score:.2f}",
    ]
    if summary.dominant_catalysts:
        context_parts.append(f"Catalysts: {', '.join(summary.dominant_catalysts[:3])}")
    if summary.material_risks:
        context_parts.append(f"Risks: {'; '.join(summary.material_risks[:2])}")

    context_block = "\n".join(context_parts)

    user_prompt = f"""\
Rewrite the following structured thesis into clear, professional analyst prose.

--- STRUCTURED THESIS ---
{deterministic_thesis}
--- END STRUCTURED THESIS ---

--- CONTEXT ---
{context_block}
--- END CONTEXT ---

Return ONLY the rewritten thesis. No other text."""

    return {
        "system": THESIS_SYSTEM_PROMPT,
        "user": user_prompt,
    }


async def rewrite_thesis_with_llm(
    deterministic_thesis: str,
    summary: TrendSummary,
    config: OllamaConfig,
    http_client: httpx.AsyncClient | None = None,
) -> str:
    """Rewrite a deterministic thesis using a local Ollama model.

    If the LLM call fails for any reason, returns the original
    deterministic thesis unchanged. This ensures the LLM layer is
    purely additive and never blocks recommendation generation.

    Args:
        deterministic_thesis: The rule-based thesis string.
        summary: The trend summary that produced the thesis.
        config: Ollama connection and model configuration.
        http_client: Optional shared HTTP client for connection reuse.

    Returns:
        The LLM-rewritten thesis on success, or the original on failure.
    """
    prompts = build_thesis_rewrite_prompt(deterministic_thesis, summary)

    owns_client = http_client is None
    client = http_client or httpx.AsyncClient(timeout=config.timeout)

    try:
        rewritten = await _call_ollama_thesis(client, config, prompts)
        if rewritten:
            logger.info(
                "LLM thesis rewrite succeeded for %s (%d chars → %d chars)",
                summary.entity_id,
                len(deterministic_thesis),
                len(rewritten),
            )
            return rewritten

        logger.warning(
            "LLM thesis rewrite returned empty for %s — using deterministic thesis",
            summary.entity_id,
        )
        return deterministic_thesis
    except Exception:
        logger.exception(
            "LLM thesis rewrite failed for %s — using deterministic thesis",
            summary.entity_id,
        )
        return deterministic_thesis
    finally:
        if owns_client:
            await client.aclose()


async def _call_ollama_thesis(
    client: httpx.AsyncClient,
    config: OllamaConfig,
    prompts: dict[str, str],
) -> str:
    """Make a single Ollama chat call for thesis rewriting.

    Returns the model's text response, or empty string on failure.
    """
    start = time.monotonic()

    payload = {
        "model": config.model,
        "messages": [
            {"role": "system", "content": prompts["system"]},
            {"role": "user", "content": prompts["user"]},
        ],
        "stream": False,
    }

    resp = await client.post(
        f"{config.base_url}/api/chat",
        json=payload,
    )
    _ = resp.raise_for_status()

    duration_ms = int((time.monotonic() - start) * 1000)

    body: dict[str, object] = resp.json()
    msg = body.get("message")
    content: str = msg.get("content", "") if isinstance(msg, dict) else ""

    logger.debug(
        "Ollama thesis call completed in %dms, response length=%d",
        duration_ms,
        len(content),
    )

    return content.strip()