Spaces:
Runtime error
Runtime error
File size: 6,409 Bytes
2e9eb8e | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """
LLM model client wrapper.
Supports multiple providers with unified interface.
"""
import os
from typing import Optional, Dict, Any, List
from dotenv import load_dotenv
import google.generativeai as genai
load_dotenv()
class LLMClient:
"""
Unified LLM client supporting multiple providers.
Auto-detects available API keys.
"""
def __init__(self, provider: Optional[str] = None):
"""
Args:
provider: 'gemini', 'groq', 'openai', 'huggingface', or None (auto-detect)
"""
self.provider = provider or self._detect_provider()
self.client = self._init_client()
try:
print(f"\U0001f916 LLM initialized: {self.provider}")
except UnicodeEncodeError:
print(f"[LLM] Initialized: {self.provider}")
def _detect_provider(self) -> str:
"""Auto-detect available API."""
if os.getenv("GEMINI_API_KEY"):
return "gemini"
elif os.getenv("GROQ_API_KEY"):
return "groq"
elif os.getenv("OPENAI_API_KEY"):
return "openai"
elif os.getenv("HF_TOKEN"):
return "huggingface"
else:
print("[WARN] No API key found, using mock mode")
return "mock"
def _init_client(self):
"""Initialize provider client."""
if self.provider == "gemini":
try:
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
return genai.GenerativeModel('gemini-2.5-flash')
except ImportError:
print("[WARN] google-generativeai not installed, falling back to mock")
return None
elif self.provider == "groq":
try:
from groq import Groq
return Groq(api_key=os.getenv("GROQ_API_KEY"))
except ImportError:
print("[WARN] Groq not installed, falling back to mock")
return None
elif self.provider == "openai":
try:
from openai import OpenAI
return OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
except ImportError:
print("[WARN] OpenAI not installed, falling back to mock")
return None
elif self.provider == "huggingface":
try:
from huggingface_hub import InferenceClient
return InferenceClient(token=os.getenv("HF_TOKEN"))
except ImportError:
print("[WARN] HuggingFace not installed, falling back to mock")
return None
return None
def generate(
self,
prompt: str,
system_prompt: str = "You are a helpful AI assistant.",
max_tokens: int = 1024,
temperature: float = 0.7
) -> str:
"""Generate text completion."""
if self.provider == "mock" or self.client is None:
return self._mock_generate(prompt)
try:
if self.provider == "gemini":
response = self.client.generate_content(
f"{system_prompt}\n\n{prompt}",
generation_config={"max_output_tokens": max_tokens, "temperature": temperature}
)
return response.text
elif self.provider == "groq":
response = self.client.chat.completions.create(
model="llama-3.3-70b-versatile",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}
],
max_tokens=max_tokens,
temperature=temperature
)
return response.choices[0].message.content
elif self.provider == "openai":
response = self.client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt}
],
max_tokens=max_tokens,
temperature=temperature
)
return response.choices[0].message.content
elif self.provider == "huggingface":
full_prompt = f"{system_prompt}\n\nUser: {prompt}\n\nAssistant:"
response = self.client.text_generation(
"Qwen/Qwen2.5-72B-Instruct",
prompt=full_prompt,
max_new_tokens=max_tokens,
temperature=temperature
)
return response
except Exception as e:
print(f"[ERROR] LLM Error: {e}")
return self._mock_generate(prompt)
def _mock_generate(self, prompt: str) -> str:
"""Mock generation for testing without API."""
if "github" in prompt.lower():
return "https://github.com/example/repo"
elif "error" in prompt.lower():
return "The error is likely due to missing dependencies. Try installing the required packages."
elif "metric" in prompt.lower():
return "Target accuracy: 95%"
else:
return f"Mock response for: {prompt[:50]}..."
def generate_structured(
self,
prompt: str,
system_prompt: str = "You are a helpful AI assistant.",
) -> Dict[str, Any]:
"""Generate structured JSON response."""
import json
json_prompt = f"{prompt}\n\nRespond ONLY with valid JSON."
response = self.generate(json_prompt, system_prompt)
try:
# Extract JSON
start = response.find('{')
end = response.rfind('}') + 1
if start >= 0 and end > start:
json_str = response[start:end]
return json.loads(json_str)
else:
return {"error": "No JSON found", "raw": response}
except Exception as e:
return {"error": str(e), "raw": response}
|