Use Jev for decisions

from collections.abc import Mapping
from typing import Any

from typesafe_sdk import AsyncTypeSafeClient, Choice as JevChoice, ChoiceAnswer

from harnest.decisions import (
    Choice, ChoiceResult, DecisionCapabilities, DecisionRequest,
    DecisionResponse, QuestionKind,
)


def json_value(value: Any) -> Any:
    """Convert Harnest's immutable state snapshot into SDK JSON containers."""
    if isinstance(value, Mapping):
        return {key: json_value(item) for key, item in value.items()}
    if isinstance(value, tuple):
        return [json_value(item) for item in value]
    return value


class JevProvider:
    """Adapt Jev Choice answers to Harnest's decision contract."""

    version = "jev-1.13.0"
    capabilities = DecisionCapabilities(
        frozenset({QuestionKind.CHOICE}),
        batching=True, probabilities=True, confidence=True,
    )

    def __init__(self, client: AsyncTypeSafeClient) -> None:
        """Use a lifecycle-owned client without acquiring another connection."""
        self.client = client

    async def evaluate(self, request: DecisionRequest) -> DecisionResponse:
        """Submit all choices together and preserve their native metadata."""
        if any(not isinstance(q, Choice) for q in request.definition.questions):
            raise TypeError("This adapter supports Choice questions only")
        result = await self.client.system_one(
            model=self.version,
            state=json_value(request.state),
            questions={
                q.name: JevChoice(
                    instructions=q.instructions, criteria=dict(q.options),
                )
                for q in request.definition.questions
            },
        )
        if result.model != self.version:
            raise ValueError("Jev returned a different model revision")
        answers = {}
        for name, answer in result.answers.items():
            if not isinstance(answer, ChoiceAnswer):
                raise TypeError("Expected a Jev Choice answer")
            answers[name] = ChoiceResult(
                value=answer.choice,
                probabilities=answer.probabilities,
                confidence=answer.confidence,
            )
        return DecisionResponse(answers)

from contextlib import asynccontextmanager

from typesafe_sdk import AsyncTypeSafeClient, RetryPolicy

from harnest import context, lifecycle
from harnest.decisions import (
    DecisionAction, DecisionBinding, DecisionOutcome, Decisions,
)
from harnest.lib.jev import JevProvider
from harnest.lib.triage import TRIAGE, ROUTING


@lifecycle.resource
@context.provider("decisions")
@asynccontextmanager
async def decisions():
    """Own the Jev client and send evaluation failures to manual review."""
    async with AsyncTypeSafeClient(
        timeout=5, retry=RetryPolicy(max_retries=0),
    ) as client:
        yield Decisions(
            providers={"jev": JevProvider(client)},
            bindings=(DecisionBinding(
                TRIAGE, "jev", ROUTING,
                on_error=DecisionOutcome(DecisionAction.REVIEW),
            ),),
            timeout_seconds=6,
        )