framework-agents 0.1.0__py3-none-any.whl

framework_agents/ontology_review.py

@@ -0,0 +1,321 @@
+"""Ontology self-review + SHACL induction. Parity:
+packages/typescript/src/semantic/ontology-review.ts.
+
+The agent audits what its own ontology IS — not what it wishes it were —
+directly from the store, and identifies what to add:
+
+1. review_ontology(): inventory + gaps. Used-but-undeclared terms,
+   declared-but-unused terms, unlabeled terms, properties missing
+   rdfs:domain/range. "Where is X used?", "what's missing?" — one call.
+2. propose_shapes(): SHACL INDUCTION from instance data — observe each
+   class's instances, emit candidate NodeShapes (min/maxCount, datatype,
+   nodeKind) that explain the observed data. The builder owns shapes: these
+   are PROPOSALS — adopt via adopt_shapes() (validates the existing closure
+   first) or the next migrate_t_box(shapes=...); never a silent adoption.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic import BaseModel
+
+from .semantic import (
+    InMemoryGraphStore,
+    LABEL_PREDICATES,
+    OWL_VOCAB,
+    RDF_TYPE,
+    TBOX_PREDICATES,
+    NodeShape,
+    PropertyShape,
+    Quad,
+    is_literal,
+    strip_literal,
+    term_datatype,
+)
+from .tools import Tool, ToolContext, define_tool
+
+
+@dataclass
+class TermStat:
+    term: str
+    kind: str
+    #: Explicit ABox uses: instance-of counts for classes, fact counts for
+    #: properties.
+    uses: int = 0
+    #: TBox axioms mentioning the term (declarations, subClassOf, domain/range...).
+    axiom_count: int = 0
+    label: str | None = None
+
+
+@dataclass
+class OntologyReview:
+    classes: list[TermStat] = field(default_factory=list)
+    properties: list[TermStat] = field(default_factory=list)
+    #: Used in data with zero TBox declarations — the "what to add" list.
+    undeclared: list[TermStat] = field(default_factory=list)
+    #: Declared but never used in any explicit fact.
+    unused: list[str] = field(default_factory=list)
+    #: Terms without rdfs:label / skos:prefLabel / altLabel.
+    unlabeled: list[str] = field(default_factory=list)
+    #: Properties missing rdfs:domain / rdfs:range axioms.
+    missing_domain: list[str] = field(default_factory=list)
+    missing_range: list[str] = field(default_factory=list)
+    #: Explicit quads vs closure size — how much the ontology explains.
+    explicit_quads: int = 0
+    closure_quads: int = 0
+
+
+def _is_axiom_quad(q: Quad, predicates_in_use: set[str]) -> bool:
+    """True when the quad is an ontology statement, not a fact."""
+    return q.predicate in TBOX_PREDICATES or (
+        q.predicate in RDF_TYPE and (q.subject in predicates_in_use or q.object in OWL_VOCAB)
+    )
+
+
+async def review_ontology(store: InMemoryGraphStore) -> OntologyReview:
+    """Audit the store: inventory + gap list. Pure read — never mutates."""
+    closure = await store.match()
+    explicit = [q for q in closure if q.source != "entailment"]
+    predicates_in_use = {q.predicate for q in explicit}
+
+    labels: dict[str, str] = {}
+    for q in closure:
+        if q.predicate in LABEL_PREDICATES:
+            labels[q.subject] = strip_literal(q.object)
+
+    classes: dict[str, TermStat] = {}
+    properties: dict[str, TermStat] = {}
+
+    def class_(term: str) -> TermStat:
+        if term not in classes:
+            classes[term] = TermStat(term=term, kind="class", label=labels.get(term))
+        return classes[term]
+
+    def property_(term: str) -> TermStat:
+        if term not in properties:
+            properties[term] = TermStat(term=term, kind="property", label=labels.get(term))
+        return properties[term]
+
+    domains: dict[str, int] = {}
+    ranges: dict[str, int] = {}
+    for q in explicit:
+        axiom = _is_axiom_quad(q, predicates_in_use)
+
+        # rdf:type: axiom = term declaration (TBox); fact = instance typing (ABox).
+        if q.predicate in RDF_TYPE:
+            if axiom:
+                class_(q.subject).axiom_count += 1
+            else:
+                class_(q.object).uses += 1
+            continue
+
+        if q.predicate in TBOX_PREDICATES:
+            # TBox statements introduce/declare their subject term.
+            subject = class_(q.subject) if "subClassOf" in q.predicate else property_(q.subject)
+            subject.axiom_count += 1
+            # Their object is a class term (domain/range/subClassOf targets).
+            if "inverseOf" not in q.predicate:
+                class_(q.object).axiom_count += 1
+            if "domain" in q.predicate:
+                domains[q.subject] = domains.get(q.subject, 0) + 1
+            if "range" in q.predicate:
+                ranges[q.subject] = ranges.get(q.subject, 0) + 1
+            continue
+
+        # ABox fact: a property use with actual data behind it.
+        property_(q.predicate).uses += 1
+
+    all_terms = [*classes.values(), *properties.values()]
+    undeclared = sorted(
+        (t for t in all_terms if t.uses > 0 and t.axiom_count == 0),
+        key=lambda t: -t.uses,
+    )
+    unused = [t.term for t in all_terms if t.axiom_count > 0 and t.uses == 0]
+    unlabeled = [t.term for t in all_terms if not t.label]
+    known_props = [t for t in properties.values() if t.axiom_count > 0 or t.uses > 0]
+    missing_domain = [t.term for t in known_props if t.term not in domains]
+    missing_range = [t.term for t in known_props if t.term not in ranges]
+
+    return OntologyReview(
+        classes=sorted(classes.values(), key=lambda t: (-t.uses, t.term)),
+        properties=sorted(properties.values(), key=lambda t: (-t.uses, t.term)),
+        undeclared=undeclared,
+        unused=unused,
+        unlabeled=unlabeled,
+        missing_domain=missing_domain,
+        missing_range=missing_range,
+        explicit_quads=len(explicit),
+        closure_quads=len(closure),
+    )
+
+
+# ---------------------------------------------------------------------------
+# SHACL induction from instance data
+# ---------------------------------------------------------------------------
+
+
+async def propose_shapes(
+    store: InMemoryGraphStore,
+    min_instances: int = 2,
+    coverage: float = 0.5,
+) -> list[NodeShape]:
+    """Propose NodeShapes by observing explicit instance data: per class,
+    per-property coverage, cardinality, datatypes, node kinds across its
+    instances. PROPOSAL ONLY — shapes are builder-owned."""
+    explicit = [q for q in await store.match() if q.source != "entailment"]
+    predicates_in_use = {q.predicate for q in explicit}
+
+    # Class -> instance subjects, from EXPLICIT type facts only (closure types
+    # would smear subclass properties up into parent shapes).
+    instances_of: dict[str, set[str]] = {}
+    for q in explicit:
+        if q.predicate not in RDF_TYPE or _is_axiom_quad(q, predicates_in_use):
+            continue
+        instances_of.setdefault(q.object, set()).add(q.subject)
+
+    shapes: list[NodeShape] = []
+    for cls, subjects in instances_of.items():
+        if len(subjects) < min_instances:
+            continue
+        by_subject: dict[str, list[Quad]] = {}
+        for q in explicit:
+            if q.subject in subjects:
+                by_subject.setdefault(q.subject, []).append(q)
+        observed: dict[str, dict[str, Any]] = {}
+        for quads in by_subject.values():
+            by_path: dict[str, list[Quad]] = {}
+            for q in quads:
+                if q.predicate in RDF_TYPE:
+                    continue
+                by_path.setdefault(q.predicate, []).append(q)
+            for path, path_quads in by_path.items():
+                o = observed.setdefault(
+                    path, {"present": 0, "max_seen": 0, "datatypes": set(), "node_kinds": set()}
+                )
+                o["present"] += 1
+                o["max_seen"] = max(o["max_seen"], len(path_quads))
+                for q in path_quads:
+                    o["node_kinds"].add("literal" if is_literal(q.object) else "iri")
+                    dt = term_datatype(q.object)
+                    if dt:
+                        o["datatypes"].add(dt)
+        properties: list[PropertyShape] = []
+        for path, o in observed.items():
+            if o["present"] / len(subjects) < coverage:
+                continue
+            shape = PropertyShape(path=path)
+            if o["present"] == len(subjects):
+                shape.min_count = 1
+            if o["max_seen"] == 1:
+                shape.max_count = 1
+            if len(o["datatypes"]) == 1:
+                shape.datatype = next(iter(o["datatypes"]))
+            if len(o["node_kinds"]) == 1:
+                shape.node_kind = next(iter(o["node_kinds"]))
+            properties.append(shape)
+        if properties:
+            shapes.append(NodeShape(target_class=cls, properties=properties))
+    return shapes
+
+
+@dataclass
+class RejectedShape:
+    shape: NodeShape
+    violations: list[str]
+
+
+@dataclass
+class ShapeAdoption:
+    adopted: list[NodeShape] = field(default_factory=list)
+    rejected: list[RejectedShape] = field(default_factory=list)
+
+
+async def adopt_shapes(store: InMemoryGraphStore, shapes: list[NodeShape]) -> ShapeAdoption:
+    """Adopt proposed shapes: each shape is validated against the CURRENT
+    closure in isolation — a shape the existing data already violates is
+    rejected with its violations (1:1 mapping, no guesswork), the clean ones
+    are registered. Versioned adoption: migrate_t_box(shapes=...) instead."""
+    closure = await store.match()
+    adoption = ShapeAdoption()
+    for shape in shapes:
+        # Validate with ONLY this shape: violations map to it unambiguously.
+        violations = InMemoryGraphStore(shapes=[shape]).check_shapes(closure).violations
+        if violations:
+            adoption.rejected.append(RejectedShape(shape=shape, violations=[v.message for v in violations]))
+        else:
+            adoption.adopted.append(shape)
+    if adoption.adopted:
+        store.add_shapes(*adoption.adopted)
+    return adoption
+
+
+# ---------------------------------------------------------------------------
+# Agent tools
+# ---------------------------------------------------------------------------
+
+
+class _NoArgs(BaseModel):
+    pass
+
+
+def review_ontology_tool(store: InMemoryGraphStore) -> Tool:
+    """review_ontology agent tool: inventory + gaps, straight from the store."""
+
+    async def execute(args: _NoArgs, ctx: ToolContext | None) -> dict[str, Any]:
+        review = await review_ontology(store)
+
+        def stat(t: TermStat) -> dict[str, Any]:
+            return {
+                "term": t.term,
+                "kind": t.kind,
+                "uses": t.uses,
+                "axiomCount": t.axiom_count,
+                **({"label": t.label} if t.label is not None else {}),
+            }
+
+        return {
+            "classes": [stat(t) for t in review.classes],
+            "properties": [stat(t) for t in review.properties],
+            "undeclared": [stat(t) for t in review.undeclared],
+            "unused": review.unused,
+            "unlabeled": review.unlabeled,
+            "missingDomain": review.missing_domain,
+            "missingRange": review.missing_range,
+            "explicitQuads": review.explicit_quads,
+            "closureQuads": review.closure_quads,
+        }
+
+    return define_tool(
+        name="review_ontology",
+        description="Audit the knowledge graph's ontology: classes/properties with usage counts, used-but-undeclared terms, unlabeled/unused terms, properties missing domain or range. Read-only.",
+        schema=_NoArgs,
+        execute=execute,
+    )
+
+
+def propose_shapes_tool(
+    store: InMemoryGraphStore,
+    min_instances: int = 2,
+    coverage: float = 0.5,
+) -> Tool:
+    """propose_shapes agent tool: SHACL candidates induced from instance data."""
+
+    async def shape_json(s: NodeShape) -> dict[str, Any]:
+        return {
+            "targetClass": s.target_class,
+            "properties": [
+                {k: v for k, v in p.__dict__.items() if v is not None} for p in s.properties
+            ],
+        }
+
+    async def execute(args: _NoArgs, ctx: ToolContext | None) -> list[dict[str, Any]]:
+        return [await shape_json(s) for s in await propose_shapes(store, min_instances, coverage)]
+
+    return define_tool(
+        name="propose_shapes",
+        description="Induce candidate SHACL shapes from existing instance data (>= 2 instances per class, >= 50% property coverage). PROPOSAL: adoption is the builder's call via adopt_shapes(), never automatic.",
+        schema=_NoArgs,
+        execute=execute,
+    )

framework_agents/orchestration.py

@@ -0,0 +1,169 @@
+"""Multi-agent orchestration. Parity: packages/typescript/src/orchestration/index.ts.
+
+Strategies:
+- "sequential": chained — each agent gets the previous agent's text appended
+  to the original input; combined text is the LAST result.
+- "parallel": fan-out with a concurrency cap; results keep agent order.
+- "router": one model call picks an agent by index; only that agent runs.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any, Literal, Optional
+
+from .agent import Agent, AgentResult
+from .errors import AgentError
+from .models import CompleteOptions, ModelProvider
+from .types import ChatMessage
+
+OrchestrationStrategy = Literal["sequential", "parallel", "router"]
+
+ROUTER_SYSTEM = "You route user requests to specialist agents. Reply with ONLY the agent number."
+
+
+@dataclass
+class OrchestrationPlan:
+    strategy: OrchestrationStrategy
+    agents: list[Agent] = field(default_factory=list)
+    """Max concurrent agents for parallel. Default: all at once."""
+    concurrency: Optional[int] = None
+    """Model that picks the agent for router. Defaults to the first agent's model."""
+    router_model: Optional[ModelProvider] = None
+
+
+@dataclass
+class OrchestratedResult:
+    strategy: OrchestrationStrategy
+    results: list[AgentResult]
+    """Combined text: last result (sequential/router), joined (parallel)."""
+    text: str
+    """Router only: absolute index of the chosen agent."""
+    routed_to: Optional[int] = None
+
+
+def _require_agents(plan: OrchestrationPlan) -> None:
+    if not plan.agents:
+        raise AgentError("ORCHESTRATION_EMPTY", "Orchestration needs at least one agent", retryable=False)
+
+
+def _description_of(agent: Any) -> str:
+    opts = getattr(agent, "options", None)
+    desc = getattr(opts, "description", "") if opts is not None else ""
+    if desc:
+        return str(desc)
+    fallback = getattr(agent, "description", "")
+    return str(fallback) if fallback else ""
+
+
+def _model_of(agent: Any) -> Optional[ModelProvider]:
+    opts = getattr(agent, "options", None)
+    model = getattr(opts, "model", None) if opts is not None else None
+    if model is None:
+        model = getattr(agent, "model", None)
+    return model
+
+
+async def _call_run(agent: Any, text: str, invocation_state: Optional[dict[str, Any]]) -> Any:
+    """Agent.run takes no invocation-state knob; pass it when the callee does.
+
+    Members of the plan may be plain Agent instances (run(text)) or richer
+    runtimes that accept an invocation-state keyword.
+    """
+    if invocation_state is not None:
+        try:
+            sig = inspect.signature(agent.run)
+        except (TypeError, ValueError):
+            sig = None
+        if sig is not None and any(
+            p.kind is inspect.Parameter.VAR_KEYWORD or p.name == "invocation_state"
+            for p in sig.parameters.values()
+        ):
+            return await agent.run(text, invocation_state=invocation_state)
+    return await agent.run(text)
+
+
+async def _run_sequential(
+    plan: OrchestrationPlan, input_text: str, invocation_state: Optional[dict[str, Any]]
+) -> OrchestratedResult:
+    """Chain: each agent receives the previous agent's text appended to the input."""
+    _require_agents(plan)
+    results: list[AgentResult] = []
+    current = input_text
+    for i, agent in enumerate(plan.agents):
+        res = await _call_run(agent, current, invocation_state)
+        results.append(res)
+        current = f"{input_text}\n\nPrevious result ({agent.name}):\n{res.text}"
+    return OrchestratedResult("sequential", results, results[-1].text if results else "")
+
+
+async def _run_parallel(
+    plan: OrchestrationPlan, input_text: str, invocation_state: Optional[dict[str, Any]]
+) -> OrchestratedResult:
+    """Fan-out with a concurrency cap; results keep agent order."""
+    _require_agents(plan)
+    count = len(plan.agents)
+    limit = plan.concurrency if plan.concurrency is not None else count
+    semaphore = asyncio.Semaphore(max(1, min(max(1, limit), count)))
+    results: list[AgentResult] = [None] * count  # type: ignore[list-item]
+
+    async def run_one(i: int, agent: Any) -> None:
+        async with semaphore:
+            results[i] = await _call_run(agent, input_text, invocation_state)
+
+    await asyncio.gather(*(run_one(i, a) for i, a in enumerate(plan.agents)))
+    return OrchestratedResult("parallel", results, "\n\n".join(r.text for r in results))
+
+
+async def _run_router(
+    plan: OrchestrationPlan, input_text: str, invocation_state: Optional[dict[str, Any]]
+) -> OrchestratedResult:
+    """One model call picks the agent by index; only that agent runs."""
+    _require_agents(plan)
+    router_model = plan.router_model or _model_of(plan.agents[0])
+    if router_model is None:
+        raise AgentError("ROUTER_NO_MATCH", "Router strategy needs a model", retryable=False)
+    menu = "\n".join(f"{i}: {a.name}" + (f" — {_description_of(a)}" if _description_of(a) else "")
+                     for i, a in enumerate(plan.agents))
+    pick = await router_model.complete(
+        CompleteOptions(
+            messages=[
+                ChatMessage(role="system", content=ROUTER_SYSTEM),
+                ChatMessage(role="user", content=f"Agents:\n{menu}\n\nRequest: {input_text}"),
+            ],
+            tools=[],
+        )
+    )
+    match = re.search(r"\d+", pick.text or "")
+    idx = int(match.group(0)) if match else -1
+    if not (0 <= idx < len(plan.agents)):
+        raise AgentError("ROUTER_NO_MATCH", f"Router picked invalid agent: {json.dumps(pick.text)}", retryable=False)
+    res = await _call_run(plan.agents[idx], input_text, invocation_state)
+    return OrchestratedResult("router", [res], getattr(res, "text", ""), routed_to=idx)
+
+
+async def run_orchestrated(
+    plan: OrchestrationPlan,
+    input_text: str,
+    invocation_state: Optional[dict[str, Any]] = None,
+) -> OrchestratedResult:
+    """Run a multi-agent plan. invocation_state threads to agents that accept it."""
+    if plan.strategy == "sequential":
+        return await _run_sequential(plan, input_text, invocation_state)
+    if plan.strategy == "parallel":
+        return await _run_parallel(plan, input_text, invocation_state)
+    if plan.strategy == "router":
+        return await _run_router(plan, input_text, invocation_state)
+    raise AgentError("ORCHESTRATION_STRATEGY", f"Unknown strategy: {plan.strategy}", retryable=False)
+
+
+__all__ = [
+    "OrchestrationPlan",
+    "OrchestrationStrategy",
+    "OrchestratedResult",
+    "run_orchestrated",
+]

