Spaces:
Paused
Paused
| from stores.llm.LLMInterface import LLMInterface | |
| import logging | |
| import requests | |
| import re | |
| import os | |
| class GroqProvider(LLMInterface): | |
| def __init__(self, url: str = None, model: str = None, | |
| default_input_max_characters: int = 1000, | |
| default_generation_max_output_tokens: int = 1000, | |
| default_generation_temperature: float = 0.1, api_key: str = None): | |
| self.url = url or "https://api.groq.com/openai/v1" | |
| self.api_key = api_key or os.getenv("GROQ_API_KEY") | |
| self.model = model | |
| self.generation_model_id = None | |
| self.embedding_model = None | |
| self.embedding_model_id = None | |
| self.embedding_size = None | |
| self.default_input_max_characters = default_input_max_characters | |
| self.default_generation_max_output_tokens = default_generation_max_output_tokens | |
| self.default_generation_temperature = default_generation_temperature | |
| self.logger = logging.getLogger(__name__) | |
| def set_generation_model(self, model_id: str): | |
| if model_id: | |
| self.model = model_id | |
| def set_embedding_model(self, model_id: str, embedding_size: int): | |
| if model_id: | |
| self.embedding_model = model_id | |
| self.embedding_size = embedding_size | |
| self.embedding_model_id = model_id | |
| def process_text(self, text: str): | |
| if not text: | |
| return "" | |
| return str(text).strip() | |
| def generate_text(self, prompt: str, chat_history: list = None, | |
| max_output_tokens: int = None, temperature: float = None): | |
| try: | |
| chat_history = chat_history or [] | |
| clean_prompt = self.process_text(prompt) | |
| messages = [] | |
| for entry in chat_history: | |
| messages.append({ | |
| "role": entry.get("role", "user"), | |
| "content": entry.get("content", "") | |
| }) | |
| messages.append({"role": "user", "content": clean_prompt}) | |
| payload = { | |
| "model": self.model, | |
| "messages": messages, | |
| "max_tokens": int(max_output_tokens or self.default_generation_max_output_tokens), | |
| "temperature": float(temperature or self.default_generation_temperature), | |
| } | |
| url = self.url.rstrip("/") + "/chat/completions" | |
| headers = { | |
| "Authorization": f"Bearer {self.api_key}", | |
| "Content-Type": "application/json", | |
| } | |
| resp = requests.post(url, json=payload, headers=headers, timeout=6000) | |
| if resp.status_code != 200: | |
| self.logger.error("Groq generate failed: %s %s", resp.status_code, resp.text) | |
| return None | |
| data = resp.json() | |
| try: | |
| generated_text = data["choices"][0]["message"]["content"].strip() | |
| except (KeyError, IndexError, TypeError): | |
| self.logger.error("Unexpected Groq response structure: %s", data) | |
| return None | |
| if not generated_text: | |
| return None | |
| usage = data.get("usage", {}) | |
| # Groq exposes x_groq.usage.total_time in seconds | |
| total_time_ms = None | |
| try: | |
| total_time_ms = round(data["x_groq"]["usage"]["total_time"] * 1000, 2) | |
| except (KeyError, TypeError): | |
| pass | |
| return { | |
| "model": data.get("model"), | |
| "response": generated_text, | |
| "tokens_generated": usage.get("completion_tokens"), | |
| "total_duration_ms": total_time_ms, | |
| "prompt_eval_tokens": usage.get("prompt_tokens"), | |
| } | |
| except Exception as e: | |
| self.logger.exception("Error in GroqProvider.generate_text: %s", e) | |
| return None | |
| def embed_text(self, text: str, document_type: str = None): | |
| """Groq does not support embeddings — returns None.""" | |
| self.logger.warning("GroqProvider does not support embeddings.") | |
| return None | |
| def construct_prompt(self, prompt: str, role: str): | |
| return { | |
| "role": role, | |
| "content": self.process_text(prompt) | |
| } | |
| def embed_text_batch(self, texts: list[str], batch_size: int = 32): | |
| """Groq does not support embeddings — returns None.""" | |
| self.logger.warning("GroqProvider does not support embeddings.") | |
| return None | |
| def clean_content(self, text: str) -> str: | |
| text = re.sub(r'\[.*?\]\(.*?\)', '', text) | |
| text = re.sub(r'\[[^\]]*\]', '', text) | |
| text = re.sub(r'\n+', '\n', text).strip() | |
| return text | |
| def web_search(self, query: str): | |
| """Groq has no native web search — returns a not-supported notice.""" | |
| self.logger.warning("GroqProvider.web_search is not natively supported.") | |
| return { | |
| "text": "Web search is not natively supported by the Groq API.", | |
| "references": [] | |
| } | |