File size: 1,821 Bytes
f960c36 dfe06b1 7234f57 3e19edc dfe06b1 034cea3 dfe06b1 3e19edc e2e2a15 034cea3 e2e2a15 7234f57 dfe06b1 3e19edc dfe06b1 3e19edc dfe06b1 3e19edc ffca4ae 3e19edc 034cea3 dfe06b1 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 |
import os
import openai
from huggingface_hub import InferenceClient
from tavily import TavilyClient
# === Environment Setup ===
HF_TOKEN = os.getenv('HF_TOKEN')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not HF_TOKEN:
raise RuntimeError(
"HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token."
)
# === Dynamic Inference Client ===
def get_inference_client(model_id: str, provider: str = "auto"):
"""
Return an inference client depending on model ID.
Uses Groq's native API for specific models, otherwise HuggingFace InferenceClient.
"""
if model_id == "moonshotai/Kimi-K2-Instruct":
if not GROQ_API_KEY:
raise RuntimeError("GROQ_API_KEY is required for Groq models.")
# Configure OpenAI client for Groq
openai.api_key = GROQ_API_KEY
openai.api_base = "https://api.groq.com/openai/v1"
def chat(messages, temperature=0.7, max_tokens=1024):
response = openai.ChatCompletion.create(
model="mixtral-8x7b-32768", # You can map the model here
messages=messages,
temperature=temperature,
max_tokens=max_tokens
)
return response["choices"][0]["message"]["content"]
return chat # Return callable interface
# Fallback to Hugging Face
return InferenceClient(
model=model_id,
provider=provider,
api_key=HF_TOKEN,
bill_to="huggingface"
)
# === Tavily 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
|