mgbam commited on
Commit
4f22252
·
verified ·
1 Parent(s): bbec1c1

Update hf_client.py

Browse files
Files changed (1) hide show
  1. hf_client.py +21 -10
hf_client.py CHANGED
@@ -1,25 +1,36 @@
 
 
1
  import os
 
2
 
3
- from huggingface_hub import InferenceClient, HfApi
4
  from tavily import TavilyClient
5
 
6
  # HF Inference Client
 
7
  HF_TOKEN = os.getenv('HF_TOKEN')
8
- if not HF_TOKEN:
9
- raise RuntimeError("HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token.")
10
 
11
- def get_inference_client(model_id, provider="auto"):
12
- """Return an InferenceClient with provider based on model_id and user selection."""
 
 
 
 
 
 
 
 
 
 
13
  if model_id == "moonshotai/Kimi-K2-Instruct":
14
  provider = "groq"
15
 
16
- # When using a third-party provider, simply specify the provider.
17
- # Billing will automatically default to the account associated with the HF_TOKEN.
18
- # The 'bill_to' parameter is NOT needed for personal billing.
19
  return InferenceClient(
20
  provider=provider,
21
- api_key=HF_TOKEN
22
- # The `bill_to` parameter has been removed.
23
  )
24
 
25
  # Tavily Search Client
 
1
+ # hf_client.py (REVISED)
2
+
3
  import os
4
+ from typing import Optional
5
 
6
+ from huggingface_hub import InferenceClient
7
  from tavily import TavilyClient
8
 
9
  # HF Inference Client
10
+ # This is your token, used as a fallback for local testing or when a user token isn't available.
11
  HF_TOKEN = os.getenv('HF_TOKEN')
 
 
12
 
13
+ def get_inference_client(model_id: str, provider: str = "auto", user_token: Optional[str] = None):
14
+ """
15
+ Return an InferenceClient.
16
+ Prioritizes the user's token for billing if provided.
17
+ Falls back to the environment HF_TOKEN otherwise.
18
+ """
19
+ # Determine which token to use for the API call
20
+ token_to_use = user_token or HF_TOKEN
21
+
22
+ if not token_to_use:
23
+ raise ValueError("An API token must be provided either by the user or as an environment variable (HF_TOKEN).")
24
+
25
  if model_id == "moonshotai/Kimi-K2-Instruct":
26
  provider = "groq"
27
 
28
+ # Initialize the client with the selected token.
29
+ # The Hugging Face Hub automatically bills the account associated with the token.
30
+ # We do NOT use the 'bill_to' parameter for this.
31
  return InferenceClient(
32
  provider=provider,
33
+ api_key=token_to_use
 
34
  )
35
 
36
  # Tavily Search Client