hf_client.py

# hf_client.py (REVISED)

import os
from typing import Optional

from huggingface_hub import InferenceClient
from tavily import TavilyClient

# HF Inference Client
# This is your token, used as a fallback for local testing or when a user token isn't available.
HF_TOKEN = os.getenv('HF_TOKEN')

def get_inference_client(model_id: str, provider: str = "auto", user_token: Optional[str] = None):
    """
    Return an InferenceClient.
    Prioritizes the user's token for billing if provided.
    Falls back to the environment HF_TOKEN otherwise.
    """
    # Determine which token to use for the API call
    token_to_use = user_token or HF_TOKEN
    
    if not token_to_use:
        raise ValueError("An API token must be provided either by the user or as an environment variable (HF_TOKEN).")

    if model_id == "moonshotai/Kimi-K2-Instruct":
        provider = "groq"
    
    # Initialize the client with the selected token.
    # The Hugging Face Hub automatically bills the account associated with the token.
    # We do NOT use the 'bill_to' parameter for this.
    return InferenceClient(
        provider=provider,
        api_key=token_to_use
    )

# Tavily Search Client
TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
tavily_client = None
if TAVILY_API_KEY:
    try:
        tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
    except Exception as e:
        print(f"Failed to initialize Tavily client: {e}")
        tavily_client = None