File size: 2,125 Bytes
f960c36 dfe06b1 7234f57 3e19edc 6a1db5c 034cea3 6a1db5c 0adaacb dfe06b1 3e19edc 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c 0adaacb 6a1db5c dfe06b1 6a1db5c dfe06b1 3e19edc dfe06b1 6a1db5c 0adaacb dfe06b1 3e19edc dfe06b1 3e19edc ffca4ae 3e19edc 034cea3 0adaacb 034cea3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os
from huggingface_hub import InferenceClient
from tavily import TavilyClient
# === API Keys ===
HF_TOKEN = os.getenv('HF_TOKEN')
GROQ_API_KEY = os.getenv('GROQ_API_KEY')
TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
if not HF_TOKEN:
raise RuntimeError("HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token.")
# === GROQ-Compatible Wrapper ===
class GroqChatClient:
def __init__(self, api_key: str):
import openai
openai.api_key = api_key
openai.api_base = "https://api.groq.com/openai/v1"
self.client = openai
self.chat = self.Chat(openai)
class Chat:
def __init__(self, openai_client):
self.completions = self.Completions(openai_client)
class Completions:
def __init__(self, openai_client):
self.client = openai_client
def create(self, model, messages, temperature=0.7, max_tokens=1024, **kwargs):
return self.client.ChatCompletion.create(
model=model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
**kwargs
)
# === Inference Client Selector ===
def get_inference_client(model_id: str, provider: str = "auto"):
"""
Returns a unified interface:
- For 'moonshotai/Kimi-K2-Instruct', uses Groq with OpenAI-compatible API
- For others, uses Hugging Face InferenceClient
"""
if model_id == "moonshotai/Kimi-K2-Instruct":
if not GROQ_API_KEY:
raise RuntimeError("GROQ_API_KEY is required for Groq-hosted models.")
return GroqChatClient(api_key=GROQ_API_KEY)
return InferenceClient(
model=model_id,
provider=provider,
api_key=HF_TOKEN,
bill_to="huggingface"
)
# === Tavily Search Client ===
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
|