agent-learning 0.4.1__py3-none-any.whl

agent_learning/__init__.py

@@ -0,0 +1,160 @@
+"""Native reinforcement learning SDK for AI agents.
+
+Replaces the agent-lightning LLM fine-tuning loop with a fully
+native, in-process learner. The SDK is organised into five layers:
+
+- ``agent_learning.types``    - durable record types (``Episode``,
+  ``Reward``, ``PolicySnapshot``, ...).
+- ``agent_learning.storage``  - pluggable persistence (Cosmos DB,
+  local file system, and in-memory).
+- ``agent_learning.metrics``  - score-based metrics that wrap the
+  Azure AI Evaluation evaluators for Intent Resolution, Task
+  Adherence, and Task Completion.
+- ``agent_learning.rewards``  - reward shaping + persistence.
+- ``agent_learning.policy``   - discrete softmax bandit policy.
+- ``agent_learning.learners`` - REINFORCE-with-baseline learner.
+- ``agent_learning.training`` - end-to-end :class:`LearningRunner`.
+
+Quick start::
+
+    from agent_learning import (
+        Action,
+        EpisodeCapture,
+        LearningRunner,
+        SoftmaxPolicy,
+    )
+
+    actions = [Action(id="prompt_A"), Action(id="prompt_B")]
+    policy = SoftmaxPolicy.from_actions(actions, agent_id="dq", task_id="sales-summary")
+
+    # Capture
+    capture = EpisodeCapture()
+    decision = policy.choose()
+    ctx = capture.start(
+        "Tell me my Q3 sales summary",
+        task_id="sales-summary",
+        intent_summary="Summarize Q3 sales",
+        action_type="chat",
+        action_name=decision.action.id,
+        expected_outcome="An accurate Q3 sales summary",
+        policy_id=policy.snapshot().id,
+        policy_version=policy.snapshot().version,
+        action_id=decision.action.id,
+        action_logprob=decision.logprob,
+    )
+    # ... agent runs, records tool calls, produces output ...
+    episode = capture.end(
+        ctx,
+        assistant_output="...",
+        execution_status="completed",
+        result_summary="Returned the sales summary",
+    )
+
+    # Train
+    runner = LearningRunner(policy=policy)
+    run = runner.run_offline_batch("dq", task_id="sales-summary", episode_limit=200)
+"""
+
+from ._version import __version__
+from .capture import CaptureContext, EpisodeCapture, get_capture
+from .classifiers import (
+    AdherenceScorer,
+    Classifier,
+    ClassifierResult,
+    CompletionScorer,
+    IntentScorer,
+    RouterClassifier,
+)
+from .config import (
+    CaptureConfig,
+    CosmosConfig,
+    ScoreConfig,
+    LearnerConfig,
+    ShapingConfig,
+)
+from .learners import Learner, LearnerResult, ReinforceLearner
+from .metrics import (
+    IntentResolutionMetric,
+    MetricEvaluator,
+    MetricRequest,
+    TaskAdherenceMetric,
+    TaskCompletionMetric,
+    default_metrics,
+    evaluate_all,
+)
+from .policy import ContextualSoftmaxPolicy, Policy, SoftmaxPolicy
+from .rewards import RewardShaper, RewardWriter, shape_episode_reward
+from .storage import (
+    CosmosStore,
+    InMemoryStore,
+    LearningStore,
+    LocalFileStore,
+    get_default_store,
+)
+from .training import LearningRunner
+from .types import (
+    Action,
+    AgentSummary,
+    AgentTaskSummary,
+    Episode,
+    MetricName,
+    MetricResult,
+    PolicySnapshot,
+    Reward,
+    RewardSource,
+    ToolCall,
+    TrainingRun,
+    TrainingStatus,
+)
+
+__all__ = [
+    "Action",
+    "AdherenceScorer",
+    "AgentSummary",
+    "AgentTaskSummary",
+    "CaptureConfig",
+    "CaptureContext",
+    "Classifier",
+    "ClassifierResult",
+    "CompletionScorer",
+    "ContextualSoftmaxPolicy",
+    "CosmosConfig",
+    "CosmosStore",
+    "Episode",
+    "EpisodeCapture",
+    "InMemoryStore",
+    "IntentScorer",
+    "IntentResolutionMetric",
+    "ScoreConfig",
+    "Learner",
+    "LearnerConfig",
+    "LearnerResult",
+    "LearningRunner",
+    "LearningStore",
+    "LocalFileStore",
+    "MetricEvaluator",
+    "MetricName",
+    "MetricRequest",
+    "MetricResult",
+    "Policy",
+    "PolicySnapshot",
+    "ReinforceLearner",
+    "Reward",
+    "RewardShaper",
+    "RewardSource",
+    "RewardWriter",
+    "RouterClassifier",
+    "ShapingConfig",
+    "SoftmaxPolicy",
+    "TaskAdherenceMetric",
+    "TaskCompletionMetric",
+    "ToolCall",
+    "TrainingRun",
+    "TrainingStatus",
+    "__version__",
+    "default_metrics",
+    "evaluate_all",
+    "get_capture",
+    "get_default_store",
+    "shape_episode_reward",
+]

