bioai/agent/fireworks_client.py

"""bioai.agent.fireworks_client -- thin real Fireworks AI API client.

Uses ``requests`` only (no fireworks-ai SDK dependency -- keeps the install
small and the code portable to the ROCm container). Caches every response to
``.fireworks_cache/`` under the project root, keyed by ``sha256(prompt)`` with a
24-hour TTL so demo runs don't re-spend credits when the prompts are
identical.

If ``FIREWORKS_API_KEY`` is not set in the environment, the constructor
raises a clear ``RuntimeError("Set FIREWORKS_API_KEY env var")`` -- the
orchestrator catches this and falls back to a degraded-mode response so
the demo still runs end-to-end without an API key.
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from pathlib import Path
from typing import Dict, List, Optional

import requests

# --------------------------------------------------------------------------- #
# Constants
# --------------------------------------------------------------------------- #
FIREWORKS_ENDPOINT = "https://api.fireworks.ai/inference/v1/chat/completions"
DEFAULT_MODEL = "accounts/fireworks/models/llama-v3p1-70b-instruct"
CACHE_DIR = Path(__file__).resolve().parents[2] / ".fireworks_cache"
CACHE_TTL_SECONDS = 24 * 60 * 60  # 24 hours


# --------------------------------------------------------------------------- #
# FireworksClient
# --------------------------------------------------------------------------- #
class FireworksClient:
    """Real Fireworks AI chat-completions client with disk caching.

    Parameters
    ----------
    model:
        Fireworks model id. Defaults to Llama-3.1-70B-Instruct.
    api_key:
        Optional explicit API key. If ``None``, reads ``FIREWORKS_API_KEY``
        from the environment and raises if missing.
    cache_dir:
        Where to store cached responses.
    cache_ttl:
        Cache time-to-live in seconds (default 24 hours).
    timeout:
        HTTP timeout per request, in seconds.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        cache_dir: Path | str = CACHE_DIR,
        cache_ttl: int = CACHE_TTL_SECONDS,
        timeout: int = 60,
    ):
        self.model = model or DEFAULT_MODEL
        self.api_key = api_key or os.environ.get("FIREWORKS_API_KEY")
        if not self.api_key:
            raise RuntimeError("Set FIREWORKS_API_KEY env var")
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.cache_ttl = cache_ttl
        self.timeout = timeout
        # Reuse a session for connection pooling across calls.
        self._session = requests.Session()
        self._session.headers.update({
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        })

    # ------------------------------------------------------------------ #
    # Low-level chat
    # ------------------------------------------------------------------ #
    def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 2048,
    ) -> str:
        """Send a chat-completions request. Returns the assistant message text.

        ``messages`` is the standard OpenAI-style list of
        ``{"role": ..., "content": ...}`` dicts.
        """
        cache_key = self._cache_key(messages, temperature, max_tokens)
        cached = self._cache_get(cache_key)
        if cached is not None:
            print(
                f"[fireworks] CACHE HIT  model={self.model} "
                f"prompt_len={sum(len(m['content']) for m in messages)} "
                f"response_len={len(cached)}"
            )
            return cached

        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        prompt_len = sum(len(m["content"]) for m in messages)
        print(f"[fireworks] API CALL  model={self.model} prompt_len={prompt_len}")
        t0 = time.time()
        resp = self._session.post(
            FIREWORKS_ENDPOINT, data=json.dumps(payload), timeout=self.timeout
        )
        dt = time.time() - t0
        if resp.status_code != 200:
            # Surface the error body so the caller can log it.
            raise RuntimeError(
                f"Fireworks API returned {resp.status_code}: {resp.text[:500]}"
            )
        data = resp.json()
        text = (
            data.get("choices", [{}])[0]
            .get("message", {})
            .get("content", "")
        )
        # Cache and log
        self._cache_set(cache_key, text)
        print(
            f"[fireworks] API OK   model={self.model} "
            f"response_len={len(text)} elapsed={dt:.2f}s cached=False"
        )
        return text

    # ------------------------------------------------------------------ #
    # High-level helpers
    # ------------------------------------------------------------------ #
    def parse_pest_report(self, user_text: str) -> Dict:
        """Extract pest species, crop, severity, location from free text.

        Returns a dict with keys ``pest_species``, ``crop``, ``severity``,
        ``location``. On parse failure, returns a dict with ``_raw`` set to
        the raw model output and best-effort defaults.
        """
        system_prompt = (
            "You are an agricultural pest identification assistant. "
            "Extract structured data from the user's pest report and return "
            "STRICT JSON ONLY (no markdown fences, no commentary) with keys: "
            '"pest_species" (string), "crop" (string), "severity" (one of '
            '"low","moderate","high","severe"), "location" (string), '
            '"notes" (string, optional). If a field is unknown, use null.'
        )
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_text},
        ]
        raw = self.chat(messages, temperature=0.1, max_tokens=512)
        try:
            parsed = json.loads(raw)
        except json.JSONDecodeError:
            # Try to find a JSON block in the response.
            import re
            m = re.search(r"\{.*\}", raw, re.DOTALL)
            if m:
                try:
                    parsed = json.loads(m.group(0))
                except json.JSONDecodeError:
                    parsed = {}
            else:
                parsed = {}
        # Ensure all expected keys exist
        for key in ("pest_species", "crop", "severity", "location"):
            parsed.setdefault(key, None)
        parsed["_raw"] = raw
        return parsed

    def generate_safety_card(
        self,
        sirna_seq: str,
        offtarget_risks: Dict[str, float],
        half_life_hours: float,
    ) -> str:
        """Generate a markdown safety card for one siRNA candidate."""
        ot_str = "\n".join(
            f"  - {sp}: {risk:.3f}" for sp, risk in offtarget_risks.items()
        ) or "  (no off-target hits detected)"
        system_prompt = (
            "You are a regulatory toxicology writer. Produce a concise "
            "markdown SAFETY CARD for a dsRNA-based biopesticide siRNA "
            "candidate. Use only the data provided. Do not invent numbers. "
            "Sections: Sequence, Off-Target Profile, Environmental Fate, "
            "Overall Risk Tier (low/moderate/high). Keep it under 200 words."
        )
        user_prompt = (
            f"siRNA sequence (21 nt): {sirna_seq}\n\n"
            f"Off-target risks per species (0..1, fraction of 21-mers hit):\n{ot_str}\n\n"
            f"Predicted environmental half-life: {half_life_hours:.2f} hours\n"
        )
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
        return self.chat(messages, temperature=0.3, max_tokens=800)

    def generate_regulatory_memo(
        self,
        pest_species: str,
        candidates: List[Dict],
    ) -> str:
        """Generate an EPA-style regulatory memo summarising the top candidates.

        Each candidate dict should contain at least ``sirna_seq``,
        ``efficacy``, ``offtarget_max``, ``half_life_hours``, ``final_score``.
        """
        cand_lines = []
        for i, c in enumerate(candidates, start=1):
            cand_lines.append(
                f"  {i}. {c.get('sirna_seq', '?')}  "
                f"efficacy={c.get('efficacy', 0):.3f}  "
                f"offtarget_max={c.get('offtarget_max', 0):.3f}  "
                f"half_life={c.get('half_life_hours', 0):.2f}h  "
                f"score={c.get('final_score', 0):.3f}"
            )
        cand_block = "\n".join(cand_lines) or "  (no candidates provided)"

        system_prompt = (
            "You are an EPA FIFRA regulatory affairs consultant. Produce a "
            "concise markdown MEMO (under 400 words) recommending whether "
            "the listed dsRNA biopesticide candidates are suitable for an "
            "experimental use permit against the named pest. Sections: "
            "Pest & Crop, Candidate Summary, Risk Assessment, Recommendation. "
            "Be conservative; if any candidate has high off-target risk or "
            "very short half-life, flag it."
        )
        user_prompt = (
            f"Pest species: {pest_species}\n\n"
            f"Top candidates (sorted by final_score):\n{cand_block}\n"
        )
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
        return self.chat(messages, temperature=0.3, max_tokens=1500)

    # ------------------------------------------------------------------ #
    # Cache helpers
    # ------------------------------------------------------------------ #
    def _cache_key(
        self,
        messages: List[Dict[str, str]],
        temperature: float,
        max_tokens: int,
    ) -> str:
        blob = json.dumps(
            {"model": self.model, "messages": messages,
             "temperature": temperature, "max_tokens": max_tokens},
            sort_keys=True,
        )
        return hashlib.sha256(blob.encode("utf-8")).hexdigest()

    def _cache_get(self, key: str) -> Optional[str]:
        path = self.cache_dir / f"{key}.txt"
        if not path.exists():
            return None
        age = time.time() - path.stat().st_mtime
        if age > self.cache_ttl:
            return None
        return path.read_text(encoding="utf-8")

    def _cache_set(self, key: str, value: str) -> None:
        path = self.cache_dir / f"{key}.txt"
        path.write_text(value, encoding="utf-8")