tunarag-python 0.2.1__py3-none-any.whl

tunarag/integrations/runnables.py

@@ -0,0 +1,270 @@
+"""Adapters for LangChain runnables and compiled LangGraph graphs."""
+
+from __future__ import annotations
+
+import inspect
+import time
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, Protocol
+
+from ..dataset import EvaluationExample
+from ..domain import Candidate, UsageRecord
+
+
+class AsyncRunnable(Protocol):
+    """Shared asynchronous invocation surface implemented by both frameworks."""
+
+    async def ainvoke(
+        self,
+        input: Any,
+        config: Any = None,
+        **kwargs: Any,
+    ) -> Any: ...
+
+
+RunnableFactory = Callable[[Candidate], Any]
+RunnableInputMapper = Callable[[Any, Candidate], Any]
+RunnableConfigMapper = Callable[[Candidate], Mapping[str, Any] | None]
+RunnableOutputMapper = Callable[[Any], Any]
+RunnableUsageMapper = Callable[[Any], Sequence[UsageRecord]]
+
+
+class _RunnableAdapter:
+    def __init__(
+        self,
+        runnable: Any,
+        *,
+        factory: RunnableFactory | None,
+        input_key: str,
+        input_mapper: RunnableInputMapper | None,
+        config_mapper: RunnableConfigMapper | None,
+        output_mapper: RunnableOutputMapper | None,
+        usage_mapper: RunnableUsageMapper | None,
+        capture_message_usage: bool,
+        identity: Any,
+        framework: str,
+    ) -> None:
+        if (runnable is None) == (factory is None):
+            raise ValueError("provide exactly one of runnable or factory")
+        if runnable is not None:
+            _require_runnable(runnable)
+        if factory is not None and not callable(factory):
+            raise TypeError("runnable factory must be callable")
+        if not input_key.strip():
+            raise ValueError("runnable input key must not be empty")
+        for name, mapper in (
+            ("input_mapper", input_mapper),
+            ("config_mapper", config_mapper),
+            ("output_mapper", output_mapper),
+            ("usage_mapper", usage_mapper),
+        ):
+            if mapper is not None and not callable(mapper):
+                raise TypeError(f"{name} must be callable")
+        if not isinstance(capture_message_usage, bool):
+            raise TypeError("capture_message_usage must be a boolean")
+        self._runnable = runnable
+        self._factory = factory
+        self._input_key = input_key
+        self._input_mapper = input_mapper
+        self._config_mapper = config_mapper
+        self._output_mapper = output_mapper
+        self._usage_mapper = usage_mapper
+        self._capture_message_usage = capture_message_usage
+        self._identity = identity
+        self._framework = framework
+
+    async def run(self, candidate: Candidate, example: Any) -> tuple[Any, tuple[UsageRecord, ...]]:
+        """Invoke one framework runnable for a TunaRAG candidate and example."""
+
+        runnable = self._resolve_runnable(candidate)
+        runnable_input = self._map_input(example, candidate)
+        config = self._map_config(candidate)
+        started_at = time.perf_counter()
+        raw_output = await runnable.ainvoke(runnable_input, config=config)
+        latency = time.perf_counter() - started_at
+        output = raw_output if self._output_mapper is None else self._output_mapper(raw_output)
+        usage = list(self._map_usage(raw_output))
+        usage.append(
+            UsageRecord(
+                component=f"{self._framework}.runnable",
+                latency_seconds=latency,
+            )
+        )
+        return output, tuple(usage)
+
+    def fingerprint(self) -> dict[str, Any]:
+        """Return stable, provider-secret-free adapter semantics."""
+
+        runnable_identity = _type_name(self._runnable) if self._runnable is not None else None
+        return {
+            "integration": self._framework,
+            "schema_version": 1,
+            "runnable": runnable_identity,
+            "factory": _callable_identity(self._factory),
+            "input_key": self._input_key,
+            "input_mapper": _callable_identity(self._input_mapper),
+            "config_mapper": _callable_identity(self._config_mapper),
+            "output_mapper": _callable_identity(self._output_mapper),
+            "usage_mapper": _callable_identity(self._usage_mapper),
+            "capture_message_usage": self._capture_message_usage,
+            "identity": self._identity,
+        }
+
+    def _resolve_runnable(self, candidate: Candidate) -> Any:
+        if self._runnable is not None:
+            return self._runnable
+        if self._factory is None:
+            raise RuntimeError("runnable adapter is not configured")
+        runnable = self._factory(candidate)
+        _require_runnable(runnable)
+        return runnable
+
+    def _map_input(self, example: Any, candidate: Candidate) -> Any:
+        if self._input_mapper is not None:
+            return self._input_mapper(example, candidate)
+        if isinstance(example, EvaluationExample):
+            return {self._input_key: example.query}
+        return example
+
+    def _map_config(self, candidate: Candidate) -> Mapping[str, Any] | None:
+        if self._config_mapper is None:
+            return {"configurable": dict(candidate.parameters)}
+        config = self._config_mapper(candidate)
+        if config is not None and not isinstance(config, Mapping):
+            raise TypeError("config_mapper must return a mapping or None")
+        return config
+
+    def _map_usage(self, raw_output: Any) -> tuple[UsageRecord, ...]:
+        if self._usage_mapper is not None:
+            usage = tuple(self._usage_mapper(raw_output))
+            if any(not isinstance(item, UsageRecord) for item in usage):
+                raise TypeError("usage_mapper must return only UsageRecord values")
+            return usage
+        if not self._capture_message_usage:
+            return ()
+        return _message_usage(raw_output, framework=self._framework)
+
+
+class LangChainAdapter(_RunnableAdapter):
+    """Expose a LangChain runnable through TunaRAG's standard adapter contract."""
+
+    def __init__(
+        self,
+        runnable: Any = None,
+        *,
+        factory: RunnableFactory | None = None,
+        input_key: str = "question",
+        input_mapper: RunnableInputMapper | None = None,
+        config_mapper: RunnableConfigMapper | None = None,
+        output_mapper: RunnableOutputMapper | None = None,
+        usage_mapper: RunnableUsageMapper | None = None,
+        capture_message_usage: bool = True,
+        identity: Any = None,
+    ) -> None:
+        super().__init__(
+            runnable,
+            factory=factory,
+            input_key=input_key,
+            input_mapper=input_mapper,
+            config_mapper=config_mapper,
+            output_mapper=output_mapper,
+            usage_mapper=usage_mapper,
+            capture_message_usage=capture_message_usage,
+            identity=identity,
+            framework="langchain",
+        )
+
+
+class LangGraphAdapter(_RunnableAdapter):
+    """Expose a compiled LangGraph graph through TunaRAG's adapter contract."""
+
+    def __init__(
+        self,
+        graph: Any = None,
+        *,
+        factory: RunnableFactory | None = None,
+        input_key: str = "question",
+        input_mapper: RunnableInputMapper | None = None,
+        config_mapper: RunnableConfigMapper | None = None,
+        output_mapper: RunnableOutputMapper | None = None,
+        usage_mapper: RunnableUsageMapper | None = None,
+        capture_message_usage: bool = True,
+        identity: Any = None,
+    ) -> None:
+        super().__init__(
+            graph,
+            factory=factory,
+            input_key=input_key,
+            input_mapper=input_mapper,
+            config_mapper=config_mapper,
+            output_mapper=output_mapper,
+            usage_mapper=usage_mapper,
+            capture_message_usage=capture_message_usage,
+            identity=identity,
+            framework="langgraph",
+        )
+
+
+def _message_usage(output: Any, *, framework: str) -> tuple[UsageRecord, ...]:
+    records: list[UsageRecord] = []
+    visited: set[int] = set()
+
+    def collect(value: Any) -> None:
+        if value is None or isinstance(value, (str, bytes, bytearray, bool, int, float)):
+            return
+        value_id = id(value)
+        if value_id in visited:
+            return
+        visited.add(value_id)
+        metadata = (
+            value.get("usage_metadata")
+            if isinstance(value, Mapping)
+            else getattr(value, "usage_metadata", None)
+        )
+        if isinstance(metadata, Mapping):
+            input_tokens = _token_count(metadata.get("input_tokens"))
+            output_tokens = _token_count(metadata.get("output_tokens"))
+            if input_tokens is not None or output_tokens is not None:
+                records.append(
+                    UsageRecord(
+                        component=f"{framework}.message",
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                    )
+                )
+        if isinstance(value, Mapping):
+            for item in value.values():
+                collect(item)
+        elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
+            for item in value:
+                collect(item)
+
+    collect(output)
+    return tuple(records)
+
+
+def _token_count(value: Any) -> int | None:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        return None
+    return value
+
+
+def _require_runnable(value: Any) -> None:
+    method = getattr(value, "ainvoke", None)
+    if method is None or not callable(method):
+        raise TypeError("framework object must provide an async ainvoke method")
+
+
+def _callable_identity(value: Callable[..., Any] | None) -> str | None:
+    if value is None:
+        return None
+    if inspect.ismethod(value) or inspect.isfunction(value):
+        return f"{value.__module__}.{value.__qualname__}"
+    return _type_name(value)
+
+
+def _type_name(value: Any) -> str:
+    value_type = type(value)
+    return f"{value_type.__module__}.{value_type.__qualname__}"

