aiva-agent 0.2.1__py3-none-any.whl

aiva_agent/__init__.py

@@ -0,0 +1,8 @@
+import os
+
+os.environ.setdefault("OPENAI_AGENTS_DISABLE_TRACING", "1")
+
+from .agent import run_once
+from .notebook import aiva_agent, reset_session
+
+__all__ = ["run_once", "aiva_agent", "reset_session"]

aiva_agent/agent.py

@@ -0,0 +1,255 @@
+"""Agent runtime: builds an OpenAI Agents SDK Agent from --disable, runs it once.
+
+Multi-provider works through any OpenAI-compatible endpoint:
+- User passes --model, --base-url, --api-key (or sets LLM_BASE_URL / LLM_API_KEY)
+- We construct AsyncOpenAI(base_url=, api_key=) and OpenAIChatCompletionsModel(model=, ...)
+- Pass that as model= to Agent(...)
+
+No PROVIDERS dict, no provider-prefix parsing, no LiteLLM. The user already knows
+their provider's URL and has their key; we don't need to enumerate.
+"""
+
+import asyncio
+import os
+from pathlib import Path
+from typing import Iterable
+
+from agents import Agent, OpenAIChatCompletionsModel, Runner, SQLiteSession, function_tool
+from openai import AsyncOpenAI
+
+from .classification_prompts import ACMG_AMP_SYSTEM_PROMPT
+from .prompts import get_main_system_prompt
+from .tools import (
+    build_annotate_tool,
+    build_literature_tool,
+    build_phen2gene_tool,
+    build_trials_tool,
+    build_vcf_tools,
+    build_web_tool,
+)
+
+DEFAULT_MODEL = "gpt-5.5"
+DEFAULT_MAX_TURNS = 25
+
+ALL_TOOLS = ("vcf", "annotate", "literature", "trials", "phen2gene", "web", "classify")
+
+
+def _resolve_max_turns() -> int:
+    raw = os.environ.get("AIVA_MAX_TURNS", "").strip()
+    if not raw:
+        return DEFAULT_MAX_TURNS
+    try:
+        n = int(raw)
+    except ValueError:
+        return DEFAULT_MAX_TURNS
+    return n if n > 0 else DEFAULT_MAX_TURNS
+
+
+def _resolve_disabled(disabled_tools: Iterable[str] | None) -> list[str]:
+    """Return the active tool list = ALL_TOOLS minus the disabled set.
+
+    Default (None / empty) → every tool is on. Unknown names are ignored
+    (the CLI already validates), preserving the lenient behavior of the
+    previous `_resolve_enabled`.
+    """
+    if not disabled_tools:
+        return list(ALL_TOOLS)
+    drop = {name.strip() for name in disabled_tools if name and name.strip()}
+    return [t for t in ALL_TOOLS if t not in drop]
+
+
+def build_model(
+    model: str,
+    base_url: str | None = None,
+    api_key: str | None = None,
+) -> OpenAIChatCompletionsModel:
+    """Construct an OpenAIChatCompletionsModel pointing at any OpenAI-compatible endpoint.
+
+    Resolution order for base_url and api_key:
+    1. The argument passed in (CLI flag).
+    2. LLM_BASE_URL / LLM_API_KEY env vars.
+    """
+    resolved_base_url = base_url or os.environ.get("LLM_BASE_URL")
+    resolved_api_key = api_key or os.environ.get("LLM_API_KEY")
+    client = AsyncOpenAI(base_url=resolved_base_url, api_key=resolved_api_key)
+    return OpenAIChatCompletionsModel(model=model, openai_client=client)
+
+
+def build_classifier_tool(
+    classifier_tools: list,
+    model: str,
+    base_url: str | None,
+    api_key: str | None,
+):
+    """Wrap the variant-classifier sub-agent as a @function_tool callable from the parent.
+
+    Mirrors genomiq's SubAgentRunner pattern: build a fresh Agent with the classifier's
+    own tool palette, run it via Runner.run, return the JSON output.
+    """
+    sub_model = build_model(model, base_url, api_key)
+    sub_agent = Agent(
+        name="variant-classifier",
+        instructions=ACMG_AMP_SYSTEM_PROMPT,
+        model=sub_model,
+        tools=classifier_tools,
+    )
+
+    # `variant_data` is a heterogeneous dict (rsid OR hgvs OR chrom/pos/ref/alt) — strict
+    # JSON Schema doesn't model that cleanly without a TypedDict per shape, so we relax
+    # strict mode for this single tool. The shape contract is documented in the docstring
+    # and enforced by the sub-agent's system prompt.
+    @function_tool(strict_mode=False)
+    async def classify_variant(
+        variant_data: dict,
+        classification_type: str,
+        assembly: str,
+        phenotype_terms: str = "",
+        description: str = "",
+        additional_context: str = "",
+    ) -> str:
+        """Classify a genomic variant using ACMG/AMP 2015 guidelines (germline) or
+        AMP/ASCO/CAP 2017 guidelines (somatic). Runs as a sub-agent that gathers
+        annotations, literature, and trial evidence on its own, returning a JSON
+        classification with criteria_met, evidence_summary, and sources.
+
+        Args:
+            variant_data: One of {"rsid": "rs..."} or {"hgvs": "..."} or
+                {"chrom": "...", "pos": ..., "ref": "...", "alt": "..."}.
+            classification_type: 'acmg' for germline or 'amp' for somatic/cancer.
+            assembly: 'GRCh37' or 'GRCh38'.
+            phenotype_terms: e.g. 'melanoma', 'hereditary breast cancer'.
+            description: Sample/patient context (optional).
+            additional_context: Verified clinical context (de novo status, family history,
+                zygosity, segregation). Apply criteria strictly to what's stated here.
+        """
+        formatted = (
+            "Classify this variant.\n\n"
+            f"- variant_data: {variant_data}\n"
+            f"- classification_type: {classification_type}\n"
+            f"- assembly: {assembly}\n"
+            f"- phenotype_terms: {phenotype_terms}\n"
+            f"- description: {description}\n"
+            f"- additional_context: {additional_context}\n"
+        )
+        result = await Runner.run(sub_agent, formatted, max_turns=_resolve_max_turns())
+        return result.final_output
+
+    return classify_variant
+
+
+def build_tools_and_cleanup(
+    enabled: list[str],
+    vcf_specs: list[tuple[str, Path]] | None,
+    model: str,
+    base_url: str | None,
+    api_key: str | None,
+) -> tuple[list, list]:
+    """Construct the tool list for an agent run plus any cleanup callables.
+
+    Raises ValueError if 'vcf' is in `enabled` but `vcf_specs` is empty. The CLI
+    auto-disables vcf with a warning before reaching this function; the guard
+    is the contract for direct library callers.
+
+    `vcf_specs` is a list of `(alias, path)` pairs; each becomes a named DuckDB
+    view. Single-file invocations should pass `[("vcf", path)]`.
+
+    classify is special: it needs annotate / literature / trials available to its
+    sub-agent. We build those tools once and share the instances between the parent's
+    palette and the classifier's palette.
+    """
+    tools: list = []
+    cleanups: list = []
+
+    annotate_tool = literature_tool = trials_tool = None
+
+    if "vcf" in enabled:
+        if not vcf_specs:
+            raise ValueError("--vcf is required when 'vcf' is enabled")
+        vcf_tools, vcf_cleanup = build_vcf_tools(vcf_specs)
+        tools.extend(vcf_tools)
+        cleanups.append(vcf_cleanup)
+
+    if "annotate" in enabled or "classify" in enabled:
+        annotate_tool = build_annotate_tool()
+        if "annotate" in enabled:
+            tools.append(annotate_tool)
+
+    if "literature" in enabled or "classify" in enabled:
+        literature_tool, literature_cleanup = build_literature_tool()
+        cleanups.append(literature_cleanup)
+        if "literature" in enabled:
+            tools.append(literature_tool)
+
+    if "trials" in enabled or "classify" in enabled:
+        trials_tool, trials_cleanup = build_trials_tool()
+        cleanups.append(trials_cleanup)
+        if "trials" in enabled:
+            tools.append(trials_tool)
+
+    if "phen2gene" in enabled:
+        phen2gene_tool, phen2gene_cleanup = build_phen2gene_tool()
+        cleanups.append(phen2gene_cleanup)
+        tools.append(phen2gene_tool)
+
+    if "web" in enabled:
+        tools.append(build_web_tool())
+
+    if "classify" in enabled:
+        classifier_tools = [t for t in (annotate_tool, literature_tool, trials_tool) if t is not None]
+        tools.append(
+            build_classifier_tool(
+                classifier_tools=classifier_tools,
+                model=model,
+                base_url=base_url,
+                api_key=api_key,
+            )
+        )
+
+    return tools, cleanups
+
+
+async def run_once(
+    prompt: str,
+    *,
+    vcf_specs: list[tuple[str, Path]] | None = None,
+    disabled_tools: Iterable[str] | None = None,
+    model: str = DEFAULT_MODEL,
+    base_url: str | None = None,
+    api_key: str | None = None,
+    session_id: str | None = None,
+) -> str:
+    enabled = _resolve_disabled(disabled_tools)
+
+    tools, cleanups = build_tools_and_cleanup(
+        enabled=enabled,
+        vcf_specs=vcf_specs,
+        model=model,
+        base_url=base_url,
+        api_key=api_key,
+    )
+
+    session = None
+    if session_id:
+        db_path = Path.home() / ".aiva" / "sessions.db"
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+        session = SQLiteSession(session_id, str(db_path))
+
+    try:
+        agent = Agent(
+            name="aiva",
+            instructions=get_main_system_prompt(enabled, vcf_specs=vcf_specs),
+            model=build_model(model, base_url, api_key),
+            tools=tools,
+        )
+        result = await Runner.run(
+            agent, prompt, max_turns=_resolve_max_turns(), session=session
+        )
+        return result.final_output
+    finally:
+        for cleanup in cleanups:
+            try:
+                result = cleanup()
+                if asyncio.iscoroutine(result):
+                    await result
+            except Exception:
+                pass

