llm_router.py

from typing import Optional, List
import cohere
from settings import COHERE_API_KEY, COHERE_MODEL_PRIMARY, MODEL_SETTINGS
from local_llm import LocalLLM

_local = None

def _local_llm() -> LocalLLM:
    global _local
    if _local is None:
        _local = LocalLLM()
    return _local

def cohere_chat(prompt: str) -> Optional[str]:
    if not COHERE_API_KEY:
        return None
    try:
        cli = cohere.Client(api_key=COHERE_API_KEY)
        resp = cli.chat(
            model=COHERE_MODEL_PRIMARY,
            message=prompt,
            temperature=MODEL_SETTINGS["temperature"],
            max_tokens=MODEL_SETTINGS["max_new_tokens"],
        )
        if hasattr(resp, "text") and resp.text: return resp.text
        if hasattr(resp, "reply") and resp.reply: return resp.reply
        if hasattr(resp, "generations") and resp.generations: return resp.generations[0].text
    except Exception:
        return None
    return None

def open_fallback_chat(prompt: str) -> Optional[str]:
    return _local_llm().chat(prompt)

def generate_narrative(scenario_text: str, structured_sections_md: str, rag_snippets: List[str]) -> str:
    grounding = "\n\n".join([f"[RAG {i+1}]\n{t}" for i, t in enumerate(rag_snippets or [])])
    prompt = f"""You are a Canadian healthcare operations copilot.
Follow the scenario's requested deliverables exactly. Use the structured computations provided (already calculated deterministically) and the RAG snippets for grounding.

# Scenario
{scenario_text}

# Deterministic Results (already computed)
{structured_sections_md}

# Grounding (Canadian sources, snippets)
{grounding}

Write a concise, decision-ready report tailored to provincial operations leaders.
Do not invent numbers. If data are missing, say so clearly.
"""
    out = cohere_chat(prompt)
    if out: return out
    out = open_fallback_chat(prompt)
    if out: return out
    return "Unable to generate narrative at this time."