hf_endpoint_pydantic_ai_model.py

from __future__ import annotations

from dataclasses import dataclass

import jsonref
from huggingface_hub import AsyncInferenceClient
from pydantic_ai.models.huggingface import HuggingFaceModel
from pydantic_ai.profiles import ModelProfile
from pydantic_ai.profiles._json_schema import JsonSchemaTransformer
from pydantic_ai.providers.huggingface import HuggingFaceProvider

from .config import settings


@dataclass(init=False)
class ExpandedSchemaTransformer(JsonSchemaTransformer):
    def transform(self, schema):
        return schema

    def walk(self):
        """
        Expand out all $defs in complex schemas. 
        """
        schema = super().walk()
        items = jsonref.replace_refs(schema, proxies=False, jsonschema=True)
        # delete "$defs" if it exists
        if "$defs" in items:  # type: ignore
            del items["$defs"]  # type: ignore
        return items


hf_endpoint_model = HuggingFaceModel(
    model_name="olmo-2-1124-7b-instruct",  # nothing uses this this, but for the record
    profile=ModelProfile(
        json_schema_transformer=ExpandedSchemaTransformer,
    ),
    provider=HuggingFaceProvider(
        hf_client=AsyncInferenceClient(
            api_key=settings.HF_TOKEN, model=settings.HF_INFERENCE_ENDPOINT
        )
    ),
)