Spaces:
Running
Running
| from groq import Groq | |
| from utils.config import settings | |
| import json | |
| import re | |
| class QueryDecomposerService: | |
| def __init__(self): | |
| self.provider = "groq" # Default | |
| if settings.GROQ_API_KEY: | |
| self.provider = "groq" | |
| self.client = Groq(api_key=settings.GROQ_API_KEY) | |
| self.model_name = getattr(settings, 'GROQ_FAST_MODEL', settings.GROQ_MODEL) | |
| print(f"QueryDecomposerService initialized with Groq model: {self.model_name}") | |
| else: | |
| raise ValueError("GROQ_API_KEY is not set.") | |
| def _get_json_response(self, prompt: str) -> dict: | |
| """Robust method to get JSON response, handling API errors and Markdown""" | |
| import re | |
| import time | |
| for attempt in range(3): | |
| try: | |
| # Primary Strategy: Strict JSON Mode | |
| try: | |
| chat_completion = self.client.chat.completions.create( | |
| messages=[ | |
| {"role": "system", "content": "You are a helpful assistant that outputs JSON."}, | |
| {"role": "user", "content": prompt} | |
| ], | |
| model=self.model_name, | |
| temperature=0, | |
| response_format={"type": "json_object"} | |
| ) | |
| content = chat_completion.choices[0].message.content | |
| return json.loads(content.strip()) | |
| except Exception as api_err: | |
| # Fallback Strategy: Text Mode if JSON validation fails | |
| if "400" in str(api_err) or "json_validate_failed" in str(api_err): | |
| print(f"JSON Mode failed. Falling back to Text Mode (Attempt {attempt+1})...") | |
| chat_completion = self.client.chat.completions.create( | |
| messages=[ | |
| {"role": "system", "content": "You are a helpful assistant that outputs JSON."}, | |
| {"role": "user", "content": prompt + "\n\nOUTPUT RAW JSON ONLY. NO MARKDOWN."} | |
| ], | |
| model=self.model_name, | |
| temperature=0 | |
| ) | |
| content = chat_completion.choices[0].message.content | |
| # Strip Markdown Code Blocks | |
| match = re.search(r"```(?:json)?(.*?)```", content, re.DOTALL) | |
| if match: | |
| content = match.group(1) | |
| return json.loads(content.strip()) | |
| raise api_err | |
| except Exception as e: | |
| print(f"Error in LLM request (Attempt {attempt+1}): {e}") | |
| if attempt < 2: | |
| time.sleep(2) | |
| else: | |
| return {} | |
| return {} | |
| def decompose(self, query: str) -> dict: | |
| prompt = f""" | |
| You are a constitutional query classifier. Analyze the user's question and extract structured information. | |
| User Query: "{query}" | |
| Your task: | |
| 1. Identify the query TYPE (one of: simple_article, amendment_impact, multi_hop, temporal_evolution, cross_reference, comparison) | |
| 2. Extract ENTITIES: article numbers, amendment numbers, years mentioned | |
| 3. Identify INTENT: what user wants (explain, compare, list, trace, relate) | |
| 4. Estimate COMPLEXITY: simple/medium/complex | |
| Return ONLY valid JSON, no other text: | |
| {{ | |
| "query_type": "amendment_impact|simple_article|multi_hop|temporal_evolution|cross_reference|comparison", | |
| "intent": "explain|compare|list|trace|relate", | |
| "entities": {{ | |
| "articles": [list of article numbers as strings], | |
| "amendments": [list of amendment numbers], | |
| "years": [list of years] | |
| }}, | |
| "complexity": "simple|medium|complex", | |
| "raw_query": "{query}" | |
| }} | |
| """ | |
| result = self._get_json_response(prompt) | |
| if not result: | |
| return { | |
| "query_type": "simple_article", | |
| "intent": "explain", | |
| "entities": {"articles": [], "amendments": [], "years": []}, | |
| "complexity": "simple", | |
| "raw_query": query, | |
| "error": "Error in query decomposition" | |
| } | |
| return result | |