Spaces:
Sleeping
Sleeping
| import os, json | |
| import httpx | |
| from typing import AsyncGenerator | |
| class LLMClient: | |
| def __init__(self, api_key=None, model=None): | |
| self.api_key = api_key or os.getenv("GROQ_API_KEY") | |
| self.model = model or os.getenv("GROQ_MODEL", "llama-3.1-8b-instant") | |
| self.http = httpx.AsyncClient(timeout=30) | |
| async def simple_call(self, prompt: str, temperature: float = 0.0) -> str: | |
| if not self.api_key: | |
| raise RuntimeError( | |
| "Groq API key not configured. Set GROQ_API_KEY environment variable. " | |
| "Get a free API key at https://console.groq.com" | |
| ) | |
| if not self.model: | |
| raise RuntimeError("Groq model not configured. Set GROQ_MODEL environment variable.") | |
| try: | |
| # Groq uses OpenAI-compatible API | |
| r = await self.http.post( | |
| "https://api.groq.com/openai/v1/chat/completions", | |
| headers={ | |
| "Authorization": f"Bearer {self.api_key}", | |
| "Content-Type": "application/json" | |
| }, | |
| json={ | |
| "model": self.model, | |
| "messages": [ | |
| {"role": "user", "content": prompt} | |
| ], | |
| "temperature": temperature, | |
| "stream": False | |
| } | |
| ) | |
| r.raise_for_status() | |
| response_data = r.json() | |
| return response_data["choices"][0]["message"]["content"] | |
| except httpx.HTTPStatusError as e: | |
| error_detail = "Unknown error" | |
| try: | |
| error_json = e.response.json() | |
| error_detail = error_json.get("error", {}).get("message", str(error_json)) | |
| except: | |
| error_detail = e.response.text | |
| raise RuntimeError(f"Groq API error: HTTP {e.response.status_code} - {error_detail}") | |
| except Exception as e: | |
| raise RuntimeError(f"Groq API call failed: {str(e)}") | |
| async def stream_call(self, prompt: str, temperature: float = 0.0) -> AsyncGenerator[str, None]: | |
| """Stream LLM response token by token.""" | |
| if not self.api_key: | |
| raise RuntimeError( | |
| "Groq API key not configured. Set GROQ_API_KEY environment variable. " | |
| "Get a free API key at https://console.groq.com" | |
| ) | |
| if not self.model: | |
| raise RuntimeError("Groq model not configured. Set GROQ_MODEL environment variable.") | |
| try: | |
| async with httpx.AsyncClient(timeout=300.0) as client: | |
| async with client.stream( | |
| "POST", | |
| "https://api.groq.com/openai/v1/chat/completions", | |
| headers={ | |
| "Authorization": f"Bearer {self.api_key}", | |
| "Content-Type": "application/json" | |
| }, | |
| json={ | |
| "model": self.model, | |
| "messages": [ | |
| {"role": "user", "content": prompt} | |
| ], | |
| "temperature": temperature, | |
| "stream": True | |
| } | |
| ) as response: | |
| response.raise_for_status() | |
| async for line in response.aiter_lines(): | |
| if line: | |
| # Groq uses Server-Sent Events format | |
| if line.startswith("data: "): | |
| data_str = line[6:] # Remove "data: " prefix | |
| if data_str.strip() == "[DONE]": | |
| break | |
| try: | |
| data = json.loads(data_str) | |
| delta = data.get("choices", [{}])[0].get("delta", {}) | |
| token = delta.get("content", "") | |
| if token: | |
| yield token | |
| except json.JSONDecodeError: | |
| continue | |
| except httpx.HTTPStatusError as e: | |
| error_detail = "Unknown error" | |
| try: | |
| error_json = e.response.json() | |
| error_detail = error_json.get("error", {}).get("message", str(error_json)) | |
| except: | |
| error_detail = e.response.text | |
| raise RuntimeError(f"Groq API streaming error: HTTP {e.response.status_code} - {error_detail}") | |
| except Exception as e: | |
| raise RuntimeError(f"Groq API streaming failed: {str(e)}") | |