bioai/agent/ollama_client.py

"""bioai.agent.ollama_client -- thin local Ollama API client.

Drop-in replacement for the Fireworks client that talks to a local Ollama
server (https://ollama.com) instead of a remote API. This eliminates the
need for any cloud LLM credits: Ollama runs on the same machine as the
pipeline (laptop, AMD cloud VM, or Docker container) and serves the Llama
3.2 3B model locally.

Setup (one-time):
    1. Install Ollama:  curl -fsSL https://ollama.com/install.sh | sh
    2. Pull the model:  ollama pull llama3.2:3b
    3. Start the server: ollama serve   (or it auto-starts on demand)

The client uses ``requests`` only (no ollama-python SDK dependency -- keeps
the install small and portable). Caches every response to
``<project_root>/.ollama_cache/`` keyed by ``sha256(prompt)`` with a 24-hour
TTL so demo runs don't re-spend compute when the prompts are identical.

If Ollama is unreachable (e.g., server not running, model not pulled), the
constructor raises a clear ``RuntimeError`` -- the orchestrator catches this
and falls back to a degraded-mode response so the demo still runs end-to-end
without a working Ollama server.

Env vars (all optional):
    OLLAMA_HOST       Override the Ollama server URL (default http://localhost:11434)
    OLLAMA_MODEL      Override the model name (default llama3.2:3b)
"""

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
# --------------------------------------------------------------------------- #
DEFAULT_OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
DEFAULT_MODEL = os.environ.get("OLLAMA_MODEL", "llama3.2:3b")
CACHE_TTL_SECONDS = 24 * 60 * 60  # 24 hours

# Portable cache dir (resolved from bioai.paths so it's not hardcoded)
from bioai.paths import CACHE_DIR as _PROJECT_CACHE_DIR  # noqa: E402
CACHE_DIR = _PROJECT_CACHE_DIR.parent / ".ollama_cache"


# --------------------------------------------------------------------------- #
# OllamaClient
# --------------------------------------------------------------------------- #
class OllamaClient:
    """Local Ollama chat-completions client with disk caching.

    Same API surface as the old FireworksClient so it's a drop-in replacement.

    Parameters
    ----------
    model:
        Ollama model tag. Defaults to ``llama3.2:3b``.
    host:
        Ollama server URL. Defaults to ``OLLAMA_HOST`` env var or
        ``http://localhost:11434``.
    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. Local Ollama is fast for 3B
        but cold-start of the model can take 10-15s, so default is 120.
    """

    def __init__(
        self,
        model: Optional[str] = None,
        host: Optional[str] = None,
        cache_dir: Path | str = CACHE_DIR,
        cache_ttl: int = CACHE_TTL_SECONDS,
        timeout: int = 120,
    ):
        self.model = model or DEFAULT_MODEL
        self.host = (host or DEFAULT_OLLAMA_HOST).rstrip("/")
        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()
        # Probe the server once at construction; raise if unreachable so the
        # orchestrator can fall back to degraded mode.
        self._verify_server()

    def _verify_server(self) -> None:
        """Hit ``GET /api/tags`` to confirm Ollama is running and the model is available.

        Raises RuntimeError with a helpful message on failure.
        Uses a short timeout (3s) so the demo doesn't hang when Ollama isn't running
        (e.g., on HuggingFace Spaces where Ollama can't be installed).
        """
        try:
            resp = self._session.get(f"{self.host}/api/tags", timeout=3)
            if resp.status_code != 200:
                raise RuntimeError(
                    f"Ollama server at {self.host} returned HTTP {resp.status_code}. "
                    f"Is `ollama serve` running?"
                )
            data = resp.json()
            available = {m.get("name", "") for m in data.get("models", [])}
            # Some Ollama versions append :latest, some don't. Check both.
            candidates = {self.model, self.model.split(":")[0], self.model + ":latest"}
            if not (available & candidates):
                # Don't hard-fail: the model might be pullable on demand. Just warn.
                print(
                    f"[ollama] WARNING: model '{self.model}' not found in available models "
                    f"({sorted(available)[:5]}...). Attempting to use it anyway -- "
                    f"run `ollama pull {self.model}` if the call fails."
                )
        except requests.exceptions.ConnectionError as e:
            raise RuntimeError(
                f"Cannot reach Ollama server at {self.host}. "
                f"Start it with `ollama serve` or set OLLAMA_HOST env var. "
                f"Original error: {e}"
            ) from e
        except Exception as e:
            raise RuntimeError(
                f"Ollama server probe failed: {type(e).__name__}: {e}"
            ) from e

    # ------------------------------------------------------------------ #
    # 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.

        Uses Ollama's OpenAI-compatible endpoint ``/v1/chat/completions`` so
        the request/response shape matches what we used for Fireworks.
        """
        cache_key = self._cache_key(messages, temperature, max_tokens)
        cached = self._cache_get(cache_key)
        if cached is not None:
            print(
                f"[ollama] CACHE HIT  model={self.model} "
                f"prompt_len={sum(len(m['content']) for m in messages)} "
                f"response_len={len(cached)}"
            )
            return cached

        # Use the OpenAI-compatible endpoint (Ollama supports both /api/chat
        # and /v1/chat/completions; the latter matches Fireworks' shape so we
        # don't need to change the response parsing).
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": False,
        }
        prompt_len = sum(len(m["content"]) for m in messages)
        print(f"[ollama] API CALL  model={self.model} prompt_len={prompt_len}")
        t0 = time.time()
        resp = self._session.post(
            f"{self.host}/v1/chat/completions",
            data=json.dumps(payload),
            timeout=self.timeout,
            headers={"Content-Type": "application/json"},
        )
        dt = time.time() - t0
        if resp.status_code != 200:
            raise RuntimeError(
                f"Ollama 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"[ollama] API OK   model={self.model} "
            f"response_len={len(text)} elapsed={dt:.2f}s cached=False"
        )
        return text

    # ------------------------------------------------------------------ #
    # High-level helpers (identical to FireworksClient)
    # ------------------------------------------------------------------ #
    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, MUST be the scientific name in lowercase '
            'with underscores, e.g. "nilaparvata_lugens" for brown planthopper, '
            '"spodoptera_frugiperda" for fall armyworm, "schistocerca_gregaria" '
            'for desert locust, "chilo_suppressalis" for striped stem borer, '
            '"myzus_persicae" for peach-potato aphid, '
            '"leptinotarsa_decemlineata" for Colorado potato beetle, '
            '"bemisia_tabaci" for tobacco whitefly), '
            '"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 (identical to FireworksClient)
    # ------------------------------------------------------------------ #
    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")