tunarag/objective.py

@@ -0,0 +1,75 @@
+"""Objective aggregation for quality, cost, and latency metrics."""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+
+from .domain import MetricValue, ValueStatus
+
+
+class ObjectiveDirection(str, Enum):
+    """Whether larger or smaller metric values are beneficial."""
+
+    MAXIMIZE = "maximize"
+    MINIMIZE = "minimize"
+
+
+@dataclass(frozen=True, slots=True)
+class ObjectiveTerm:
+    """One weighted metric contribution to an objective score."""
+
+    metric: str
+    weight: float = 1.0
+    direction: ObjectiveDirection = ObjectiveDirection.MAXIMIZE
+    minimum: float | None = None
+    maximum: float | None = None
+
+    def __post_init__(self) -> None:
+        if not self.metric.strip():
+            raise ValueError("objective metric must not be empty")
+        if self.weight < 0:
+            raise ValueError("objective weight must not be negative")
+        if self.minimum is not None and self.maximum is not None and self.minimum >= self.maximum:
+            raise ValueError("objective minimum must be less than maximum")
+        if (self.minimum is None) != (self.maximum is None):
+            raise ValueError("objective minimum and maximum must be provided together")
+
+
+@dataclass(frozen=True, slots=True)
+class ObjectiveResult:
+    """Aggregated score and confidence state."""
+
+    score: float | None
+    status: ValueStatus
+
+
+@dataclass(frozen=True, slots=True)
+class Objective:
+    """Combine metric observations into a benefit-oriented weighted score."""
+
+    terms: tuple[ObjectiveTerm, ...]
+
+    def __post_init__(self) -> None:
+        if not self.terms:
+            raise ValueError("objective must contain at least one term")
+        if sum(term.weight for term in self.terms) == 0:
+            raise ValueError("objective must contain a positive weight")
+
+    def evaluate(self, metrics: Sequence[MetricValue]) -> ObjectiveResult:
+        by_name = {metric.name: metric for metric in metrics}
+        score = 0.0
+        status = ValueStatus.EXACT
+        for term in self.terms:
+            metric = by_name.get(term.metric)
+            if metric is None or metric.value is None:
+                return ObjectiveResult(None, ValueStatus.UNAVAILABLE)
+            value = float(metric.value)
+            if term.minimum is not None and term.maximum is not None:
+                value = (value - term.minimum) / (term.maximum - term.minimum)
+                value = max(0.0, min(1.0, value))
+            if term.direction is ObjectiveDirection.MINIMIZE:
+                value = 1.0 - value if term.minimum is not None else -value
+            score += term.weight * value
+            if metric.status is ValueStatus.ESTIMATED:
+                status = ValueStatus.ESTIMATED
+        return ObjectiveResult(score, status)

