agent-learning 0.4.1__py3-none-any.whl

agent_learning/learners/reinforce.py

@@ -0,0 +1,153 @@
+"""REINFORCE-style policy gradient learner with an EMA value baseline.
+
+The native learner consumes (episode, aggregate reward) pairs and
+applies one batched gradient step per call. The gradient of a
+softmax bandit's logit ``z_a`` w.r.t. the log-likelihood of action
+``a`` is ``1_{action == a} - π(a)``, so the policy-gradient update
+for one episode is:
+
+    Δz_a = lr * (R - b) * (1_{action == a} - π(a))
+
+We accumulate that delta over the batch, average by the batch size,
+add a small entropy bonus to keep exploration alive, and update the
+EMA baseline ``b`` afterwards. Importance sampling weights are
+applied when the snapshot's policy version is older than the
+in-memory policy version, to keep offline updates unbiased without
+exploding gradients (the weight is clipped).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Dict, Iterable, List, Optional
+
+from ..config import LearnerConfig
+from ..policy.base import Policy
+from ..policy.softmax_bandit import SoftmaxPolicy
+from ..types import Episode, Reward, RewardSource
+from .base import Learner, LearnerResult
+
+
+class ReinforceLearner(Learner):
+    """REINFORCE-with-baseline learner for :class:`SoftmaxPolicy`."""
+
+    def __init__(self, config: Optional[LearnerConfig] = None) -> None:
+        self._config = config or LearnerConfig()
+
+    @property
+    def config(self) -> LearnerConfig:
+        return self._config
+
+    def update(
+        self,
+        policy: Policy,
+        episodes: Iterable[Episode],
+        rewards: Iterable[Reward],
+    ) -> LearnerResult:
+        if not isinstance(policy, SoftmaxPolicy):
+            raise TypeError("ReinforceLearner requires a SoftmaxPolicy")
+
+        # Index rewards by episode id - we only consume aggregate rewards
+        episode_list = list(episodes)
+        aggregate_rewards: Dict[str, float] = {}
+        for r in rewards:
+            if r.source != RewardSource.AGGREGATE:
+                continue
+            # Keep the most recent aggregate per episode (sorted by created_at desc)
+            current = aggregate_rewards.get(r.episode_id)
+            if current is None:
+                aggregate_rewards[r.episode_id] = r.value
+
+        snapshot = policy.snapshot()
+        baseline_before = snapshot.baseline
+        probs_now = policy.probabilities()
+        action_ids = [a.id for a in snapshot.actions]
+        action_index = {aid: idx for idx, aid in enumerate(action_ids)}
+
+        deltas: Dict[str, float] = {aid: 0.0 for aid in action_ids}
+        used = 0
+        reward_sum = 0.0
+        entropy = -sum(p * math.log(max(p, 1e-12)) for p in probs_now)
+
+        for episode in episode_list:
+            if episode.action_id is None or episode.action_id not in action_index:
+                continue
+            reward_value = aggregate_rewards.get(episode.id)
+            if reward_value is None:
+                continue
+            advantage = reward_value - baseline_before
+
+            # Importance sampling weight when the episode was logged under a
+            # behaviour policy that differs from the target policy.
+            iw = self._importance_weight(episode, probs_now, action_index)
+            advantage *= iw
+
+            chosen_idx = action_index[episode.action_id]
+            for idx, aid in enumerate(action_ids):
+                indicator = 1.0 if idx == chosen_idx else 0.0
+                grad = indicator - probs_now[idx]
+                deltas[aid] += advantage * grad
+
+            # Entropy regulariser pushes towards the maximum-entropy
+            # distribution: its gradient w.r.t. logit i is -log p_i - H.
+            for idx, aid in enumerate(action_ids):
+                deltas[aid] += self._config.entropy_bonus * (-math.log(max(probs_now[idx], 1e-12)) - entropy)
+
+            used += 1
+            reward_sum += reward_value
+
+        if used == 0:
+            return LearnerResult(
+                episodes_used=0,
+                mean_reward=0.0,
+                baseline_before=baseline_before,
+                baseline_after=baseline_before,
+                logit_deltas={aid: 0.0 for aid in action_ids},
+                extra={"reason": "no actionable episodes"},
+            )
+
+        # Average gradients across the batch and scale by learning rate
+        for aid in deltas:
+            deltas[aid] = self._config.learning_rate * deltas[aid] / used
+
+        # Update EMA baseline using the mean batch reward
+        mean_reward = reward_sum / used
+        decay = self._config.baseline_decay
+        baseline_after = decay * baseline_before + (1.0 - decay) * mean_reward
+
+        policy.apply_update(
+            deltas,
+            baseline=baseline_after,
+            episodes_seen=snapshot.episodes_seen + used,
+        )
+
+        return LearnerResult(
+            episodes_used=used,
+            mean_reward=mean_reward,
+            baseline_before=baseline_before,
+            baseline_after=baseline_after,
+            logit_deltas=deltas,
+            extra={"entropy": entropy},
+        )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _importance_weight(
+        self,
+        episode: Episode,
+        probs_now: List[float],
+        action_index: Dict[str, int],
+    ) -> float:
+        """Return clipped π_target / π_behaviour for an episode."""
+        if episode.action_logprob is None or episode.action_id is None:
+            return 1.0
+        idx = action_index[episode.action_id]
+        p_target = max(probs_now[idx], 1e-12)
+        p_behaviour = max(math.exp(episode.action_logprob), 1e-12)
+        ratio = p_target / p_behaviour
+        return min(ratio, self._config.importance_clip)
+
+
+__all__ = ["ReinforceLearner"]

agent_learning/metrics/__init__.py

@@ -0,0 +1,22 @@
+"""Score-based evaluation metrics for native RL reward shaping.
+
+Each metric is a thin wrapper around an evaluator from
+``azure-ai-evaluation``. The wrapper normalises the raw evaluator score
+into the ``[0, 1]`` range expected by the reward shaper.
+"""
+
+from .base import MetricEvaluator, MetricRequest
+from .intent_resolution import IntentResolutionMetric
+from .task_adherence import TaskAdherenceMetric
+from .task_completion import TaskCompletionMetric
+from .registry import default_metrics, evaluate_all
+
+__all__ = [
+    "IntentResolutionMetric",
+    "MetricEvaluator",
+    "MetricRequest",
+    "TaskAdherenceMetric",
+    "TaskCompletionMetric",
+    "default_metrics",
+    "evaluate_all",
+]

agent_learning/metrics/base.py

@@ -0,0 +1,233 @@
+"""Base classes for native score-based metrics."""
+
+from __future__ import annotations
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from ..config import ScoreConfig
+from ..types import Episode, MetricName, MetricResult
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetricRequest:
+    """Input passed to a metric evaluator.
+
+    The fields mirror the prompty schemas used by azure-ai-evaluation:
+    ``query`` is the user message (or the full conversation history
+    formatted as a string), ``response`` is the agent reply, and
+    ``tool_calls`` / ``system_message`` are populated when required by
+    the underlying evaluator.
+    """
+
+    query: str
+    response: str
+    system_message: Optional[str] = None
+    tool_calls: Optional[str] = None
+    tool_definitions: Optional[str] = None
+    extra: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_episode(cls, episode: Episode) -> "MetricRequest":
+        """Build a metric request from an :class:`Episode`."""
+        query = _format_query(episode)
+        return cls(
+            query=query,
+            response=episode.assistant_output or "",
+            system_message=episode.system_message,
+            tool_calls=_format_tool_calls(episode),
+            tool_definitions=episode.metadata.get("tool_definitions"),
+        )
+
+
+def _format_query(episode: Episode) -> str:
+    """Render the conversation history as the ``query`` field.
+
+    The Azure prompts accept a free-form ``query`` string. We prefer
+    the structured conversation history when present, falling back to
+    the single ``user_input`` field.
+    """
+    if episode.conversation_history:
+        lines: List[str] = []
+        for msg in episode.conversation_history:
+            role = msg.get("role", "user").upper()
+            content = msg.get("content", "")
+            lines.append(f"{role}: {content}")
+        if episode.user_input and (
+            not lines or not lines[-1].lower().startswith("user:")
+        ):
+            lines.append(f"USER: {episode.user_input}")
+        return "\n".join(lines)
+    return episode.user_input or ""
+
+
+def _format_tool_calls(episode: Episode) -> str:
+    """Render tool calls as a JSON string consumable by the scorer."""
+    if not episode.tool_calls:
+        return "[]"
+    payload = [
+        {
+            "name": tc.name,
+            "arguments": tc.arguments,
+            "result": tc.result,
+            "error": tc.error,
+        }
+        for tc in episode.tool_calls
+    ]
+    return json.dumps(payload, ensure_ascii=False)
+
+
+class MetricEvaluator(ABC):
+    """Abstract base for all native scoring metrics."""
+
+    #: Stable identifier - subclasses must override.
+    NAME: MetricName
+
+    def __init__(
+        self,
+        score_config: Optional[ScoreConfig] = None,
+        *,
+        evaluator: Optional[Any] = None,
+    ) -> None:
+        self._score_config = score_config or ScoreConfig()
+        self._evaluator = evaluator  # Allow injection for tests
+
+    # ------------------------------------------------------------------
+    # Construction hooks (subclasses override _build_evaluator)
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    def _build_evaluator(self) -> Any:
+        """Instantiate the underlying azure-ai-evaluation evaluator."""
+
+    @abstractmethod
+    def _normalize(self, raw: Dict[str, Any]) -> Optional[float]:
+        """Map the evaluator's raw output to a ``[0, 1]`` score."""
+
+    @abstractmethod
+    def _build_kwargs(self, request: MetricRequest) -> Dict[str, Any]:
+        """Build the call kwargs for the underlying evaluator."""
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    @property
+    def evaluator(self) -> Any:
+        """The lazily-constructed underlying evaluator."""
+        if self._evaluator is None:
+            if not self._score_config.enabled:
+                raise RuntimeError(
+                    "Score model is not configured. Set AGENT_LEARNING_SCORE_ENDPOINT "
+                    "and AGENT_LEARNING_SCORE_DEPLOYMENT (or pass an explicit "
+                    "ScoreConfig) before evaluating metrics."
+                )
+            self._evaluator = self._build_evaluator()
+        return self._evaluator
+
+    def evaluate(self, request: MetricRequest) -> MetricResult:
+        """Run the underlying scorer and return a normalised result.
+
+        Errors are caught and returned as ``status=skipped`` so that a
+        flaky scorer does not poison the learner. Callers that want
+        strict semantics can inspect ``result.status``.
+        """
+        if not (request.query and request.response):
+            return MetricResult(
+                metric=self.NAME,
+                score=None,
+                normalized=None,
+                status="skipped",
+                reason="query or response is empty",
+                evaluator=self._score_config.azure_deployment or None,
+            )
+
+        try:
+            kwargs = self._build_kwargs(request)
+            raw = self.evaluator(**kwargs)
+            if not isinstance(raw, dict):
+                raw = dict(raw)
+        except Exception as exc:  # pragma: no cover - scoring runtime errors
+            logger.warning("Metric %s failed: %s", self.NAME.value, exc)
+            return MetricResult(
+                metric=self.NAME,
+                score=None,
+                normalized=None,
+                status="skipped",
+                reason=f"evaluator error: {exc}",
+                evaluator=self._score_config.azure_deployment or None,
+            )
+
+        return self._build_result(raw)
+
+    # ------------------------------------------------------------------
+    # Result extraction
+    # ------------------------------------------------------------------
+
+    def _build_result(self, raw: Dict[str, Any]) -> MetricResult:
+        score = _extract_score(raw, self.NAME)
+        normalized = self._normalize(raw)
+        status = "skipped" if score is None else "completed"
+
+        properties = raw.get("properties") if isinstance(raw.get("properties"), dict) else None
+        reason_keys = (
+            f"{self.NAME.value}_reason",
+            "reason",
+        )
+        reason: Optional[str] = None
+        for key in reason_keys:
+            if isinstance(raw.get(key), str):
+                reason = raw[key]
+                break
+
+        return MetricResult(
+            metric=self.NAME,
+            score=score,
+            normalized=normalized,
+            status=status,
+            reason=reason,
+            properties=properties,
+            evaluator=self._score_config.azure_deployment or None,
+            metadata={"raw": _safe_subset(raw)},
+        )
+
+
+def _extract_score(raw: Dict[str, Any], name: MetricName) -> Optional[float]:
+    """Return the numeric score from a scorer response, regardless of key."""
+    keys = (
+        "score",
+        f"{name.value}_score",
+        f"{name.value}",
+    )
+    for key in keys:
+        if key in raw and raw[key] is not None:
+            try:
+                return float(raw[key])
+            except (TypeError, ValueError):
+                continue
+    return None
+
+
+def _safe_subset(raw: Dict[str, Any]) -> Dict[str, Any]:
+    """Trim the raw score dict to a Cosmos-safe size for persistence."""
+    keep_keys = {
+        "score",
+        "status",
+        "reason",
+        "properties",
+        "intent_resolution_score",
+        "intent_resolution_reason",
+        "task_adherence_score",
+        "task_adherence_reason",
+        "task_completion_score",
+        "task_completion_reason",
+    }
+    return {k: v for k, v in raw.items() if k in keep_keys}
+
+
+__all__ = ["MetricEvaluator", "MetricRequest"]

