managed_model.py

import asyncio
from typing import List

from ldai import log
from ldai.models import AICompletionConfig
from ldai.providers.runner import Runner
from ldai.providers.types import JudgeResult, ManagedResult, RunnerResult
from ldai.tracker import LDAIConfigTracker


class ManagedModel:
    """
    LaunchDarkly managed wrapper for AI model invocations.

    Holds a Runner. Handles judge evaluation dispatch and tracking
    automatically via ``create_tracker()``. Conversation history is
    managed by the runner. Obtain an instance via
    ``LDAIClient.create_model()``.
    """

    def __init__(
        self,
        ai_config: AICompletionConfig,
        model_runner: Runner,
    ):
        self._ai_config = ai_config
        self._model_runner = model_runner

    async def run(self, prompt: str) -> ManagedResult:
        """
        Run the model with a prompt string.

        Delegates to the runner, then dispatches judge evaluations and
        records tracking metrics.

        :param prompt: The user prompt to send to the model
        :return: ManagedResult containing the model's response, metric summary,
            and an optional evaluations task
        """
        tracker = self._ai_config.create_tracker()

        result: RunnerResult = await tracker.track_metrics_of_async(
            lambda r: r.metrics,
            lambda: self._model_runner.run(prompt),
        )

        evaluations_task = await self._track_judge_results(tracker, prompt, result.content)

        return ManagedResult(
            content=result.content,
            metrics=tracker.get_summary(),
            raw=result.raw,
            parsed=result.parsed,
            evaluations=evaluations_task,
        )

    async def _track_judge_results(
        self,
        tracker: LDAIConfigTracker,
        input_text: str,
        output_text: str,
    ) -> asyncio.Task[List[JudgeResult]]:
        evaluator_task = self._ai_config.evaluator.evaluate(input_text, output_text)

        async def _run_and_track(eval_task: asyncio.Task) -> List[JudgeResult]:
            results = await eval_task
            for r in results:
                if not r.sampled:
                    continue
                if r.success:
                    try:
                        tracker.track_judge_result(r)
                    except Exception as exc:
                        log.warning("Judge evaluation failed: %s", exc)
                else:
                    log.warning("Judge evaluation failed: %s", r.error_message)
            return results

        return asyncio.create_task(_run_and_track(evaluator_task))

    def get_model_runner(self) -> Runner:
        """
        Return the underlying runner for advanced use.

        :return: The Runner instance.
        """
        return self._model_runner

    def get_config(self) -> AICompletionConfig:
        """Return the AI completion config."""
        return self._ai_config