CoolFace
Apppublic

mrme77/dfars-assistant

sourceHugging Faceupdated 1mo agoView on Hugging Face
0likes
answer.py130 linesDownload Raw Back to generation
1"""Answer generation interface for OpenRouter-backed models."""2 3import os4 5import httpx6from dotenv import load_dotenv7 8DEFAULT_OPENROUTER_MODEL = "google/gemini-2.5-flash"9DEFAULT_OPENROUTER_APP_TITLE = "DFARS App"10DEFAULT_OPENROUTER_REFERER = "http://127.0.0.1:8501"11OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"12 13 14def answer_with_openrouter(context: str) -> str:15    """Generate an answer from an OpenRouter chat model.16 17    Args:18        context: Fully assembled context package.19 20    Returns:21        Generated answer text.22 23    Raises:24        RuntimeError: If required environment variables are missing or the API fails.25    """26    load_dotenv()27    api_key = _load_openrouter_api_key()28    if not api_key:29        raise RuntimeError(30            "OPENROUTER_API_KEY is required. Set it in `.env` or configure "31            "OPENROUTER_KEYCHAIN_SERVICE and OPENROUTER_KEYCHAIN_ACCOUNT."32        )33 34    try:35        response = httpx.post(36            OPENROUTER_URL,37            headers=_openrouter_headers(api_key),38            json={39                "model": DEFAULT_OPENROUTER_MODEL,40                "messages": [41                    {42                        "role": "system",43                        "content": _system_prompt(),44                    },45                    {"role": "user", "content": context},46                ],47            },48            timeout=60,49        )50        response.raise_for_status()51    except httpx.HTTPStatusError as exc:52        raise RuntimeError(53            f"OpenRouter request failed with status {exc.response.status_code}: "54            f"{exc.response.text}"55        ) from exc56    except httpx.HTTPError as exc:57        raise RuntimeError(f"OpenRouter request failed: {exc}") from exc58 59    data = response.json()60    try:61        return str(data["choices"][0]["message"]["content"])62    except (KeyError, IndexError, TypeError) as exc:63        raise RuntimeError("OpenRouter response did not contain an answer.") from exc64 65 66def _system_prompt() -> str:67    """Return the fixed behavior rules for DFARS answer generation."""68    return "\n".join(69        [70            "You are a DFARS research assistant.",71            "Answer only from the provided retrieved DFARS context.",72            "Cite DFARS section identifiers and page ranges for every substantive claim.",73            "If the excerpts do not answer the question, say so.",74            "Do not provide legal advice or final contract determinations.",75        ]76    )77 78 79def _openrouter_headers(api_key: str) -> dict[str, str]:80    """Build OpenRouter request headers with app attribution.81 82    Args:83        api_key: OpenRouter API key.84 85    Returns:86        Headers for the OpenRouter chat completion request.87    """88    return {89        "Authorization": f"Bearer {api_key}",90        "Content-Type": "application/json",91        "X-Title": DEFAULT_OPENROUTER_APP_TITLE,92        "HTTP-Referer": DEFAULT_OPENROUTER_REFERER,93    }94 95 96def _load_openrouter_api_key() -> str | None:97    """Load the OpenRouter key from environment or macOS Keychain.98 99    Returns:100        API key string when configured, otherwise `None`.101 102    Raises:103        RuntimeError: If keychain lookup is configured but the keyring library fails.104    """105    api_key = os.getenv("OPENROUTER_API_KEY")106    if api_key:107        return api_key108 109    service = os.getenv("OPENROUTER_KEYCHAIN_SERVICE", "openrouter")110    account = os.getenv("OPENROUTER_KEYCHAIN_ACCOUNT", "OPENROUTER_API_KEY")111    if not service or not account:112        return None113 114    try:115        import keyring116    except ImportError as exc:117        raise RuntimeError(118            "The `keyring` package is required to read OPENROUTER_API_KEY "119            "from macOS Keychain. Install dependencies with "120            "`uv pip install -r requirements.txt --python dfars-env/bin/python`."121        ) from exc122 123    try:124        return keyring.get_password(service, account)125    except Exception as exc:126        raise RuntimeError(127            f"Could not read OpenRouter key from keychain service "128            f"`{service}` and account `{account}`."129        ) from exc130