agent_learning/metrics/intent_resolution.py

@@ -0,0 +1,50 @@
+"""Intent resolution metric wrapper.
+
+Score range from the underlying scorer is ``1`` (very poor) to ``5``
+(excellent). We normalise to ``[0, 1]`` via ``(score - 1) / 4`` so
+that a perfect intent resolution becomes a maximum positive reward
+signal after shaping.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from ..types import MetricName
+from .base import MetricEvaluator, MetricRequest
+
+
+class IntentResolutionMetric(MetricEvaluator):
+    """Wraps :class:`azure.ai.evaluation.IntentResolutionEvaluator`."""
+
+    NAME = MetricName.INTENT_RESOLUTION
+
+    def _build_evaluator(self) -> Any:
+        from azure.ai.evaluation import IntentResolutionEvaluator  # type: ignore
+
+        return IntentResolutionEvaluator(
+            model_config=self._score_config.to_model_config(),
+            threshold=self._score_config.threshold,
+        )
+
+    def _build_kwargs(self, request: MetricRequest) -> Dict[str, Any]:
+        kwargs: Dict[str, Any] = {
+            "query": request.query,
+            "response": request.response,
+        }
+        if request.tool_definitions is not None:
+            kwargs["tool_definitions"] = request.tool_definitions
+        return kwargs
+
+    def _normalize(self, raw: Dict[str, Any]) -> Optional[float]:
+        # The scorer produces an integer 1..5
+        from .base import _extract_score  # type: ignore
+
+        score = _extract_score(raw, self.NAME)
+        if score is None:
+            return None
+        clamped = max(1.0, min(5.0, score))
+        return (clamped - 1.0) / 4.0
+
+
+__all__ = ["IntentResolutionMetric"]

agent_learning/metrics/registry.py

@@ -0,0 +1,42 @@
+"""Convenience helpers to evaluate an episode across all default metrics."""
+
+from __future__ import annotations
+
+from typing import Iterable, List, Optional
+
+from ..config import ScoreConfig
+from ..types import Episode, MetricResult
+from .base import MetricEvaluator, MetricRequest
+from .intent_resolution import IntentResolutionMetric
+from .task_adherence import TaskAdherenceMetric
+from .task_completion import TaskCompletionMetric
+
+
+def default_metrics(score_config: Optional[ScoreConfig] = None) -> List[MetricEvaluator]:
+    """Return the three native metrics wired with the same score config."""
+    cfg = score_config or ScoreConfig()
+    return [
+        IntentResolutionMetric(cfg),
+        TaskAdherenceMetric(cfg),
+        TaskCompletionMetric(cfg),
+    ]
+
+
+def evaluate_all(
+    episode: Episode,
+    metrics: Optional[Iterable[MetricEvaluator]] = None,
+    *,
+    score_config: Optional[ScoreConfig] = None,
+) -> List[MetricResult]:
+    """Evaluate one episode against every supplied metric.
+
+    Each metric is given its own try/except inside :meth:`evaluate`,
+    so a single failing scorer will not prevent the others from
+    producing scores.
+    """
+    metric_list = list(metrics) if metrics is not None else default_metrics(score_config)
+    request = MetricRequest.from_episode(episode)
+    return [m.evaluate(request) for m in metric_list]
+
+
+__all__ = ["default_metrics", "evaluate_all"]

agent_learning/metrics/task_adherence.py

@@ -0,0 +1,42 @@
+"""Task adherence metric wrapper.
+
+The underlying scorer emits a binary score ``0`` (material failure) or
+``1`` (pass). It is already in ``[0, 1]`` so normalisation is a no-op.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from ..types import MetricName
+from .base import MetricEvaluator, MetricRequest
+
+
+class TaskAdherenceMetric(MetricEvaluator):
+    """Wraps :class:`azure.ai.evaluation.TaskAdherenceEvaluator`."""
+
+    NAME = MetricName.TASK_ADHERENCE
+
+    def _build_evaluator(self) -> Any:
+        from azure.ai.evaluation import TaskAdherenceEvaluator  # type: ignore
+
+        return TaskAdherenceEvaluator(model_config=self._score_config.to_model_config())
+
+    def _build_kwargs(self, request: MetricRequest) -> Dict[str, Any]:
+        return {
+            "query": request.query,
+            "response": request.response,
+            "system_message": request.system_message or "",
+            "tool_calls": request.tool_calls or "[]",
+        }
+
+    def _normalize(self, raw: Dict[str, Any]) -> Optional[float]:
+        from .base import _extract_score  # type: ignore
+
+        score = _extract_score(raw, self.NAME)
+        if score is None:
+            return None
+        return max(0.0, min(1.0, score))
+
+
+__all__ = ["TaskAdherenceMetric"]

