feat(huggingface): add HuggingFace provider (#156)

Author: Kludex

deploy/example.config.ts

@@ -2,7 +2,7 @@ import { env } from 'cloudflare:workers'
 import type { Config } from '@deploy/types'
 
 // can be whatever you want, just used to make linking apiKeys to providers typesafe.
-type ProviderKeys = 'openai' | 'anthropic' | 'google-vertex' | 'bedrock' | 'groq' | 'azure'
+type ProviderKeys = 'openai' | 'anthropic' | 'google-vertex' | 'bedrock' | 'groq' | 'azure' | 'huggingface'
 
 // projects, users and keys must have numeric keys, using constants here to make it easier to understand
 // of course, keys must be unique within a type (e.g. project ids must be unique) but users and projects can have the same id
@@ -98,6 +98,12 @@ export const config: Config<ProviderKeys> = {
       injectCost: true,
       credentials: env.AWS_BEARER_TOKEN_BEDROCK,
     },
+    huggingface: {
+      providerId: 'huggingface',
+      baseUrl: 'https://router.huggingface.co/v1',
+      injectCost: true,
+      credentials: env.HF_TOKEN,
+    },
   },
   // individual apiKeys
   apiKeys: {

deploy/example.env.local

@@ -23,5 +23,8 @@ AWS_BEARER_TOKEN_BEDROCK=...
 # python -c "import json;print(json.dumps(json.loads(open(input('Service account JSON file path: ')).read())))"
 GOOGLE_SERVICE_ACCOUNT_KEY=full service google service account key...
 
+# same for Hugging Face, generate a token (you would use env.HF_TOKEN in config.ts)
+HF_TOKEN=...
+
 # password for viewing /status/
 STATUS_AUTH_API_KEY="change-me!"

examples/ex_huggingface.py

@@ -0,0 +1,15 @@
+import os
+
+from openai import OpenAI
+
+api_key = os.getenv('PYDANTIC_AI_GATEWAY_API_KEY')
+assert api_key is not None
+
+client = OpenAI(api_key=api_key, base_url='http://localhost:8787/huggingface/v1')
+
+completion = client.chat.completions.create(
+    model='openai/gpt-oss-20b:hyperbolic',
+    messages=[{'role': 'user', 'content': 'What is the capital of France?'}],
+)
+
+print(completion.choices[0].message)

examples/pai_huggingface.py

@@ -0,0 +1,48 @@
+import os
+from datetime import date
+
+import logfire
+from huggingface_hub import AsyncInferenceClient
+from pydantic import BaseModel, field_validator
+from pydantic_ai import Agent, __version__
+from pydantic_ai.models.huggingface import HuggingFaceModel
+from pydantic_ai.providers.huggingface import HuggingFaceProvider
+
+logfire.configure(service_name='testing')
+logfire.instrument_pydantic_ai()
+logfire.instrument_aiohttp_client(capture_all=True)
+print('pydantic-ai version:', __version__)
+
+
+class Person(BaseModel, use_attribute_docstrings=True):
+    name: str
+    """The name of the person."""
+    dob: date
+    """The date of birth of the person. MUST BE A VALID ISO 8601 date."""
+    city: str
+    """The city where the person lives."""
+
+    @field_validator('dob')
+    def validate_dob(cls, v: date) -> date:
+        if v >= date(1900, 1, 1):
+            raise ValueError('The person must be born in the 19th century')
+        return v
+
+
+api_key = os.getenv('PYDANTIC_AI_GATEWAY_API_KEY')
+# api_key = os.getenv('HF_TOKEN')
+assert api_key is not None
+base_url = 'http://localhost:8787/huggingface'
+# base_url = None
+
+hf_client = AsyncInferenceClient(api_key=api_key, provider='novita', base_url=base_url)
+provider = HuggingFaceProvider(hf_client=hf_client)
+model = HuggingFaceModel('moonshotai/Kimi-K2-Thinking', provider=provider)
+
+person_agent = Agent(
+    model=model,
+    output_type=Person,
+    instructions='Extract information about the person',
+)
+result = person_agent.run_sync("Samuel lived in London and was born on Jan 28th '87")
+print(repr(result.output))

examples/pyproject.toml

@@ -4,9 +4,10 @@ version = "0.1.0"
 dependencies = [
     "boto3>=1.40.28",
     "devtools>=0.12.2",
-    "logfire[httpx]>=4.3.3",
+    "logfire[httpx,aiohttp]>=4.3.3",
     "opentelemetry-instrumentation-botocore>=0.57b0",
-    "pydantic-ai>=1.10.0",
+    "pydantic-ai[huggingface]>=1.10.0",
+    "huggingface-hub<1.0",
     "types-boto3[bedrock-runtime]",
     "mypy-boto3-bedrock-runtime",
 ]

gateway/src/api/base.ts

@@ -1,3 +1,4 @@
+import type { Provider as UsageProvider } from '@pydantic/genai-prices'
 import { extractUsage, findProvider, type Usage } from '@pydantic/genai-prices'
 import type { GenAIAttributes, GenAIAttributesExtractor } from '../otel/attributes'
 import type { InputMessages, OutputMessages, TextPart } from '../otel/genai'
@@ -48,6 +49,8 @@ export interface SafeExtractor<RequestBody, ResponseBody, StreamChunk> {
 export abstract class BaseAPI<RequestBody, ResponseBody, StreamChunk = JsonData>
   implements GenAIAttributesExtractor<RequestBody, ResponseBody>, SafeExtractor<RequestBody, ResponseBody, StreamChunk>
 {
+  private usageProvider: UsageProvider | undefined
+
   /** @apiFlavor: the flavor of the API, used to determine the response model and usage */
   apiFlavor: string | undefined = undefined
 
@@ -57,9 +60,10 @@ export abstract class BaseAPI<RequestBody, ResponseBody, StreamChunk = JsonData>
   extractedRequest: ExtractedRequest = {}
   extractedResponse: Partial<ExtractedResponse> = {}
 
-  constructor(providerId: ProviderID, requestModel?: string) {
+  constructor(providerId: ProviderID, requestModel?: string, options?: { usageProvider?: UsageProvider }) {
     this.providerId = providerId
     this.requestModel = requestModel
+    this.usageProvider = options?.usageProvider
   }
 
   requestExtractors: ExtractorConfig<RequestBody, ExtractedRequest> = {}
@@ -86,7 +90,7 @@ export abstract class BaseAPI<RequestBody, ResponseBody, StreamChunk = JsonData>
   }
 
   extractUsage(responseBody: ResponseBody | StreamChunk): Usage | undefined {
-    const provider = findProvider({ providerId: this.providerId })
+    const provider = this.usageProvider ?? findProvider({ providerId: this.providerId })
     // This should never happen because we know the provider ID is valid, but we will throw an error to be safe.
     if (!provider) throw new Error(`Provider not found for provider ID: ${this.providerId}`)
     const { usage } = extractUsage(provider, responseBody, this.apiFlavor)

gateway/src/providers/default.ts

@@ -334,7 +334,8 @@ export class DefaultProviderProxy {
     const url = this.url()
 
     // Validate that it's possible to calculate the price for the request model.
-    if (requestModel && this.providerProxy.disableKey) {
+    // HuggingFace is an exception because we will only know the real provider in the response headers.
+    if (requestModel && this.providerProxy.disableKey && this.providerId() !== 'huggingface') {
       const price = calcPrice({ input_tokens: 0, output_tokens: 0 }, requestModel, { provider: this.usageProvider() })
       if (!price) {
         return { modelNotFound: true, requestModel }

gateway/src/providers/huggingface.ts

@@ -0,0 +1,28 @@
+import { findProvider, type Provider as UsageProvider } from '@pydantic/genai-prices'
+import type { ModelAPI } from '../api'
+import { ChatCompletionAPI } from '../api/chat'
+import { DefaultProviderProxy } from './default'
+
+export class HuggingFaceProvider extends DefaultProviderProxy {
+  // This provider refers to the provider that will be used to calculate the price.
+  protected provider: string | null = null
+
+  protected modelAPI(): ModelAPI {
+    return new ChatCompletionAPI('huggingface', undefined, { usageProvider: this.usageProvider() })
+  }
+
+  apiFlavor(): string | undefined {
+    return 'chat'
+  }
+
+  // We need to do this magic, because the `provider` is only set in the response headers.
+  protected usageProvider(): UsageProvider | undefined {
+    return findProvider({ providerId: `${this.providerId()}-${this.provider ?? 'unknown'}` })
+  }
+
+  protected responseHeaders(headers: Headers): Headers {
+    const newHeaders = super.responseHeaders(headers)
+    this.provider = headers.get('x-inference-provider')
+    return newHeaders
+  }
+}

gateway/src/providers/index.ts

@@ -23,6 +23,7 @@ import { BedrockProvider } from './bedrock'
 import { DefaultProviderProxy, type ProviderOptions } from './default'
 import { GoogleVertexProvider } from './google'
 import { GroqProvider } from './groq'
+import { HuggingFaceProvider } from './huggingface'
 import { OpenAIProvider } from './openai'
 import { TestProvider } from './test'
 
@@ -42,6 +43,8 @@ export function getProvider(providerId: ProviderID): ProviderSig {
       return AnthropicProvider
     case 'bedrock':
       return BedrockProvider
+    case 'huggingface':
+      return HuggingFaceProvider
     case 'test':
       return TestProvider
     default:

gateway/src/types.ts

@@ -38,7 +38,15 @@ export interface ApiKeyInfo<ProviderKey extends string = string> {
   otelSettings?: OtelSettings
 }
 
-export type ProviderID = 'groq' | 'openai' | 'google-vertex' | 'anthropic' | 'test' | 'bedrock' | 'azure'
+export type ProviderID =
+  | 'azure'
+  | 'groq'
+  | 'openai'
+  | 'google-vertex'
+  | 'anthropic'
+  | 'test'
+  | 'bedrock'
+  | 'huggingface'
 // TODO | 'fireworks' | 'mistral' | 'cohere'
 
 const providerIds: Record<ProviderID, boolean> = {
@@ -48,6 +56,7 @@ const providerIds: Record<ProviderID, boolean> = {
   anthropic: true,
   test: true,
   bedrock: true,
+  huggingface: true,
   azure: true,
 }