Spaces:
Sleeping
Sleeping
File size: 2,102 Bytes
d40b9df | 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 | import os
from typing import Optional
from litellm import completion
from dotenv import load_dotenv
load_dotenv()
class LLMClient:
"""LLM client using LiteLLM"""
def __init__(
self,
model: str = "groq/llama-3.3-70b-versatile",
api_key: Optional[str] = None,
temperature: float = 0.1
):
"""
Initialize LLM client
Args:
model: Model identifier (e.g., "groq/llama-3.3-70b-versatile")
api_key: API key (if None, uses GROQ_API_KEY env var)
temperature: Sampling temperature
"""
self.model = model
self.temperature = temperature
if api_key:
os.environ["GROQ_API_KEY"] = api_key
elif "GROQ_API_KEY" not in os.environ:
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
raise ValueError(
"GROQ_API_KEY not found. Please set it as environment variable "
"or pass as api_key parameter. Get free key from https://console.groq.com/"
)
def generate(
self,
prompt: str,
max_tokens: int = 512,
system_prompt: Optional[str] = None
) -> str:
"""
Generate text using LLM
Args:
prompt: User prompt
max_tokens: Maximum tokens to generate
system_prompt: Optional system prompt
Returns:
Generated text
"""
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": prompt})
try:
response = completion(
model=self.model,
messages=messages,
temperature=self.temperature,
max_tokens=max_tokens
)
return response.choices[0].message.content
except Exception as e:
raise Exception(f"Error calling LLM: {str(e)}")
|