diff --git a/.env.example b/.env.example index 21daf4db..7f8bb5d9 100644 --- a/.env.example +++ b/.env.example @@ -49,3 +49,6 @@ AGNES_API_KEY = "" # SILICONFLOW_API_KEY SILICONFLOW_API_KEY = "" + +# Groq API Configuration +GROQ_API_KEY = "" diff --git a/llm_inference/model_inference.py b/llm_inference/model_inference.py index ad231c18..a2baa35c 100644 --- a/llm_inference/model_inference.py +++ b/llm_inference/model_inference.py @@ -33,6 +33,7 @@ def __init__(self): self.deepseek_api_key = os.getenv("DEEPSEEK_API_KEY") self.perplexity_api_key = os.getenv("PERPLEXITY_API_KEY") self.replicate_api_key = os.getenv("REPLICATE_API_KEY") + self.groq_api_key = os.getenv("GROQ_API_KEY") # AWS credentials self.aws_access_key_id = os.getenv("AWS_ACCESS_KEY_ID") @@ -97,7 +98,8 @@ def infer( return self._call_agnes(model_name, prompt) elif provider == "siliconflow": return self._call_siliconflow(model_name, prompt) - + elif provider == "groq": + return self._call_groq(model_name, prompt) else: # Default to Together API for most open-source models return self._call_together(model_name, prompt) @@ -217,6 +219,12 @@ def _get_provider(self, model_name: str) -> str: # Siliconflow "THUDM/GLM-4-9B-0414": "siliconflow", "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B": "siliconflow", + # Groq + "llama-3.3-70b-versatile": "groq", + "llama-3.1-8b-instant": "groq", + "llama-3.1-70b-versatile": "groq", + "mixtral-8x7b-32768": "groq", + "gemma2-9b-it": "groq", } # Check if exact model name is in mapping @@ -920,3 +928,42 @@ def _call_aws(self, model_name: str, prompt: str) -> Dict[str, Any]: "model_used": model_name, "provider": "aws", } + + def _call_groq(self, model_name: str, prompt: str) -> Dict[str, Any]: + """Call Groq API.""" + import openai + + # Groq uses OpenAI-compatible API + client = openai.OpenAI( + api_key=self.groq_api_key, base_url="https://api.groq.com/openai/v1" + ) + + response = client.chat.completions.create( + model=model_name, + messages=[{"role": "user", "content": prompt}], + max_tokens=2048, + temperature=0.7, + ) + + usage = getattr(response, "usage", None) + input_tokens = getattr(usage, "prompt_tokens", 0) if usage is not None else 0 + completion_tokens = ( + getattr(usage, "completion_tokens", 0) if usage is not None else 0 + ) + total_tokens = ( + getattr(usage, "total_tokens", 0) + if usage is not None + else input_tokens + completion_tokens + ) + + return { + "response": response.choices[0].message.content, + "success": True, + "token_usage": { + "input_tokens": input_tokens, + "output_tokens": completion_tokens, + "total_tokens": total_tokens, + }, + "model_used": model_name, + "provider": "groq", + } diff --git a/model_cost/model_cost.json b/model_cost/model_cost.json index ae82c1d1..a1848863 100644 --- a/model_cost/model_cost.json +++ b/model_cost/model_cost.json @@ -386,5 +386,25 @@ "google/gemma-4-31b-it": { "input_token_price_per_million": 0.08, "output_token_price_per_million": 0.35 + }, + "llama-3.3-70b-versatile": { + "input_token_price_per_million": 0.59, + "output_token_price_per_million": 0.79 + }, + "llama-3.1-8b-instant": { + "input_token_price_per_million": 0.05, + "output_token_price_per_million": 0.08 + }, + "llama-3.1-70b-versatile": { + "input_token_price_per_million": 0.59, + "output_token_price_per_million": 0.79 + }, + "mixtral-8x7b-32768": { + "input_token_price_per_million": 0.24, + "output_token_price_per_million": 0.24 + }, + "gemma2-9b-it": { + "input_token_price_per_million": 0.20, + "output_token_price_per_million": 0.20 } } diff --git a/universal_model_names.py b/universal_model_names.py index 884554a2..ae7a6877 100644 --- a/universal_model_names.py +++ b/universal_model_names.py @@ -138,6 +138,12 @@ "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B", # KT-ModelRouter pool additions "google/gemma-4-31b-it", + # Groq + "llama-3.3-70b-versatile", + "llama-3.1-8b-instant", + "llama-3.1-70b-versatile", + "mixtral-8x7b-32768", + "gemma2-9b-it", ]