pen-stack 0.1.0__py3-none-any.whl

pen_stack/rag/llm.py

@@ -0,0 +1,215 @@
+"""Provider-agnostic LLM layer for PEN-STACK services (RAG, agent, monitor).
+
+Hybrid backend: a strong hosted model for reasoning/agent/Q&A (default NVIDIA Nemotron, OpenAI-compatible)
+with automatic fallback to a local, free, private model (Ollama). The single switch is `configs/llm.yaml`.
+
+This is strictly an orchestration/phrasing layer. Every quantitative claim and every citation still comes
+from the deterministic validated-tool path; the LLM never introduces a number, gene, or citation. The
+choice of model therefore does not affect scientific reproducibility - only the quality of orchestration
+and prose. If no provider is reachable, the callers fall back to the deterministic answer (LLM optional).
+
+Secrets: the API key is read from the env var named in `api_key_env`, then from the gitignored
+`api_key_file`. Keys are NEVER committed.
+"""
+from __future__ import annotations
+
+import json
+import os
+import urllib.request
+from pathlib import Path
+
+import yaml
+
+_CFG = Path(__file__).resolve().parents[2] / "configs" / "llm.yaml"
+_ROOT = Path(__file__).resolve().parents[2]
+
+_SYSTEM = ("You rephrase already-verified genome-writing facts into one clear paragraph for a wet-lab "
+           "scientist. Use ONLY the facts provided. Do NOT invent or alter any number, gene, or citation. "
+           "Do not give clinical advice. Keep it under 90 words.")
+
+
+def load_llm_config(path: str | Path = _CFG) -> dict:
+    return yaml.safe_load(Path(path).read_text(encoding="utf-8"))
+
+
+def _provider_cfg(cfg: dict, name: str) -> dict | None:
+    return (cfg.get("providers") or {}).get(name)
+
+
+def _resolve_key(pcfg: dict) -> str | None:
+    env = pcfg.get("api_key_env")
+    if env and os.environ.get(env):
+        return os.environ[env].strip()
+    f = pcfg.get("api_key_file")
+    if f:
+        p = Path(f)
+        if not p.is_absolute():
+            p = _ROOT / f
+        if p.exists():
+            return p.read_text(encoding="utf-8").strip()
+    return pcfg.get("api_key")
+
+
+def _norm_tool_calls(raw: list | None) -> list:
+    out = []
+    for c in raw or []:
+        fn = c.get("function", {})
+        args = fn.get("arguments", {})
+        if isinstance(args, str):
+            try:
+                args = json.loads(args or "{}")
+            except json.JSONDecodeError:
+                args = {}
+        out.append({"function": {"name": fn.get("name"), "arguments": args}})
+    return out
+
+
+def _chat_openai(pcfg: dict, messages: list, tools: list | None, temperature: float,
+                 timeout: int) -> dict | None:
+    """OpenAI-compatible /v1/chat/completions (NVIDIA NIM, OpenAI, vLLM, Ollama /v1)."""
+    base = pcfg["api_base"].rstrip("/")
+    key = _resolve_key(pcfg)
+    payload = {"model": pcfg["model"], "messages": messages, "temperature": temperature,
+               "max_tokens": int(pcfg.get("max_tokens", 1024))}
+    if tools:
+        payload["tools"] = tools
+        payload["tool_choice"] = "auto"
+    headers = {"Content-Type": "application/json"}
+    if key:
+        headers["Authorization"] = f"Bearer {key}"
+    req = urllib.request.Request(f"{base}/chat/completions", data=json.dumps(payload).encode(), headers=headers)
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        d = json.load(r)
+    msg = d["choices"][0]["message"]
+    return {"content": (msg.get("content") or "").strip(), "tool_calls": _norm_tool_calls(msg.get("tool_calls")),
+            "raw": msg, "style": "openai"}
+
+
+def _chat_ollama(pcfg: dict, messages: list, tools: list | None, temperature: float,
+                 timeout: int) -> dict | None:
+    """Ollama native /api/chat."""
+    base = pcfg["api_base"].rstrip("/")
+    payload = {"model": str(pcfg["model"]).split("/")[-1], "messages": messages, "stream": False,
+               "options": {"temperature": temperature}}
+    if tools:
+        payload["tools"] = tools
+    req = urllib.request.Request(f"{base}/api/chat", data=json.dumps(payload).encode(),
+                                 headers={"Content-Type": "application/json"})
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        d = json.load(r)
+    msg = d.get("message", {})
+    return {"content": (msg.get("content") or "").strip(), "tool_calls": _norm_tool_calls(msg.get("tool_calls")),
+            "raw": msg, "style": "ollama"}
+
+
+def _chat_anthropic(pcfg: dict, messages: list, tools: list | None, temperature: float,
+                    timeout: int) -> dict | None:
+    """Anthropic Messages API (/v1/messages). System prompts are a top-level field (not a message role),
+    and tool calls come back as `tool_use` content blocks; both are normalised to the common shape."""
+    base = pcfg.get("api_base", "https://api.anthropic.com/v1").rstrip("/")
+    key = _resolve_key(pcfg)
+    sys_txt = "\n".join(m["content"] for m in messages if m.get("role") == "system")
+    msgs = [{"role": m["role"], "content": m["content"]} for m in messages if m.get("role") != "system"]
+    payload = {"model": pcfg["model"], "max_tokens": int(pcfg.get("max_tokens", 1024)),
+               "temperature": temperature, "messages": msgs}
+    if sys_txt:
+        payload["system"] = sys_txt
+    if tools: # convert OpenAI-style function tools to Anthropic tool schema
+        payload["tools"] = [{"name": t["function"]["name"], "description": t["function"].get("description", ""),
+                             "input_schema": t["function"].get("parameters", {"type": "object", "properties": {}})}
+                            for t in tools]
+    headers = {"content-type": "application/json", "anthropic-version": "2023-06-01"}
+    if key:
+        headers["x-api-key"] = key
+    req = urllib.request.Request(f"{base}/messages", data=json.dumps(payload).encode(), headers=headers)
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        d = json.load(r)
+    content, tool_calls = "", []
+    for block in d.get("content", []):
+        if block.get("type") == "text":
+            content += block.get("text", "")
+        elif block.get("type") == "tool_use": # keep the tool_use id so the caller can build a tool_result
+            tool_calls.append({"id": block.get("id"),
+                               "function": {"name": block.get("name"), "arguments": block.get("input", {})}})
+    # `raw` must be a valid Anthropic assistant message (role + content blocks) so a multi-turn tool loop can
+    # append it verbatim; the full response object would be rejected on the next call.
+    raw = {"role": "assistant", "content": d.get("content", [])}
+    return {"content": content.strip(), "tool_calls": tool_calls, "raw": raw, "style": "anthropic"}
+
+
+def _call_provider(name: str, cfg: dict, messages: list, tools: list | None, timeout: int) -> dict | None:
+    pcfg = _provider_cfg(cfg, name)
+    if not pcfg:
+        return None
+    temp = float(cfg.get("temperature", 0.1))
+    style = pcfg.get("style", "openai")
+    try:
+        if style == "ollama":
+            return _chat_ollama(pcfg, messages, tools, temp, timeout)
+        if style == "anthropic":
+            return _chat_anthropic(pcfg, messages, tools, temp, timeout)
+        return _chat_openai(pcfg, messages, tools, temp, timeout)
+    except Exception: # noqa: BLE001 - any provider failure -> let the caller try the fallback
+        return None
+
+
+# Cooldown cache: once a provider fails (e.g. Ollama not installed on the laptop tier), skip it for
+# `health_ttl` seconds instead of re-attempting it on every call. This is what prevents the multi-minute
+# stalls when a configured provider is absent/slow - we pay one failed attempt, then bypass it.
+_COOLDOWN: dict[str, float] = {}
+
+
+def chat(messages: list, tools: list | None = None, cfg: dict | None = None,
+         timeout: int | None = None) -> dict | None:
+    """Provider-agnostic chat. Tries the active provider, then the configured fallback, skipping any
+    provider in cooldown (recently unreachable). Returns {content, tool_calls, provider} or None if every
+    provider fails (callers then degrade deterministically - the LLM is non-load-bearing)."""
+    import time
+    cfg = cfg or load_llm_config()
+    timeout = timeout if timeout is not None else int(cfg.get("call_timeout", 60))
+    ttl = float(cfg.get("health_ttl", 120))
+    order = [cfg.get("provider", "nvidia")]
+    fb = cfg.get("fallback")
+    if fb and fb not in order:
+        order.append(fb)
+    now = time.time()
+    tried_any = False
+    for name in order:
+        if _COOLDOWN.get(name, 0) > now: # provider recently failed -> skip without waiting
+            continue
+        tried_any = True
+        res = _call_provider(name, cfg, messages, tools, timeout)
+        if res is not None:
+            res["provider"] = name
+            _COOLDOWN.pop(name, None)
+            return res
+        _COOLDOWN[name] = now + ttl # mark unreachable; don't retry for ttl seconds
+    if not tried_any: # every provider in cooldown -> one cheap retry of the first
+        name = order[0]
+        res = _call_provider(name, cfg, messages, tools, min(timeout, int(cfg.get("health_timeout", 8))))
+        if res is not None:
+            res["provider"] = name
+            _COOLDOWN.pop(name, None)
+            return res
+    return None
+
+
+def active_provider(cfg: dict | None = None, timeout: int | None = None) -> str | None:
+    """Name of the first reachable provider (active, then fallback), or None. Uses the config `health_timeout`
+    by default so an absent provider is detected quickly (and then cooled down by chat())."""
+    cfg = cfg or load_llm_config()
+    timeout = timeout if timeout is not None else int(cfg.get("health_timeout", 8))
+    r = chat([{"role": "user", "content": "ok"}], cfg=cfg, timeout=timeout)
+    return r.get("provider") if r else None
+
+
+def available(cfg: dict | None = None, timeout: int = 30) -> bool:
+    return active_provider(cfg, timeout) is not None
+
+
+def phrase(facts: str, cfg: dict | None = None, timeout: int = 120) -> str | None:
+    """Rephrase grounded facts. Returns None on any failure (caller keeps the deterministic answer)."""
+    msgs = [{"role": "system", "content": _SYSTEM},
+            {"role": "user", "content": f"Facts:\n{facts}\n\nRephrase as one paragraph."}]
+    r = chat(msgs, cfg=cfg, timeout=timeout)
+    return (r.get("content") or None) if r else None