framework_agents/pricing.py

@@ -0,0 +1,63 @@
+"""Per-model prices + cost math. Parity: src/models/pricing.ts (post-fix).
+
+Pinned semantics:
+- unknown models -> None (NEVER a misleading zero)
+- local ollama: free first, name collisions with cloud models lose to it
+- reasoning tokens are already inside output_tokens: never added on top
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from .types import TokenUsage as _TokenUsage  # type alias used by call sites
+
+TokenUsage = _TokenUsage
+
+
+class ModelPrice(BaseModel):
+    input_per_million: float
+    output_per_million: float
+
+
+# USD per 1M tokens. Providers send token counts ONLY — cost is always
+# computed client-side against a table. Verify against your provider bill.
+_PRICES: dict[str, ModelPrice] = {
+    "Qwen/Qwen3.5-9B": ModelPrice(input_per_million=0.2, output_per_million=0.8),
+    "meta-llama/Llama-3.3-70B-Instruct-Turbo": ModelPrice(input_per_million=0.88, output_per_million=0.88),
+    "deepseek-ai/DeepSeek-V4.1-Flash": ModelPrice(input_per_million=0.2, output_per_million=0.6),
+    "zai-org/GLM-5.3-Flash": ModelPrice(input_per_million=0.2, output_per_million=0.6),
+    "gpt-4o-mini": ModelPrice(input_per_million=0.15, output_per_million=0.6),
+    "claude-sonnet-4-6": ModelPrice(input_per_million=3, output_per_million=15),
+    "qwen2.5": ModelPrice(input_per_million=0, output_per_million=0),
+}
+
+_FREE = ModelPrice(input_per_million=0, output_per_million=0)
+
+
+def price_for(model_id: str) -> Optional[ModelPrice]:
+    """Look up a price by bare model id (`provider:` prefixes are stripped).
+
+    Local ollama first: even a name colliding with a cloud model (e.g.
+    `ollama:gpt-4o-mini`) runs free, so it must win over the priced table.
+    """
+    if model_id.startswith("ollama:"):
+        return _FREE
+    bare = model_id.split(":", 1)[1] if ":" in model_id else model_id
+    return _PRICES.get(bare)
+
+
+def usage_cost(usage, model_id: str, override: Optional[ModelPrice] = None) -> Optional[float]:
+    """USD cost for usage; unknown models return None (never a fake zero).
+
+    `output_tokens` already includes reasoning tokens on every supported
+    provider, so they are billed at the output rate with no extra addition.
+    """
+    price = override or price_for(model_id)
+    if price is None:
+        return None
+    return (
+        usage.input_tokens * price.input_per_million + usage.output_tokens * price.output_per_million
+    ) / 1_000_000

framework_agents/prompts.py

@@ -0,0 +1,132 @@
+"""Prompts: versioned, renderable templates. No magic strings in agent code.
+
+Parity: packages/typescript/src/prompts/index.ts. Sources: inline
+PromptTemplate, markdown files with frontmatter (BOM/CRLF tolerant), A/B
+variants (weighted random — unbiased via uniform roll over the weight total).
+"""
+
+from __future__ import annotations
+
+import random
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional
+
+FRONTMATTER_RE = re.compile(r"^---\r?\n(.*?)\r?\n---\r?\n?(.*)$", re.DOTALL)
+VAR_RE = re.compile(r"\{\{\s*(\w+)\s*\}\}")
+
+
+@dataclass
+class PromptTemplate:
+    name: str
+    template: str
+    version: str = "0.1.0"
+
+    def render(self, vars: Optional[dict[str, Any]] = None) -> str:
+        values = vars or {}
+
+        def sub(match: re.Match[str]) -> str:
+            key = match.group(1)
+            if key not in values:
+                raise ValueError(f"Missing prompt var: {key} (prompt: {self.name})")
+            return str(values[key])
+
+        return VAR_RE.sub(sub, self.template)
+
+
+def define_prompt(name: str, template: str, version: str = "0.1.0") -> PromptTemplate:
+    """TS parity helper: template + render raising on a missing {{var}}."""
+    return PromptTemplate(name=name, template=template, version=version)
+
+
+def _strip_bom(content: str) -> str:
+    return content[1:] if content.startswith("") else content
+
+
+def _parse_meta(fm: str) -> dict[str, str]:
+    meta: dict[str, str] = {}
+    for line in fm.split("\n"):
+        i = line.find(":")
+        if i > 0:
+            key = line[:i].strip()
+            value = line[i + 1 :].strip()
+            if (value.startswith('"') and value.endswith('"')) or (
+                value.startswith("'") and value.endswith("'")
+            ):
+                value = value[1:-1]
+            meta[key] = value
+    return meta
+
+
+def load_prompt_file(path: str | Path) -> PromptTemplate:
+    """Load `--- name, version ---` frontmatter + markdown body from disk."""
+    raw = _strip_bom(Path(path).read_text(encoding="utf-8"))
+    match = FRONTMATTER_RE.match(raw)
+    if not match:
+        raise ValueError(f"Invalid prompt file (missing frontmatter): {path}")
+    meta = _parse_meta(match.group(1))
+    if not meta.get("name"):
+        raise ValueError(f"Invalid prompt file (missing name): {path}")
+    return define_prompt(meta["name"], match.group(2).strip(), meta.get("version") or "0.1.0")
+
+
+def _variant_fields(variant: Any) -> tuple[str, str, Optional[str], Optional[float]]:
+    if isinstance(variant, dict):
+        return (
+            variant.get("name", ""),
+            variant.get("template", ""),
+            variant.get("version"),
+            variant.get("weight"),
+        )
+    return (
+        getattr(variant, "name", ""),
+        getattr(variant, "template", ""),
+        getattr(variant, "version", None),
+        getattr(variant, "weight", None),
+    )
+
+
+def ab_prompts(
+    variants: list[Any],
+    pick: Optional[str] = None,
+    weights: Optional[list[float]] = None,
+) -> PromptTemplate:
+    """A/B variants: pick by weighted traffic split, or an explicit name.
+
+    `pick` matches "name" or "name@version". `weights` (parallel to variants)
+    overrides per-variant weights; default weight is 1. The roll is a single
+    uniform draw over the weight total — the interval walk is unbiased.
+    """
+    if not variants:
+        raise ValueError("ab_prompts needs at least one variant")
+    if pick:
+        for variant in variants:
+            name, template, version, _ = _variant_fields(variant)
+            label = f"{name}@{version}" if version else name
+            if label == pick or name == pick:
+                return define_prompt(name, template, version or "0.1.0")
+        raise ValueError(f"Unknown prompt variant: {pick}")
+
+    effective = [
+        weights[i] if weights is not None and i < len(weights) else (_variant_fields(v)[3] or 1)
+        for i, v in enumerate(variants)
+    ]
+    total = sum(effective)
+    roll = random.uniform(0, total)
+    for variant, weight in zip(variants, effective):
+        roll -= weight
+        if roll <= 0:
+            name, template, version, _ = _variant_fields(variant)
+            return define_prompt(name, template, version or "0.1.0")
+    # Floating-point misses are astronomically rare; deterministic tail.
+    name, template, version, _ = _variant_fields(variants[-1])
+    return define_prompt(name, template, version or "0.1.0")
+
+
+__all__ = [
+    "PromptTemplate",
+    "ab_prompts",
+    "define_prompt",
+    "load_prompt_file",
+]