agent_learning/metrics/task_completion.py

@@ -0,0 +1,54 @@
+"""Task completion metric wrapper.
+
+The underlying scorer emits ``1`` (task completed) or ``0`` (task
+failed or incomplete). It is already in ``[0, 1]``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from ..types import MetricName
+from .base import MetricEvaluator, MetricRequest
+
+
+class TaskCompletionMetric(MetricEvaluator):
+    """Wraps the (experimental) ``TaskCompletionEvaluator`` if available.
+
+    Some releases of ``azure-ai-evaluation`` ship ``TaskCompletionEvaluator``
+    and others ship ``CompletenessEvaluator``. We try the former first
+    and fall back to the latter so the SDK stays compatible across
+    minor versions.
+    """
+
+    NAME = MetricName.TASK_COMPLETION
+
+    def _build_evaluator(self) -> Any:
+        try:  # pragma: no cover - both paths are exercised in CI matrix
+            from azure.ai.evaluation import TaskCompletionEvaluator  # type: ignore
+
+            return TaskCompletionEvaluator(model_config=self._score_config.to_model_config())
+        except ImportError:
+            from azure.ai.evaluation import CompletenessEvaluator  # type: ignore
+
+            return CompletenessEvaluator(model_config=self._score_config.to_model_config())
+
+    def _build_kwargs(self, request: MetricRequest) -> Dict[str, Any]:
+        kwargs: Dict[str, Any] = {
+            "query": request.query,
+            "response": request.response,
+        }
+        if request.tool_definitions is not None:
+            kwargs["tool_definitions"] = request.tool_definitions
+        return kwargs
+
+    def _normalize(self, raw: Dict[str, Any]) -> Optional[float]:
+        from .base import _extract_score  # type: ignore
+
+        score = _extract_score(raw, self.NAME)
+        if score is None:
+            return None
+        return max(0.0, min(1.0, score))
+
+
+__all__ = ["TaskCompletionMetric"]

agent_learning/policy/__init__.py

@@ -0,0 +1,7 @@
+"""Policy abstractions and the default softmax bandit policy."""
+
+from .base import Policy
+from .contextual_softmax import ContextualSoftmaxPolicy
+from .softmax_bandit import SoftmaxPolicy
+
+__all__ = ["ContextualSoftmaxPolicy", "Policy", "SoftmaxPolicy"]

agent_learning/policy/base.py

@@ -0,0 +1,59 @@
+"""Abstract policy interface."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, List, Optional, Sequence
+
+from ..types import Action, PolicySnapshot
+
+
+@dataclass
+class Decision:
+    """Output of :meth:`Policy.choose`: which action and with what probability."""
+
+    action: Action
+    logprob: float
+    probabilities: List[float]  # Probability assigned to every action (same order as snapshot.actions)
+
+
+class Policy(ABC):
+    """Abstract base class for native RL policies.
+
+    The Policy is stateless from the caller's perspective: the durable
+    state lives in the :class:`PolicySnapshot` returned by
+    :meth:`snapshot`. Callers persist the snapshot themselves; the
+    Policy never talks to the store directly.
+
+    Concrete policies may use an optional ``state`` argument on
+    :meth:`choose` to condition action selection on contextual features
+    (a request embedding, a session summary, etc.). Marginal policies
+    such as :class:`SoftmaxPolicy` ignore the argument.
+    """
+
+    @abstractmethod
+    def choose(self, state: Optional[Any] = None) -> Decision:
+        """Sample an action according to the current parameters.
+
+        Args:
+            state: Optional contextual feature vector. Contextual
+                policies (e.g. :class:`ContextualSoftmaxPolicy`) require
+                a 1-D array of length ``feature_dim``. Marginal policies
+                ignore the argument.
+        """
+
+    @abstractmethod
+    def actions(self) -> Sequence[Action]:
+        """Return the action space."""
+
+    @abstractmethod
+    def snapshot(self) -> PolicySnapshot:
+        """Return a deep-copyable snapshot of the policy state."""
+
+    @abstractmethod
+    def apply_update(self, deltas: dict, *, baseline: float, episodes_seen: int) -> None:
+        """Apply per-action logit deltas and update bookkeeping."""
+
+
+__all__ = ["Decision", "Policy"]