Source code for denia.agents

"""Agentic layer: LLM agents at the edges of the discovery loop.

DENIA's architecture puts agents where language understanding adds value and
keeps the experiment-selection core classical and reproducible:

* :class:`SchemaAgent` - reads arbitrary column headers/sample rows and returns
  a machine-actionable :class:`~denia.schema.CampaignSchema` (roles, units,
  category/fraction pairings, optimization direction).
* :class:`ReporterAgent` - turns a recommendation batch into a natural-language
  briefing an experimentalist can act on.

All agents degrade gracefully: without the ``anthropic`` package or an
``ANTHROPIC_API_KEY``, DENIA falls back to deterministic heuristics
(:func:`denia.ingest.infer_schema`) and template-based reports, so the full
loop runs offline. Agent output is always validated against the dataframe
before use; the LLM proposes, the code verifies.
"""

from __future__ import annotations

import json
import os
from typing import List, Optional

import pandas as pd

from .ingest import infer_schema
from .schema import CampaignSchema, ColumnRole

_DEFAULT_MODEL = "claude-sonnet-4-6"


def _get_client():
    """Return an Anthropic client, or None if unavailable (offline mode)."""
    if not os.environ.get("ANTHROPIC_API_KEY"):
        return None
    try:
        import anthropic  # type: ignore

        return anthropic.Anthropic()
    except Exception:
        return None


[docs] class SchemaAgent: """Infer a campaign schema from raw headers and sample rows via an LLM. The agent receives only column names, dtypes, and a few sample values - never the full dataset - and must return strict JSON mapping each column to a role. The result is validated column-by-column against the dataframe; any invalid assignment falls back to the heuristic inference for that column. This "propose-verify" pattern is what makes an LLM safe to place inside a scientific pipeline. """ def __init__(self, model: str = _DEFAULT_MODEL): self.model = model def infer( self, df: pd.DataFrame, target_hint: Optional[str] = None, objective_hint: str = "maximize", ) -> CampaignSchema: client = _get_client() baseline = self._safe_baseline(df, target_hint) if client is None: return baseline prompt = self._build_prompt(df, target_hint, objective_hint) try: msg = client.messages.create( model=self.model, max_tokens=2000, messages=[{"role": "user", "content": prompt}], ) text = "".join(b.text for b in msg.content if getattr(b, "type", "") == "text") data = json.loads(text.strip().removeprefix("```json").removesuffix("```").strip()) schema = CampaignSchema.from_dict({"columns": data["columns"], "provenance": "agent"}) return self._validate(schema, df, baseline) except Exception: return baseline # ------------------------------------------------------------------ internals @staticmethod def _safe_baseline(df: pd.DataFrame, target_hint: Optional[str]) -> CampaignSchema: try: return infer_schema(df, target=target_hint) except ValueError: # No target identifiable: return schema with everything as inputs; # Campaign will demand an explicit target. schema = infer_schema(df.assign(__denia_target__=0.0), target="__denia_target__") schema.columns = [c for c in schema.columns if c.name != "__denia_target__"] return schema @staticmethod def _build_prompt(df: pd.DataFrame, target_hint: Optional[str], objective_hint: str) -> str: sample = df.head(4).to_dict(orient="list") cols_desc = [ {"name": str(c), "dtype": str(df[c].dtype), "sample": [str(v) for v in sample[c]]} for c in df.columns ] roles = [r.value for r in ColumnRole] hint = f'The user indicates the optimization target is "{target_hint}" ({objective_hint}).' if target_hint else "" return ( "You are the schema-inference agent of DENIA, an active-learning framework " "for experimental campaigns stored in spreadsheets. Classify every column of " "the table below into exactly one role from this list: " f"{roles}. For fraction_input columns (e.g. 'Cation 1 mol%'), set 'partner' to the " "category column they quantify. Extract units from names like 'Temperature, K'. " "For target columns set 'maximize' true/false. " f"{hint}\n\nColumns:\n{json.dumps(cols_desc, indent=1)}\n\n" 'Respond ONLY with JSON: {"columns": [{"name": ..., "role": ..., "unit": ..., ' '"partner": ..., "maximize": ...}, ...]} - no prose, no markdown fences.' ) @staticmethod def _validate(schema: CampaignSchema, df: pd.DataFrame, baseline: CampaignSchema) -> CampaignSchema: """Column-level verification of agent output against the actual data.""" valid = [] baseline_by_name = {c.name: c for c in baseline.columns} for c in schema.columns: if c.name not in df.columns: continue # hallucinated column: drop numeric = pd.api.types.is_numeric_dtype(df[c.name]) role_needs_numeric = c.role in ( ColumnRole.NUMERIC_INPUT, ColumnRole.FRACTION_INPUT, ColumnRole.TARGET ) if role_needs_numeric and not numeric: c = baseline_by_name.get(c.name, c) # fall back for this column if c.partner is not None and c.partner not in df.columns: c.partner = None valid.append(c) # any dataframe column the agent forgot -> baseline assignment seen = {c.name for c in valid} valid += [c for n, c in baseline_by_name.items() if n not in seen] out = CampaignSchema(columns=valid, provenance="agent") return out if out.targets else baseline
[docs] class ReporterAgent: """Produce a natural-language briefing for a recommendation batch.""" def __init__(self, model: str = _DEFAULT_MODEL): self.model = model def report( self, recommendations: pd.DataFrame, schema: CampaignSchema, history_summary: str, ) -> str: client = _get_client() if client is None: return self._template_report(recommendations, schema, history_summary) try: prompt = ( "You are the reporting agent of DENIA, an active-learning framework for " "experimental design. Write a concise briefing (<250 words) for a lab " "scientist: which experiments to run next and why, referencing the " "acquisition scores and predicted values. Be concrete and cautious; " "do not invent columns or values.\n\n" f"Campaign: {history_summary}\n" f"Targets: {schema.targets}\n" f"Recommended batch:\n{recommendations.to_string()}" ) msg = client.messages.create( model=self.model, max_tokens=800, messages=[{"role": "user", "content": prompt}], ) return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text") except Exception: return self._template_report(recommendations, schema, history_summary) @staticmethod def _template_report(recommendations: pd.DataFrame, schema: CampaignSchema, history_summary: str) -> str: n = len(recommendations) t = ", ".join(schema.targets) return ( f"DENIA recommendation briefing\n" f"Campaign state: {history_summary}\n" f"Recommended batch: {n} experiments, ranked by expected improvement on {t}. " f"Top candidate: index {recommendations.index[0]} " f"(predicted {recommendations.iloc[0].get('__pred__', float('nan')):.3g} " f{recommendations.iloc[0].get('__std__', float('nan')):.2g}). " f"Batch selected for high acquisition value with diversity penalty to avoid redundancy." )