aiva_agent/classification_prompts.py

@@ -0,0 +1,143 @@
+"""Static system prompt for the variant-classifier sub-agent.
+
+The genomiq classification flow built one prompt per call that interleaved the
+variant data, criteria, and instructions. With the OpenAI Agents SDK's
+function_tool sub-agent pattern (a `@function_tool` that internally instantiates
+an `Agent` and calls `Runner.run`) we use a single, static system prompt — the
+variant data is passed in via the first user-turn message from the parent agent.
+"""
+
+ACMG_AMP_SYSTEM_PROMPT = """You are a clinical genetics expert specializing in variant interpretation
+using ACMG/AMP 2015 guidelines (germline) and AMP/ASCO/CAP 2017 guidelines (somatic/cancer).
+
+# Input contract
+
+Your invoking prompt will provide these fields (the parent agent is instructed to
+include them). Parse them out before doing any work:
+
+- **variant_data** (REQUIRED): one of
+    • `{"rsid": "rs113488022"}`
+    • `{"hgvs": "chr7:g.140453136A>T"}` or `{"hgvs": "NM_004333.6:c.1799T>A"}`
+    • `{"chrom": "7", "pos": 140453136, "ref": "A", "alt": "T"}`
+- **classification_type** (REQUIRED): `"acmg"` (germline) or `"amp"` (somatic/cancer)
+- **assembly** (REQUIRED): `"GRCh37"` or `"GRCh38"`
+- **phenotype_terms**: e.g. "melanoma", "hereditary breast cancer". Use it to scope
+  literature/trial searches and to evaluate phenotype-specific criteria (PP4, etc.).
+- **description**: free-form sample/patient context. Use it for clinical framing.
+- **additional_context**: verified clinical context (de novo status, family history,
+  zygosity, segregation). Apply ACMG criteria based STRICTLY on what is explicitly
+  stated here — do not extrapolate. PS2 only with confirmed de novo; PM6 with
+  assumed de novo; neither if no de novo info.
+
+If `variant_data`, `classification_type`, or `assembly` is missing or ambiguous,
+return `{"error": "missing required field: <field>", "expected": "..."}` instead of
+guessing.
+
+# Available tools
+
+- `annotate_variant` (myvariant / vep / civic) — gather variant annotations.
+  Pass `include_commercial=true` if you need CADD or dbNSFP scores for PP3/BP4.
+- `search_biomedical_literature` (PubTator3) — supporting literature.
+- `search_clinical_trials` (ClinicalTrials.gov) — therapeutic context for somatic.
+
+You have a strict ~120-second budget. Prioritize: (1) ONE annotate_variant call to
+gather all annotations, (2) a few targeted literature searches, (3) clinical trials
+only if relevant for somatic classification. Do NOT call annotate_variant multiple
+times — one call returns all data; only retry on error or rate-limit.
+
+# ACMG/AMP 2015 (Germline) — `acmg`
+
+**Key Interpretation Rules:**
+- DO NOT defer to ClinVar's existing classification — your classification must be based
+  on YOUR evaluation of all criteria, not ClinVar's conclusion.
+- Missing gnomAD = RARE variant (supports PM2 criterion).
+- In-silico priority: CADD and DITTO are primary. CADD≥15 or DITTO≥0.5 strongly
+  supports deleterious effect. SIFT, PolyPhen, MutationTaster, REVEL are secondary.
+- Avoid plain "VUS" — use "VUS (leaning pathogenic)" or "VUS (leaning benign)".
+
+**Critical Validation Rules:**
+- PP3/BP4 (Computational predictions): ONLY apply with actual scores (CADD, DITTO,
+  REVEL, SIFT, PolyPhen). Do NOT apply based on domain location alone.
+- PM1 (Functional domain): Requires BOTH functional domain evidence AND absence of
+  benign variation at that location.
+- PP2 (Missense constraint): Requires quantitative metrics (gnomAD mis_z, RVIS).
+- PM5: Requires that other same-position missense variants are TRULY pathogenic in
+  ClinVar/literature, not just reported.
+- PS2/PM6 (De novo): PS2 if confirmed; PM6 if assumed; otherwise neither.
+- PP4 (Phenotype specificity): Requires pathognomonic features, not general symptoms.
+- Each criterion is fully met or not_applicable — no partial application.
+- PP5/BP6 are removed per ClinGen 2018; do not apply.
+
+**Pathogenic Criteria (apply only when met):**
+PVS1 (LOF in gene where LOF causes disease — strength per Abou Tayoun 2018),
+PS1 (same AA change as established pathogenic), PS2 (confirmed de novo),
+PS3 (well-established functional studies), PS4 (case prevalence OR>5),
+PM1 (hot spot + no benign variation), PM2_Supporting (rare in gnomAD),
+PM3 (in trans with pathogenic, recessive), PM4 (length change in-frame),
+PM5 (novel missense at established pathogenic position), PM6 (assumed de novo),
+PP1 (co-segregation), PP2 (constrained gene with quantitative metrics),
+PP3 (computational evidence supporting deleterious),
+PP4 (pathognomonic phenotype).
+
+**Benign Criteria (apply only when met):**
+BA1 (>5% gnomAD), BS1 (frequency > expected), BS2 (healthy adult observed),
+BS3 (functional studies show no damage), BS4 (lack of segregation),
+BP1 (missense in truncating-only gene), BP2 (in trans/cis with pathogenic),
+BP3 (in-frame del/ins in repeat), BP4 (computational benign predictions),
+BP5 (alternate molecular cause), BP7 (synonymous, no splice impact).
+
+**Combining Rules:**
+- Pathogenic: 1 VS + 1 S; OR 2 S; OR 1 S + 3 M; OR 1 S + 2 M + 2 P; OR 1 S + 1 M + 4 P
+- Likely Pathogenic: 1 VS + 1 M; OR 1 S + 1-2 M; OR 1 S + 2 P; OR 3 M; OR 2 M + 2 P; OR 1 M + 4 P
+- Benign: 1 stand-alone (BA1) OR 2 strong
+- Likely Benign: 1 strong + 1 supporting OR 2 supporting
+- VUS: criteria don't meet thresholds (use "leaning pathogenic" / "leaning benign")
+
+**ACMG output (JSON ONLY):**
+{
+  "classification": "Pathogenic | Likely Pathogenic | VUS (leaning pathogenic) | VUS (leaning benign) | Likely Benign | Benign",
+  "confidence": "high | medium | low",
+  "acmg_score": <int>,
+  "criteria_met": ["PM2_Supporting", "PP3", ...],
+  "evidence_summary": {"PM2_Supporting": "Absent in gnomAD v4.0 ...", "PP3": "CADD=28.5 ...", "notable": "..."},
+  "classification_rationale": "...",
+  "sources": ["ClinVar:12215", "PubMed:12345678", "gnomAD v4.0"]
+}
+
+# AMP/ASCO/CAP 2017 (Somatic) — `amp`
+
+**Tier Framework:**
+- Tier I-A: FDA-approved therapy for this tumor type or in NCCN/ASCO/CAP guidelines.
+- Tier I-B: Well-powered studies with expert consensus on clinical utility.
+- Tier II-C: FDA-approved for different tumor type, or trial eligibility.
+- Tier II-D: Preclinical / case reports / biological rationale.
+- Tier III: Insufficient evidence.
+- Tier IV: Benign / common germline / no functional impact.
+
+**Evidence categories:** Therapeutic (TA/TB/TC/TD), Diagnostic (DA/DB/DC/DD),
+Prognostic (PA/PB/PC/PD). Benign: BN1 (>1% gnomAD), BN2 (no functional impact —
+requires actual scores), BN3 (germline polymorphism), BN4 (synonymous, no splice).
+
+**Critical Validation Rule:** BN2 only when actual scores show CADD<15 AND DITTO<0.5,
+SIFT=tolerated, PolyPhen=benign. Do NOT apply BN2 from domain inference.
+
+**AMP output (JSON ONLY):**
+{
+  "classification": "Tier I-A | Tier I-B | Tier II-C | Tier II-D | Tier III | Tier IV",
+  "confidence": "high | medium | low",
+  "criteria_met": ["TA", "PB", ...],
+  "evidence_summary": {"TA": "FDA-approved drug X for this tumor type ...", "PB": "..."},
+  "therapeutic_summary": "...",
+  "diagnostic_summary": "...",
+  "prognostic_summary": "...",
+  "sources": ["OncoKB", "ClinVar:12345", "PubMed:...", "NCCN Guidelines"]
+}
+
+# General
+
+- Internally evaluate ALL criteria but only output the criteria that are actually met.
+- Return ONLY valid JSON — no prose before or after, no markdown fences.
+- If evidence is missing, mark a criterion not_applicable with a brief explanation.
+- DO NOT use web_search for variant-specific data (gnomAD, CADD, ClinVar) — those
+  come from annotate_variant.
+"""