celltype-cli 0.1.0__py3-none-any.whl

ct/agent/system_prompt.py

@@ -0,0 +1,186 @@
+"""
+Unified system prompt builder for the Claude Agent SDK runner.
+
+Merges the domain knowledge primer, workflow guides, bioinformatics code-gen
+hints, synthesis rules, tool catalog, and dynamic data context into a single
+system prompt that Claude uses throughout the agentic loop.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+logger = logging.getLogger("ct.system_prompt")
+
+
+# ---------------------------------------------------------------------------
+# Identity / role preamble
+# ---------------------------------------------------------------------------
+
+_IDENTITY = """\
+You are **celltype-cli**, an autonomous drug discovery research agent.
+
+You have access to 190+ domain tools covering target discovery, chemistry,
+expression, viability, safety, clinical development, omics, genomics, literature,
+and more — plus a persistent Python sandbox (``run_python``) for custom analyses.
+
+Your job: take a research question and answer it **completely**, using the right
+tools and code, self-correcting as you go, and producing a publication-quality
+synthesis at the end.
+
+## Operating Mode
+- You are in an agentic loop: call tools, see results, call more tools, then
+  write your final answer as plain text (no tool call).
+- Think step-by-step. Use tools to gather evidence, then synthesize.
+- If a tool fails or returns unhelpful data, try a different approach or use
+  your own knowledge to fill gaps.
+- For data analysis questions, use ``run_python`` to load data, explore it,
+  and compute the answer. Variables persist between calls.
+"""
+
+
+# ---------------------------------------------------------------------------
+# Synthesis instructions (injected at the end)
+# ---------------------------------------------------------------------------
+
+_SYNTHESIS_INSTRUCTIONS = """\
+
+## When You Are Ready to Answer
+
+Write your final answer as a direct text response (do NOT call any more tools).
+Your answer should be:
+
+1. **Complete**: Address every part of the question. Decompose the question into
+   sub-parts and make sure each is answered with specifics.
+2. **Accurate**: Use tool results as primary evidence. Supplement with your
+   domain knowledge. Never fabricate data.
+3. **Data-rich**: Include specific gene names, cell lines, p-values, effect
+   sizes, IC50 values, trial names, mutation positions, etc.
+4. **Mechanistic**: Explain the biological *why*, not just the *what*.
+5. **Actionable**: End with 3-5 specific experimental next steps (named assays,
+   cell lines, concentrations, readouts).
+
+BANNED PHRASES — never write these:
+- "cannot be answered with the data retrieved"
+- "failed to retrieve" / "failed to identify"
+- "insufficient data" / "insufficient evidence"
+- "No results were obtained"
+If tools failed, pivot to answering from your knowledge instead.
+"""
+
+
+# ---------------------------------------------------------------------------
+# Builder
+# ---------------------------------------------------------------------------
+
+def build_system_prompt(
+    session,
+    tool_names: list[str] | None = None,
+    data_context: str | None = None,
+    history: str | None = None,
+) -> str:
+    """Build the unified system prompt for the Agent SDK runner.
+
+    Args:
+        session: Active ct Session.
+        tool_names: Names of tools available in the MCP server (for reference).
+        data_context: Free-text description of available data files / directories.
+        history: Prior conversation turns (for interactive multi-turn sessions).
+
+    Returns:
+        The complete system prompt string.
+    """
+    parts: list[str] = []
+
+    # 1. Identity
+    parts.append(_IDENTITY)
+
+    # 2. Tool catalog (concise reference — full descriptions are in MCP tool defs)
+    # NOTE: The Agent SDK exposes tool names+descriptions+schemas via MCP natively.
+    # We only include a brief orientation here, NOT the full tool_descriptions_for_llm()
+    # which would blow up the system prompt to 155K chars and hit ARG_MAX limits.
+    if tool_names:
+        parts.append(f"\n## Available Tools ({len(tool_names)} total)\n")
+        parts.append(
+            "You have access to all tools via MCP. Key tools:\n"
+            "- **run_python**: Execute Python code in a sandbox (pd, np, plt, scipy, sklearn, pysam, gseapy, pydeseq2, BioPython). Variables persist between calls.\n"
+            "- **run_r**: Execute R code directly. Prefer run_r over run_python for: natural splines (ns()), wilcox.test(), p.adjust(), fisher.test(), lm()/predict(), organism-specific KEGG ORA (use KEGGREST package: keggList, keggLink, keggGet for any organism code), and any analysis where R is the reference implementation. R and Python give DIFFERENT results for splines, multiple testing correction, and nonparametric tests — when the expected answer was computed in R, use R.\n"
+            "- **literature.pubmed_search**, **literature.chembl_query**, **literature.openalex_search**: Literature/DB search\n"
+            "- **data_api.opentargets_search**, **data_api.depmap_search**, **data_api.uniprot_lookup**: Platform APIs\n"
+            "- **omics.geo_search**, **omics.geo_fetch**, **omics.deseq2**, **omics.dataset_info**: Omics data discovery + analysis\n"
+            "- **expression.pathway_enrichment**, **expression.l1000_similarity**: Expression/pathway tools\n"
+            "- **viability.dose_response**, **clinical.indication_map**, **clinical.trial_search**: Viability/clinical\n"
+            "- **chemistry.descriptors**, **chemistry.sar_analyze**, **chemistry.pubchem_lookup**: Chemistry tools\n"
+            "- **target.coessentiality**, **genomics.gwas_lookup**, **protein.function_predict**: Target/genomics tools\n"
+            "- **safety.classify**, **safety.admet_predict**: Safety/ADMET tools\n"
+            "\nFor data analysis questions, prefer **run_python** — it's the most powerful tool.\n"
+            "For drug discovery questions, combine domain tools with your knowledge.\n"
+        )
+
+    # 3. Workflow guides (compact — key sequences for common tasks)
+    try:
+        from ct.agent.workflows import format_workflows_for_llm
+        workflows = format_workflows_for_llm()
+        if workflows:
+            parts.append(workflows)
+    except Exception as e:
+        logger.warning("Could not load workflows: %s", e)
+
+    # 4. Domain knowledge primer (CRITICAL for drug discovery accuracy)
+    # NOTE: The KNOWLEDGE_PRIMER contains both tool orientation AND domain facts.
+    # The tool orientation section overlaps with MCP tool descriptions but the
+    # cross-disciplinary thinking patterns and domain-specific accuracy anchors
+    # are essential. Include in full.
+    try:
+        from ct.agent.knowledge import KNOWLEDGE_PRIMER
+        parts.append("\n" + KNOWLEDGE_PRIMER)
+    except Exception as e:
+        logger.warning("Could not load knowledge primer: %s", e)
+
+    # 6. Bioinformatics code-gen hints (CRITICAL for BixBench performance)
+    try:
+        from ct.tools.code import BIOINFORMATICS_CODE_GEN_PROMPT, AGENTIC_CODE_ADDENDUM
+        # Strip the template placeholders and include the raw hints
+        hints = BIOINFORMATICS_CODE_GEN_PROMPT
+        # Remove the {namespace_description} and {data_files_description} placeholders
+        hints = hints.replace("{namespace_description}", "(see run_python tool description)")
+        hints = hints.replace("{data_files_description}", "(see data context below)")
+        parts.append("\n## Bioinformatics Code Generation Guide\n")
+        parts.append(
+            "When using ``run_python`` for bioinformatics analysis, follow these "
+            "patterns and guidelines:\n"
+        )
+        parts.append(hints)
+        parts.append(AGENTIC_CODE_ADDENDUM)
+    except Exception as e:
+        logger.warning("Could not load code-gen hints: %s", e)
+
+    # 7. Synthesis rules
+    try:
+        from ct.agent.knowledge import SYNTHESIZER_PRIMER
+        parts.append("\n## Synthesis Guidelines\n")
+        parts.append(SYNTHESIZER_PRIMER)
+    except Exception as e:
+        logger.warning("Could not load synthesizer primer: %s", e)
+
+    # 8. Synthesis instructions
+    parts.append(_SYNTHESIS_INSTRUCTIONS)
+
+    # 9. Dynamic data context
+    if data_context:
+        parts.append("\n## Data Context\n")
+        parts.append(data_context)
+
+    # 10. Session history (for multi-turn interactive mode)
+    if history:
+        parts.append("\n## Prior Conversation\n")
+        parts.append(history)
+
+    prompt = "\n".join(parts)
+    logger.info(
+        "Built system prompt: %d chars, %d sections",
+        len(prompt),
+        len(parts),
+    )
+    return prompt

ct/agent/trace_store.py

@@ -0,0 +1,228 @@
+"""
+Trace store: full-fidelity capture of agent message streams.
+
+Captures the complete sequence of text blocks, tool calls, and tool results
+from each query execution. Used by the notebook exporter and other
+downstream consumers that need richer data than the compact Trajectory.
+
+Operates independently of Trajectory — Trajectory stores compact turn
+summaries for LLM context injection; TraceStore captures the full stream
+for export/replay.
+"""
+
+import base64
+import json
+import logging
+import mimetypes
+import time
+from pathlib import Path
+
+logger = logging.getLogger("ct.trace_store")
+
+# Marker used by MCP handlers to embed structured metadata in tool results
+TRACE_META_MARKER = "\n__CT_TRACE_META__\n"
+
+# Maximum file size for base64 embedding (10 MB)
+_MAX_EMBED_BYTES = 10 * 1024 * 1024
+
+
+def _sessions_dir() -> Path:
+    """Return the sessions directory, creating it if needed."""
+    d = Path.home() / ".ct" / "sessions"
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def parse_trace_meta(result_text: str) -> dict | None:
+    """Extract __CT_TRACE_META__ JSON from a tool result string.
+
+    Returns the parsed dict, or None if no marker found.
+    """
+    if TRACE_META_MARKER not in result_text:
+        return None
+    try:
+        _, meta_json = result_text.split(TRACE_META_MARKER, 1)
+        return json.loads(meta_json.strip())
+    except (ValueError, json.JSONDecodeError) as e:
+        logger.warning("Failed to parse trace meta: %s", e)
+        return None
+
+
+def _embed_plots(event: dict) -> None:
+    """Read plot files and add base64-encoded data to the event in-place."""
+    plots = event.get("plots", [])
+    if not plots:
+        return
+
+    embedded = []
+    for plot_path_str in plots:
+        plot_path = Path(plot_path_str)
+        if not plot_path.exists():
+            logger.warning("Plot file missing at capture time: %s", plot_path)
+            continue
+        if plot_path.stat().st_size > _MAX_EMBED_BYTES:
+            logger.warning(
+                "Plot file too large for embedding (%d bytes): %s",
+                plot_path.stat().st_size,
+                plot_path,
+            )
+            continue
+
+        mime, _ = mimetypes.guess_type(str(plot_path))
+        if mime is None:
+            suffix_map = {
+                ".png": "image/png",
+                ".svg": "image/svg+xml",
+                ".jpg": "image/jpeg",
+                ".jpeg": "image/jpeg",
+                ".pdf": "application/pdf",
+            }
+            mime = suffix_map.get(plot_path.suffix.lower(), "application/octet-stream")
+
+        try:
+            data = base64.b64encode(plot_path.read_bytes()).decode("ascii")
+            embedded.append({
+                "filename": plot_path.name,
+                "mime": mime,
+                "data": data,
+            })
+        except OSError as e:
+            logger.warning("Failed to read plot file %s: %s", plot_path, e)
+
+    if embedded:
+        event["plots_base64"] = embedded
+
+
+class TraceStore:
+    """Captures and persists full agent message stream traces.
+
+    Usage::
+
+        store = TraceStore(session_id="abc-123")
+        events = []
+
+        # ... pass events list to process_messages() ...
+
+        store.add_events(events, query="my question", model="claude-sonnet-4-5")
+        store.flush()  # persist to ~/.ct/sessions/abc-123.trace.jsonl
+    """
+
+    def __init__(self, session_id: str):
+        self.session_id = session_id
+        self._events: list[dict] = []
+        self._path = _sessions_dir() / f"{session_id}.trace.jsonl"
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    @property
+    def events(self) -> list[dict]:
+        return self._events
+
+    def add_event(self, event: dict) -> None:
+        """Add a single trace event."""
+        if "timestamp" not in event:
+            event["timestamp"] = time.time()
+        self._events.append(event)
+
+    def add_events(
+        self,
+        events: list[dict],
+        query: str = "",
+        model: str = "",
+        duration_s: float = 0.0,
+        cost_usd: float = 0.0,
+    ) -> None:
+        """Wrap a list of trace events with query_start/query_end and add them.
+
+        Also performs eager base64 embedding of plots.
+        """
+        now = time.time()
+
+        # query_start
+        self._events.append({
+            "type": "query_start",
+            "session_id": self.session_id,
+            "query": query,
+            "model": model,
+            "timestamp": now,
+        })
+
+        # Add all events (with plot embedding)
+        for event in events:
+            if event.get("type") == "tool_result" and event.get("plots"):
+                _embed_plots(event)
+            self._events.append(event)
+
+        # query_end
+        self._events.append({
+            "type": "query_end",
+            "duration_s": duration_s,
+            "cost_usd": cost_usd,
+            "timestamp": time.time(),
+        })
+
+    def flush(self, path: Path | None = None) -> Path:
+        """Persist all events to a JSONL file. Appends if file exists."""
+        out = path or self._path
+        out.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(out, "a", encoding="utf-8") as f:
+            for event in self._events:
+                f.write(json.dumps(event, default=str) + "\n")
+
+        count = len(self._events)
+        self._events.clear()
+        logger.info("Flushed %d trace events to %s", count, out)
+        return out
+
+    @staticmethod
+    def load(path: Path | str) -> list[dict]:
+        """Load trace events from a JSONL file."""
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"Trace file not found: {path}")
+
+        events = []
+        with open(path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    events.append(json.loads(line))
+        return events
+
+    @staticmethod
+    def find_trace(session_id: str | None = None) -> Path | None:
+        """Find a trace file by session ID prefix, or return the most recent.
+
+        Args:
+            session_id: Session ID or prefix to match. If None, returns
+                the most recent trace file.
+
+        Returns:
+            Path to the trace file, or None if not found.
+        """
+        sessions_dir = _sessions_dir()
+        traces = sorted(
+            sessions_dir.glob("*.trace.jsonl"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        )
+        if not traces:
+            return None
+
+        if session_id is None:
+            return traces[0]
+
+        # Exact match first
+        exact = sessions_dir / f"{session_id}.trace.jsonl"
+        if exact.exists():
+            return exact
+
+        # Prefix match
+        for t in traces:
+            if t.stem.replace(".trace", "").startswith(session_id):
+                return t
+
+        return None

ct/agent/trajectory.py

@@ -0,0 +1,169 @@
+"""
+Trajectory: session memory across queries in interactive mode.
+
+Records queries, answers, entities mentioned, and tools used so the planner
+can reference prior context ("earlier you found X, now also check Y").
+Supports persistence to JSONL for session resume.
+"""
+
+import json
+import time
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+
+
+@dataclass
+class Turn:
+    """A single query-answer turn in a session."""
+    query: str
+    answer: str
+    entities: list[str] = field(default_factory=list)
+    tools_used: list[str] = field(default_factory=list)
+    timestamp: float = field(default_factory=time.time)
+
+
+class Trajectory:
+    """Session memory: records turns and provides context for the planner."""
+
+    def __init__(self, max_turns: int = 20, session_id: str = None, title: str = None):
+        self.turns: list[Turn] = []
+        self.max_turns = max_turns
+        self.session_id = session_id
+        self.title = title
+        self.created_at = time.time()
+
+    def add_turn(self, query: str, answer: str, plan=None):
+        """Record a completed turn with entity extraction."""
+        entities = []
+        tools_used = []
+        if plan and hasattr(plan, "steps"):
+            tools_used = [s.tool for s in plan.steps if s.status == "completed"]
+
+        turn = Turn(
+            query=query,
+            answer=answer,
+            entities=entities,
+            tools_used=tools_used,
+        )
+        self.turns.append(turn)
+
+        # Drop oldest turns if over limit
+        if len(self.turns) > self.max_turns:
+            self.turns = self.turns[-self.max_turns:]
+
+    def context_for_planner(self) -> str:
+        """Format recent turns as context for the planner prompt."""
+        if not self.turns:
+            return ""
+
+        recent = self.turns[-5:]
+        lines = ["## Session Context (prior queries this session)", ""]
+        for i, turn in enumerate(recent, 1):
+            entities_str = ", ".join(turn.entities) if turn.entities else "none"
+            tools_str = ", ".join(turn.tools_used) if turn.tools_used else "none"
+            # Truncate answer to first 200 chars for context
+            answer_preview = turn.answer[:200] + "..." if len(turn.answer) > 200 else turn.answer
+            lines.append(f"**Turn {i}**: {turn.query}")
+            lines.append(f"  Entities: {entities_str}")
+            lines.append(f"  Tools: {tools_str}")
+            lines.append(f"  Finding: {answer_preview}")
+            lines.append("")
+
+        all_entities = self.entities()
+        if all_entities:
+            lines.append(f"**All entities this session**: {', '.join(all_entities)}")
+            lines.append("")
+
+        lines.append(
+            "Use this context to: reference prior findings, avoid repeating work, "
+            "resolve ambiguous entity references (e.g., 'it', 'the compound')."
+        )
+        return "\n".join(lines)
+
+    def entities(self) -> list[str]:
+        """All unique entities mentioned across the session, preserving order."""
+        seen = set()
+        result = []
+        for turn in self.turns:
+            for entity in turn.entities:
+                if entity not in seen:
+                    seen.add(entity)
+                    result.append(entity)
+        return result
+
+    def save(self, path: Path):
+        """Persist trajectory to a JSONL file."""
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, "w") as f:
+            # First line: session metadata
+            meta = {
+                "type": "meta",
+                "session_id": self.session_id,
+                "title": self.title,
+                "created_at": self.created_at,
+                "n_turns": len(self.turns),
+            }
+            f.write(json.dumps(meta) + "\n")
+            # Subsequent lines: turns
+            for turn in self.turns:
+                record = {"type": "turn", **asdict(turn)}
+                f.write(json.dumps(record) + "\n")
+
+    @classmethod
+    def load(cls, path: Path) -> "Trajectory":
+        """Load a trajectory from a JSONL file."""
+        trajectory = cls()
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    record = json.loads(line)
+                except json.JSONDecodeError as e:
+                    import logging
+                    logging.getLogger("ct.trajectory").warning(
+                        "Skipping corrupted line in trajectory %s: %s", path, e,
+                    )
+                    continue
+                if record.get("type") == "meta":
+                    trajectory.session_id = record.get("session_id")
+                    trajectory.title = record.get("title")
+                    trajectory.created_at = record.get("created_at", 0)
+                elif record.get("type") == "turn":
+                    turn = Turn(
+                        query=record.get("query", ""),
+                        answer=record.get("answer", ""),
+                        entities=record.get("entities", []),
+                        tools_used=record.get("tools_used", []),
+                        timestamp=record.get("timestamp", 0),
+                    )
+                    trajectory.turns.append(turn)
+        return trajectory
+
+    @staticmethod
+    def sessions_dir() -> Path:
+        """Return the directory where sessions are stored."""
+        d = Path.home() / ".ct" / "sessions"
+        d.mkdir(parents=True, exist_ok=True)
+        return d
+
+    @classmethod
+    def list_sessions(cls) -> list[dict]:
+        """List all saved sessions, most recent first."""
+        sessions_dir = cls.sessions_dir()
+        sessions = []
+        for path in sessions_dir.glob("*.jsonl"):
+            try:
+                with open(path) as f:
+                    first_line = f.readline().strip()
+                    if first_line:
+                        meta = json.loads(first_line)
+                        if meta.get("type") == "meta":
+                            meta["path"] = str(path)
+                            sessions.append(meta)
+            except (json.JSONDecodeError, OSError):
+                continue
+        # Sort by created_at descending (most recent first)
+        sessions.sort(key=lambda s: s.get("created_at", 0), reverse=True)
+        return sessions