tunarag/py.typed

@@ -0,0 +1 @@
+

tunarag/result.py

@@ -0,0 +1,309 @@
+"""Inspectable optimization result models."""
+
+import csv
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+from .domain import Candidate, MetricValue, UsageRecord, ValueStatus
+from .errors import redact
+from .objective import ObjectiveDirection, ObjectiveResult
+from .store import StudyStatus, TrialStatus
+
+
+@dataclass(frozen=True, slots=True)
+class UsageSummary:
+    """Aggregate usage with estimate provenance retained."""
+
+    input_tokens: int
+    output_tokens: int
+    exact_cost: float
+    estimated_cost: float
+    latency_seconds: float
+    status: ValueStatus
+
+
+@dataclass(frozen=True, slots=True)
+class TrialSummary:
+    """Actionable terminal summary for one candidate trial."""
+
+    id: str
+    sequence: int
+    candidate: Candidate
+    status: TrialStatus
+    metrics: tuple[MetricValue, ...]
+    usage: tuple[UsageRecord, ...]
+    objective: ObjectiveResult | None
+    attempts: int
+    cached: bool = False
+    error: str | None = None
+
+    def metric(self, name: str) -> MetricValue | None:
+        """Return one named metric, or None when it was not observed."""
+
+        return next((metric for metric in self.metrics if metric.name == name), None)
+
+
+@dataclass(frozen=True, slots=True)
+class OptimizationResult:
+    """Completed study result, including failures and no-success outcomes."""
+
+    study_id: str
+    study_status: StudyStatus
+    stop_reason: str
+    trials: tuple[TrialSummary, ...]
+    best_trial: TrialSummary | None
+    usage: UsageSummary
+
+    @property
+    def best_candidate(self) -> Candidate | None:
+        """Return the highest-scoring candidate when one succeeded."""
+
+        return None if self.best_trial is None else self.best_trial.candidate
+
+    @property
+    def successful_trials(self) -> tuple[TrialSummary, ...]:
+        """Return successful trials in execution order."""
+
+        return tuple(trial for trial in self.trials if trial.status is TrialStatus.SUCCEEDED)
+
+    @property
+    def failed_trials(self) -> tuple[TrialSummary, ...]:
+        """Return failed trials in execution order."""
+
+        return tuple(trial for trial in self.trials if trial.status is TrialStatus.FAILED)
+
+    @property
+    def ranked_trials(self) -> tuple[TrialSummary, ...]:
+        """Return successful trials by descending objective, then sequence."""
+
+        return tuple(
+            sorted(
+                (
+                    trial
+                    for trial in self.successful_trials
+                    if trial.objective is not None and trial.objective.score is not None
+                ),
+                key=lambda trial: (-_objective_score(trial), trial.sequence),
+            )
+        )
+
+    @property
+    def metric_names(self) -> tuple[str, ...]:
+        """Return all observed metric names in lexical order."""
+
+        return tuple(sorted({metric.name for trial in self.trials for metric in trial.metrics}))
+
+    def metric_leaders(
+        self,
+        name: str,
+        *,
+        direction: ObjectiveDirection = ObjectiveDirection.MAXIMIZE,
+    ) -> tuple[TrialSummary, ...]:
+        """Rank successful trials that have an available named metric."""
+
+        if not name.strip():
+            raise ValueError("metric name must not be empty")
+        trials = [
+            trial
+            for trial in self.successful_trials
+            if (metric := trial.metric(name)) is not None and metric.value is not None
+        ]
+        multiplier = -1.0 if direction is ObjectiveDirection.MAXIMIZE else 1.0
+        return tuple(
+            sorted(
+                trials,
+                key=lambda trial: (
+                    multiplier * cast(float, cast(MetricValue, trial.metric(name)).value),
+                    trial.sequence,
+                ),
+            )
+        )
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return a redacted, JSON-compatible report."""
+
+        return {
+            "schema_version": 1,
+            "study_id": self.study_id,
+            "study_status": self.study_status.value,
+            "stop_reason": self.stop_reason,
+            "best_trial_id": None if self.best_trial is None else self.best_trial.id,
+            "counts": {
+                "trials": len(self.trials),
+                "successful": len(self.successful_trials),
+                "failed": len(self.failed_trials),
+            },
+            "usage": _usage_dict(self.usage),
+            "trials": [_trial_dict(trial) for trial in self.trials],
+        }
+
+    def to_json(self, *, indent: int | None = 2) -> str:
+        """Serialize the redacted report to deterministic JSON."""
+
+        return json.dumps(
+            self.as_dict(),
+            ensure_ascii=False,
+            allow_nan=False,
+            indent=indent,
+            sort_keys=True,
+        )
+
+    def write_json(self, path: str | Path, *, indent: int | None = 2) -> Path:
+        """Write a redacted JSON report and return its path."""
+
+        target = Path(path)
+        target.write_text(self.to_json(indent=indent) + "\n", encoding="utf-8")
+        return target
+
+    def write_csv(self, path: str | Path) -> Path:
+        """Write one redacted, flattened row per trial and return its path."""
+
+        target = Path(path)
+        metric_columns = [
+            column
+            for name in self.metric_names
+            for column in (
+                f"metric.{name}.value",
+                f"metric.{name}.status",
+                f"metric.{name}.coverage",
+            )
+        ]
+        fieldnames = [
+            "study_id",
+            "trial_id",
+            "sequence",
+            "status",
+            "objective_score",
+            "objective_status",
+            "attempts",
+            "cached",
+            "error",
+            "candidate",
+            "input_tokens",
+            "output_tokens",
+            "exact_cost",
+            "estimated_cost",
+            "latency_seconds",
+            "usage_status",
+            *metric_columns,
+        ]
+        with target.open("w", encoding="utf-8", newline="") as stream:
+            writer = csv.DictWriter(stream, fieldnames=fieldnames, extrasaction="raise")
+            writer.writeheader()
+            for trial in self.trials:
+                writer.writerow(_trial_csv_row(self.study_id, trial, self.metric_names))
+        return target
+
+
+def _usage_dict(usage: UsageSummary) -> dict[str, Any]:
+    return {
+        "input_tokens": usage.input_tokens,
+        "output_tokens": usage.output_tokens,
+        "exact_cost": usage.exact_cost,
+        "estimated_cost": usage.estimated_cost,
+        "latency_seconds": usage.latency_seconds,
+        "status": usage.status.value,
+    }
+
+
+def _trial_dict(trial: TrialSummary) -> dict[str, Any]:
+    candidate = cast(Mapping[str, Any], redact(trial.candidate.parameters))
+    return {
+        "id": trial.id,
+        "sequence": trial.sequence,
+        "status": trial.status.value,
+        "candidate": dict(candidate),
+        "metrics": [
+            {
+                "name": metric.name,
+                "value": metric.value,
+                "status": metric.status.value,
+                "coverage": metric.coverage,
+            }
+            for metric in trial.metrics
+        ],
+        "usage": [
+            {
+                "component": item.component,
+                "input_tokens": item.input_tokens,
+                "output_tokens": item.output_tokens,
+                "cost": item.cost,
+                "latency_seconds": item.latency_seconds,
+                "status": item.status.value,
+                "pricing_version": item.pricing_version,
+            }
+            for item in trial.usage
+        ],
+        "objective": (
+            None
+            if trial.objective is None
+            else {"score": trial.objective.score, "status": trial.objective.status.value}
+        ),
+        "attempts": trial.attempts,
+        "cached": trial.cached,
+        "error": trial.error,
+    }
+
+
+def _trial_csv_row(
+    study_id: str, trial: TrialSummary, metric_names: tuple[str, ...]
+) -> dict[str, Any]:
+    usage = _summarize_usage(trial.usage)
+    candidate = redact(trial.candidate.parameters)
+    row: dict[str, Any] = {
+        "study_id": study_id,
+        "trial_id": trial.id,
+        "sequence": trial.sequence,
+        "status": trial.status.value,
+        "objective_score": None if trial.objective is None else trial.objective.score,
+        "objective_status": None if trial.objective is None else trial.objective.status.value,
+        "attempts": trial.attempts,
+        "cached": trial.cached,
+        "error": trial.error,
+        "candidate": json.dumps(
+            candidate, ensure_ascii=False, sort_keys=True, separators=(",", ":")
+        ),
+        "input_tokens": usage.input_tokens,
+        "output_tokens": usage.output_tokens,
+        "exact_cost": usage.exact_cost,
+        "estimated_cost": usage.estimated_cost,
+        "latency_seconds": usage.latency_seconds,
+        "usage_status": usage.status.value,
+    }
+    for name in metric_names:
+        metric = trial.metric(name)
+        row[f"metric.{name}.value"] = None if metric is None else metric.value
+        row[f"metric.{name}.status"] = None if metric is None else metric.status.value
+        row[f"metric.{name}.coverage"] = None if metric is None else metric.coverage
+    return row
+
+
+def _objective_score(trial: TrialSummary) -> float:
+    if trial.objective is None or trial.objective.score is None:
+        raise ValueError("ranked trial has no objective score")
+    return trial.objective.score
+
+
+def _summarize_usage(records: tuple[UsageRecord, ...]) -> UsageSummary:
+    statuses = {record.status for record in records}
+    if not records or statuses == {ValueStatus.UNAVAILABLE}:
+        status = ValueStatus.UNAVAILABLE
+    elif statuses == {ValueStatus.EXACT}:
+        status = ValueStatus.EXACT
+    else:
+        status = ValueStatus.ESTIMATED
+    return UsageSummary(
+        input_tokens=sum(record.input_tokens or 0 for record in records),
+        output_tokens=sum(record.output_tokens or 0 for record in records),
+        exact_cost=sum(
+            record.cost or 0.0 for record in records if record.status is ValueStatus.EXACT
+        ),
+        estimated_cost=sum(
+            record.cost or 0.0 for record in records if record.status is not ValueStatus.EXACT
+        ),
+        latency_seconds=sum(record.latency_seconds or 0.0 for record in records),
+        status=status,
+    )