hf_client.py

"""
This module handles the creation of API clients for the application.
It includes the logic for instantiating the Hugging Face InferenceClient
and the Tavily Search client.

The get_inference_client function is critical for enabling the "user-pays"
model in a Hugging Face Space. It prioritizes the API token of the logged-in
user, ensuring their account is billed for inference costs.
"""

import os
from typing import Optional

from huggingface_hub import InferenceClient
from tavily import TavilyClient

# --- Hugging Face Inference Client ---

# This is the Space owner's token, loaded from environment secrets.
# It serves as a fallback for local development or when a user-provided token is not available.
HF_TOKEN = os.getenv('HF_TOKEN')

def get_inference_client(model_id: str, provider: str = "auto", user_token: Optional[str] = None) -> InferenceClient:
    """
    Creates and returns a Hugging Face InferenceClient.

    This function implements the "user-pays" logic. It prioritizes using the token
    provided by the logged-in user (`user_token`). If that is not available,
    it falls back to the Space owner's token (`HF_TOKEN`).

    Args:
        model_id (str): The ID of the model to be used (e.g., "mistralai/Mistral-7B-Instruct-v0.2").
        provider (str): The specific inference provider (e.g., "groq"). Defaults to "auto".
        user_token (Optional[str]): The API token of the logged-in user, passed from the Gradio app.

    Returns:
        InferenceClient: An initialized client ready for making API calls.
    
    Raises:
        ValueError: If no API token can be found (neither from the user nor the environment).
    """
    # 1. Determine which token to use for the API call. The user's token takes precedence.
    token_to_use = user_token or HF_TOKEN

    # 2. Validate that we have a token. If not, the application cannot make API calls.
    if not token_to_use:
        raise ValueError(
            "Cannot proceed without an API token. "
            "Please log into Hugging Face, or ensure the HF_TOKEN environment secret is set for this Space."
        )

    # 3. Handle any model-specific provider logic.
    if model_id == "moonshotai/Kimi-K2-Instruct":
        provider = "groq"
    
    # 4. Instantiate and return the client.
    # The Hugging Face Hub automatically bills the account associated with the provided `api_key`.
    # The `bill_to` parameter is NOT needed or used for this user-pays scenario.
    return InferenceClient(
        provider=provider,
        api_key=token_to_use
    )


# --- Tavily Search Client ---

# This client uses the Space owner's TAVILY_API_KEY, as this is a backend service.
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:
        # Log an error if the client fails to initialize, but don't crash the app.
        print(f"Warning: Failed to initialize Tavily client. Web search will be unavailable. Error: {e}")
        tavily_client = None