agent_learning/_version.py

@@ -0,0 +1,3 @@
+"""Package version."""
+
+__version__ = "0.4.1"

agent_learning/capture.py

@@ -0,0 +1,271 @@
+"""Episode capture helpers.
+
+The capture hook is the orchestrator-facing entry point. It wraps a
+single agent turn, lets the caller record tool calls during the
+turn, and persists a complete :class:`Episode` at the end.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from .config import CaptureConfig
+from .storage.base import LearningStore
+from .storage.cosmos import get_default_store
+from .types import Episode, ToolCall
+
+logger = logging.getLogger(__name__)
+
+
+# Conservative redaction patterns (mirrors lightning's set)
+_REDACT_PATTERNS = [
+    (re.compile(r"(bearer\s+)[a-zA-Z0-9\-_\.]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(api[_-]?key[\"\s:=]+)[a-zA-Z0-9\-_]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(password[\"\s:=]+)[^\s\"]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(secret[\"\s:=]+)[^\s\"]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(token[\"\s:=]+)[a-zA-Z0-9\-_\.]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(connection[_-]?string[\"\s:=]+)[^\s\"]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(AccountKey=)[^;]+", re.IGNORECASE), r"\1[REDACTED]"),
+    (re.compile(r"(SharedAccessSignature=)[^;]+", re.IGNORECASE), r"\1[REDACTED]"),
+]
+
+
+def redact(text: Optional[str]) -> Optional[str]:
+    """Strip well-known secret patterns from a string."""
+    if not text:
+        return text
+    out = text
+    for pattern, replacement in _REDACT_PATTERNS:
+        out = pattern.sub(replacement, out)
+    return out
+
+
+@dataclass
+class CaptureContext:
+    """In-flight state for a single capture."""
+
+    episode_id: str
+    agent_id: str
+    start_time: float
+    user_input: str
+    system_message: Optional[str] = None
+    conversation_history: List[Dict[str, str]] = field(default_factory=list)
+    tool_calls: List[ToolCall] = field(default_factory=list)
+    policy_id: Optional[str] = None
+    policy_version: Optional[int] = None
+    action_id: Optional[str] = None
+    action_logprob: Optional[float] = None
+    context_features: Dict[str, Any] = field(default_factory=dict)
+    model_deployment: Optional[str] = None
+    correlation_id: Optional[str] = None
+    session_id: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    agent_name: Optional[str] = None
+    task_id: str = "default"
+    task_name: Optional[str] = None
+    intent_summary: Optional[str] = None
+    action_type: Optional[str] = None
+    action_name: Optional[str] = None
+    target: Optional[str] = None
+    input_summary: Optional[str] = None
+    expected_outcome: Optional[str] = None
+
+
+class EpisodeCapture:
+    """Capture one episode at a time and persist it via the store."""
+
+    def __init__(
+        self,
+        config: Optional[CaptureConfig] = None,
+        store: Optional[LearningStore] = None,
+    ) -> None:
+        self._config = config or CaptureConfig()
+        self._store = store
+
+    # ------------------------------------------------------------------
+    # Plumbing
+    # ------------------------------------------------------------------
+
+    @property
+    def store(self) -> LearningStore:
+        if self._store is None:
+            self._store = get_default_store()
+        return self._store
+
+    @property
+    def config(self) -> CaptureConfig:
+        return self._config
+
+    def is_enabled(self) -> bool:
+        return self._config.enabled
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(
+        self,
+        user_input: str,
+        *,
+        task_id: Optional[str] = None,
+        task_name: Optional[str] = None,
+        intent_summary: Optional[str] = None,
+        action_type: Optional[str] = None,
+        action_name: Optional[str] = None,
+        target: Optional[str] = None,
+        input_summary: Optional[str] = None,
+        expected_outcome: Optional[str] = None,
+        system_message: Optional[str] = None,
+        conversation_history: Optional[List[Dict[str, str]]] = None,
+        model_deployment: Optional[str] = None,
+        correlation_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        policy_id: Optional[str] = None,
+        policy_version: Optional[int] = None,
+        action_id: Optional[str] = None,
+        action_logprob: Optional[float] = None,
+        context_features: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> CaptureContext:
+        return CaptureContext(
+            episode_id=str(uuid.uuid4()),
+            agent_id=self._config.agent_id,
+            agent_name=self._config.agent_name,
+            task_id=task_id or self._config.task_id,
+            task_name=task_name or self._config.task_name,
+            start_time=time.time(),
+            user_input=user_input,
+            intent_summary=intent_summary,
+            action_type=action_type,
+            action_name=action_name,
+            target=target,
+            input_summary=input_summary,
+            expected_outcome=expected_outcome,
+            system_message=system_message,
+            conversation_history=conversation_history or [],
+            policy_id=policy_id,
+            policy_version=policy_version,
+            action_id=action_id,
+            action_logprob=action_logprob,
+            context_features=context_features or {},
+            model_deployment=model_deployment,
+            correlation_id=correlation_id,
+            session_id=session_id,
+            metadata=metadata or {},
+        )
+
+    def record_tool_call(
+        self,
+        ctx: CaptureContext,
+        name: str,
+        arguments: Dict[str, Any],
+        result: Optional[str] = None,
+        *,
+        duration_ms: Optional[int] = None,
+        error: Optional[str] = None,
+    ) -> None:
+        if not self.is_enabled():
+            return
+
+        safe_args = arguments
+        if self._config.redact_secrets:
+            safe_args = {
+                k: redact(v) if isinstance(v, str) else v for k, v in arguments.items()
+            }
+
+        safe_result = result
+        if safe_result and self._config.redact_secrets:
+            safe_result = redact(safe_result)
+        if safe_result and len(safe_result) > self._config.max_output_length:
+            safe_result = safe_result[: self._config.max_output_length] + "...[TRUNCATED]"
+
+        ctx.tool_calls.append(
+            ToolCall(
+                name=name,
+                arguments=safe_args or {},
+                result=safe_result,
+                duration_ms=duration_ms,
+                error=error,
+            )
+        )
+
+    def end(
+        self,
+        ctx: CaptureContext,
+        assistant_output: str,
+        *,
+        execution_status: Optional[str] = None,
+        result_summary: Optional[str] = None,
+        token_usage: Optional[Dict[str, int]] = None,
+        extra_metadata: Optional[Dict[str, Any]] = None,
+    ) -> Optional[Episode]:
+        if not self.is_enabled():
+            return None
+
+        latency_ms = int((time.time() - ctx.start_time) * 1000)
+        output = assistant_output
+        if output and self._config.redact_secrets:
+            output = redact(output) or output
+
+        metadata = dict(ctx.metadata)
+        if extra_metadata:
+            metadata.update(extra_metadata)
+
+        episode = Episode(
+            id=ctx.episode_id,
+            agent_id=ctx.agent_id,
+            agent_name=ctx.agent_name,
+            task_id=ctx.task_id,
+            task_name=ctx.task_name,
+            user_input=ctx.user_input,
+            assistant_output=output or "",
+            intent_summary=ctx.intent_summary,
+            action_type=ctx.action_type,
+            action_name=ctx.action_name,
+            target=ctx.target,
+            input_summary=ctx.input_summary,
+            expected_outcome=ctx.expected_outcome,
+            execution_status=execution_status,
+            result_summary=result_summary,
+            system_message=ctx.system_message,
+            conversation_history=ctx.conversation_history,
+            tool_calls=ctx.tool_calls,
+            policy_id=ctx.policy_id,
+            policy_version=ctx.policy_version,
+            action_id=ctx.action_id,
+            action_logprob=ctx.action_logprob,
+            context_features=ctx.context_features,
+            model_deployment=ctx.model_deployment,
+            correlation_id=ctx.correlation_id,
+            session_id=ctx.session_id,
+            request_latency_ms=latency_ms,
+            token_usage=token_usage,
+            metadata=metadata,
+        )
+
+        try:
+            self.store.store_episode(episode)
+            logger.info("Captured episode %s for agent %s", episode.id, episode.agent_id)
+        except Exception as exc:  # pragma: no cover
+            logger.warning("Failed to persist episode %s: %s", episode.id, exc)
+
+        return episode
+
+
+_default_capture: Optional[EpisodeCapture] = None
+
+
+def get_capture() -> EpisodeCapture:
+    """Return a process-wide singleton capture hook."""
+    global _default_capture
+    if _default_capture is None:
+        _default_capture = EpisodeCapture()
+    return _default_capture
+
+
+__all__ = ["CaptureContext", "EpisodeCapture", "get_capture", "redact"]

agent_learning/classifiers/__init__.py

@@ -0,0 +1,37 @@
+"""Pure-Python classifiers for RL reward shaping pipelines.
+
+The reward shaper in :mod:`agent_learning.rewards.shaping` combines
+multiple classifier outputs into a scalar reward. This sub-package
+ships the classifier implementations themselves:
+
+- :class:`agent_learning.classifiers.router.RouterClassifier` —
+  multi-class classifier over a context vector. Used to route an
+  incoming request to one of a known set of class ids before the
+  policy chooses an action.
+- :class:`agent_learning.classifiers.scorers.intent.IntentScorer`,
+  :class:`~agent_learning.classifiers.scorers.adherence.AdherenceScorer`,
+  and
+  :class:`~agent_learning.classifiers.scorers.completion.CompletionScorer`
+  — binary ``{pass, fail}`` classifiers over a ``(context, action)``
+  pair. Drop-in replacements for LLM-based scorers with the same
+  call surface.
+
+All classes here are deterministic, dependency-free (stdlib only),
+and JSON-serialisable via ``to_snapshot`` / ``from_snapshot``. The
+fit loop is a plain mini-batch logistic regression so the
+classifiers are reproducible across environments and easy to ship
+through the same image as the rest of the SDK.
+"""
+
+from .base import Classifier, ClassifierResult
+from .router import RouterClassifier
+from .scorers import AdherenceScorer, CompletionScorer, IntentScorer
+
+__all__ = [
+  "AdherenceScorer",
+    "Classifier",
+    "ClassifierResult",
+    "CompletionScorer",
+    "IntentScorer",
+    "RouterClassifier",
+]

agent_learning/classifiers/base.py

@@ -0,0 +1,185 @@
+"""Base types shared by the router and the three scorers."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Protocol
+
+
+@dataclass(frozen=True)
+class ClassifierResult:
+    """Output of any classifier in this package.
+
+    ``label`` is the predicted class id (a ``str`` for the router,
+    ``"pass"`` / ``"fail"`` for the scorers). ``confidence`` is the
+    model's calibrated probability for that label, in ``[0, 1]``.
+    ``features`` exposes the per-feature contribution so callers can
+    persist a per-decision explainability trace.
+    """
+
+    label: str
+    confidence: float
+    features: Dict[str, float] = field(default_factory=dict)
+
+
+class Classifier(Protocol):
+    """Common surface across the router and the three scorers.
+
+    Implementations are deterministic — given the same fit input
+    they produce the same weights, and given the same predict
+    input they produce the same :class:`ClassifierResult`.
+    """
+
+    def fit(self, training_rows: List[dict]) -> "Classifier":
+        ...
+
+    def predict(self, features: Dict[str, float]) -> ClassifierResult:
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Pure-Python logistic regression primitives shared by the four classifiers
+# ---------------------------------------------------------------------------
+
+
+def _sigmoid(x: float) -> float:
+    """Numerically stable logistic function."""
+    if x >= 0.0:
+        z = math.exp(-x)
+        return 1.0 / (1.0 + z)
+    z = math.exp(x)
+    return z / (1.0 + z)
+
+
+def _softmax(logits: List[float]) -> List[float]:
+    """Numerically stable softmax over ``logits``."""
+    if not logits:
+        return []
+    m = max(logits)
+    exps = [math.exp(v - m) for v in logits]
+    s = sum(exps)
+    if s == 0.0:
+        n = len(exps)
+        return [1.0 / n] * n
+    return [v / s for v in exps]
+
+
+def _dot(weights: List[float], features: List[float]) -> float:
+    """Inner product of two equal-length vectors."""
+    total = 0.0
+    for w, x in zip(weights, features):
+        total += w * x
+    return total
+
+
+def fit_binary_logreg(
+    rows: List[dict],
+    feature_dim: int,
+    *,
+    epochs: int = 20,
+    learning_rate: float = 0.10,
+    weight_decay: float = 1e-4,
+    batch_size: int = 256,
+    seed: int = 42,
+) -> List[float]:
+    """Fit a binary logistic regression with closed-form gradient.
+
+    Each row in ``rows`` must carry:
+
+    - ``"features"``: list[float] of length ``feature_dim + 1`` (the
+      last entry is the bias term — caller appends 1.0).
+    - ``"label"``: ``0`` or ``1``.
+
+    Returns a weight vector of length ``feature_dim + 1``. Uses
+    standard mini-batch gradient descent with an L2 penalty; pure
+    Python, no numpy.
+    """
+    import random as _random
+
+    rng = _random.Random(seed)
+    n = feature_dim + 1  # bias column appended by caller
+    weights = [0.0] * n
+    indices = list(range(len(rows)))
+    if not indices:
+        return weights
+
+    for _epoch in range(epochs):
+        rng.shuffle(indices)
+        for start in range(0, len(indices), batch_size):
+            batch = indices[start:start + batch_size]
+            grad = [0.0] * n
+            for i in batch:
+                row = rows[i]
+                x = row["features"]
+                y = float(row["label"])
+                z = _dot(weights, x)
+                p = _sigmoid(z)
+                err = p - y
+                for j in range(n):
+                    grad[j] += err * x[j]
+            inv_b = 1.0 / max(len(batch), 1)
+            for j in range(n):
+                grad[j] = grad[j] * inv_b + weight_decay * weights[j]
+                weights[j] -= learning_rate * grad[j]
+    return weights
+
+
+def fit_multinomial_logreg(
+    rows: List[dict],
+    feature_dim: int,
+    num_classes: int,
+    *,
+    epochs: int = 20,
+    learning_rate: float = 0.10,
+    weight_decay: float = 1e-4,
+    batch_size: int = 256,
+    seed: int = 42,
+) -> List[List[float]]:
+    """Fit a multinomial logistic regression (softmax classifier).
+
+    Each row in ``rows`` must carry:
+
+    - ``"features"``: list[float] of length ``feature_dim + 1``.
+    - ``"label"``: integer in ``[0, num_classes)``.
+
+    Returns a ``num_classes x (feature_dim + 1)`` weight matrix.
+    """
+    import random as _random
+
+    rng = _random.Random(seed)
+    n = feature_dim + 1
+    weights: List[List[float]] = [[0.0] * n for _ in range(num_classes)]
+    indices = list(range(len(rows)))
+    if not indices:
+        return weights
+
+    for _epoch in range(epochs):
+        rng.shuffle(indices)
+        for start in range(0, len(indices), batch_size):
+            batch = indices[start:start + batch_size]
+            grads: List[List[float]] = [[0.0] * n for _ in range(num_classes)]
+            for i in batch:
+                row = rows[i]
+                x = row["features"]
+                y = int(row["label"])
+                logits = [_dot(weights[k], x) for k in range(num_classes)]
+                probs = _softmax(logits)
+                for k in range(num_classes):
+                    err = probs[k] - (1.0 if k == y else 0.0)
+                    for j in range(n):
+                        grads[k][j] += err * x[j]
+            inv_b = 1.0 / max(len(batch), 1)
+            for k in range(num_classes):
+                for j in range(n):
+                    grads[k][j] = grads[k][j] * inv_b + weight_decay * weights[k][j]
+                    weights[k][j] -= learning_rate * grads[k][j]
+    return weights
+
+
+__all__ = [
+    "Classifier",
+    "ClassifierResult",
+    "fit_binary_logreg",
+    "fit_multinomial_logreg",
+]