pen_stack/rag/qa.py

@@ -0,0 +1,105 @@
+"""Grounded, cited Q&A over the PEN-STACK platform.
+
+The front door for non-expert users. Contract (enforced by pen_stack.agent.guardrails):
+  * clinical-directive questions are refused;
+  * every *quantitative* claim is produced by a validated tool call (writability / atlas / cross-link),
+    never guessed by the LLM - the answer's ``provenance`` block names the tool;
+  * every factual claim carries a citation (DOIs from the curated atlas/WT-KB).
+
+An optional LLM (Ollama/Qwen via litellm) only *phrases* the grounded facts; it is never the source of a
+number or a citation. With no LLM available the deterministic tool+retrieval path still satisfies the
+contract - that is the whole point.
+"""
+from __future__ import annotations
+
+import re
+
+from pen_stack.agent.guardrails import DISCLAIMER, enforce_grounded, out_of_scope
+from pen_stack.rag.index import build_cards, retrieve
+
+_GENE_RE = re.compile(r"\b([A-Z][A-Z0-9]{2,9})\b") # crude gene-symbol cue (TRAC, CCR5, ...)
+_FAMILY_HINTS = {
+    "bridge": "bridge_IS110", "is110": "bridge_IS110", "iscro4": "bridge_IS110",
+    "seek": "seek_IS1111", "is1111": "seek_IS1111", "cast": "CAST_VK",
+    "integrase": "serine_integrase", "bxb1": "serine_integrase", "paste": "PE_integrase",
+    "prime": "PE_integrase", "cas9": "Cas9", "cas12a": "Cas12a", "tnpb": "TnpB_Fanzor",
+    "fanzor": "TnpB_Fanzor",
+}
+_WRITABLE_CUES = ("where", "writable", "safe harbour", "safe harbor", "insert", "insertion site", "locus")
+# Standing citation for tool-derived writability claims: the Writable Genome atlas
+# (TRIP durability supervision + clinical-CIS safety supervision).
+_WRITABILITY_CITATIONS = ["10.1016/j.cell.2013.07.018"] # Akhtar 2013 (TRIP) - durability supervision
+
+
+def _family_in(question: str) -> str | None:
+    q = question.lower()
+    for cue, fam in _FAMILY_HINTS.items():
+        if cue in q:
+            return fam
+    return None
+
+
+def answer(question: str, ct: str = "k562", use_llm: bool = False) -> dict:
+    refusal = out_of_scope(question)
+    if refusal:
+        return {"refused": True, "answer": refusal, "citations": [], "provenance": [],
+                "disclaimer": DISCLAIMER}
+
+    cards = build_cards()
+    retrieved = retrieve(question, cards, k=3)
+    citations = sorted({d for c in retrieved for d in c.citations})
+    provenance: list[dict] = []
+    parts: list[str] = []
+
+    # --- numeric route 1: "where can I write / writable loci for GENE" -> writability tool ---
+    if any(cue in question.lower() for cue in _WRITABLE_CUES):
+        genes = [g for g in _GENE_RE.findall(question) if g not in {"PEN", "STACK", "DNA", "RNA"}]
+        if genes:
+            try:
+                from pen_stack.atlas.crosslink import loci_for_gene
+                g = loci_for_gene(genes[0], ct)
+                if not g.empty:
+                    w = float(g["writability"].max())
+                    provenance.append({"tool": "crosslink.loci_for_gene",
+                                       "args": {"gene": genes[0], "ct": ct},
+                                       "result": {"max_writability": round(w, 3), "n_bins": int(len(g))}})
+                    parts.append(f"For {genes[0]} in {ct}, the most writable bin scores "
+                                 f"{w:.3f} (writability = safety x durability), across {len(g)} bins.")
+                    citations = sorted(set(citations) | set(_WRITABILITY_CITATIONS))
+            except FileNotFoundError:
+                parts.append("(Writable-genome atlas not loaded; numeric writability unavailable.)")
+
+    # --- numeric route 2: "which writer / tell me about FAMILY" -> atlas tool ---
+    fam = _family_in(question)
+    if fam:
+        import pandas as pd
+
+        from pen_stack.rag.index import _ATLAS
+        adf = pd.read_parquet(_ATLAS)
+        sub = adf[adf["family"] == fam]
+        rep = sub[sub["entry_kind"] == "curated_core"]
+        rep = rep.iloc[0] if len(rep) else sub.iloc[0]
+        provenance.append({"tool": "atlas.query", "args": {"family": fam},
+                           "result": {"n_systems": int(len(sub)),
+                                      "reachability_tier": rep.get("reachability_tier"),
+                                      "deliv_class": rep.get("deliv_class")}})
+        parts.append(f"{fam}: {len(sub):,} catalogued systems; reachability {rep.get('reachability_tier')}; "
+                     f"deliverability {rep.get('deliv_class')}; representative {rep['representative_system']}.")
+
+    # --- factual route: retrieval-grounded summary (always cited) ---
+    if retrieved:
+        parts.append("Relevant atlas facts: " + " | ".join(c.text for c in retrieved[:2]))
+
+    if not parts:
+        parts.append("No grounded match in the atlas. Try naming a writer family or a target gene.")
+
+    out = {"refused": False, "answer": " ".join(parts), "citations": citations,
+           "provenance": provenance, "disclaimer": DISCLAIMER}
+    out = enforce_grounded(out)
+    # optional LLM phrasing - presentation only; numbers/citations stay tool-derived (additive field)
+    if use_llm:
+        from pen_stack.rag.llm import phrase
+        phrased = phrase(out["answer"])
+        if phrased:
+            out["answer_phrased"] = phrased
+    return out

pen_stack/rag/retrieve.py

@@ -0,0 +1,48 @@
+"""RAG retriever.
+
+Loads the committed corpus + its committed embedding matrix and returns the top-k chunks for a query by EXACT
+cosine over the embeddings (semantic, via the live query embedder), with a deterministic token-Jaccard fallback
+when no embedder is reachable. Exact cosine over a committed matrix is the reproducible equivalent of a vector
+index at this corpus scale (hundreds of chunks) - FAISS is unnecessary.
+"""
+from __future__ import annotations
+
+from functools import lru_cache
+
+import numpy as np
+
+from pen_stack.rag.corpus import emb_path, load_corpus
+from pen_stack.rag.embed import embed_query, lexical_scores
+
+
+@lru_cache(maxsize=1)
+def _load():
+    df = load_corpus()
+    p = emb_path()
+    emb = None
+    if p.exists():
+        m = np.load(p).astype("float32")
+        if len(m) == len(df):  # guard against a stale embedding matrix
+            emb = m
+    return df, emb
+
+
+def retrieve(query: str, k: int = 4) -> dict:
+    df, emb = _load()
+    texts = df["text"].tolist()
+    qv = embed_query(query) if emb is not None else None  # LRU-cached query embedding
+    if qv is not None and emb is not None:
+        scores = emb @ qv  # both L2-normalised -> dot product is cosine
+        method = "semantic (nomic-embed-text)"
+    else:
+        scores = lexical_scores(query, texts)
+        method = "lexical (jaccard fallback)"
+    order = [int(i) for i in np.argsort(-scores)[:k]]
+    hits = []
+    for i in order:
+        r = df.iloc[i]
+        hits.append({"chunk_id": str(r.chunk_id), "text": str(r.text), "source_id": str(r.source_id),
+                     "doi": str(r.doi) if r.doi else "", "type": str(r.type),
+                     "scope_status": str(r.scope_status), "score": float(scores[i])})
+    return {"query": query, "method": method, "hits": hits,
+            "top_score": float(scores[order[0]]) if order else 0.0, "n_corpus": len(df)}

pen_stack/rules/__init__.py

@@ -0,0 +1,9 @@
+"""Machine-readable rules engine for genome writing. Rules are versioned data in
+configs/rules/*.yaml; evaluators (rules/evaluators.py) delegate to the existing validated functions; the
+solver (rules/solver.py) returns legality + named reasons. Import evaluators so they register on load."""
+from pen_stack.rules import evaluators as _evaluators # noqa: F401 (registers evaluators)
+from pen_stack.rules.loader import load_ruleset
+from pen_stack.rules.schema import Design, Rule, Ruleset
+from pen_stack.rules.solver import evaluate, is_legal, legality_report
+
+__all__ = ["load_ruleset", "Design", "Rule", "Ruleset", "evaluate", "is_legal", "legality_report"]