Spaces:
Running on Zero
Running on Zero
| import json | |
| import os | |
| from pathlib import Path | |
| from typing import Literal | |
| from dotenv import load_dotenv | |
| from langchain_groq import ChatGroq | |
| from langsmith import traceable | |
| from pydantic import BaseModel, Field | |
| ROOT_DIR = Path(__file__).resolve().parents[2] | |
| PROMPT_FILE = ROOT_DIR / "prompts" / "query_classification.json" | |
| ENV_FILE = ROOT_DIR / ".env" | |
| load_dotenv(ENV_FILE) | |
| class QueryClassification(BaseModel): | |
| category: Literal["fact", "detailed", "general"] = Field( | |
| description="The query category used to select the retrieval strategy." | |
| ) | |
| confidence: float = Field( | |
| ge=0.0, | |
| le=1.0, | |
| description="Confidence score for the classification." | |
| ) | |
| class QueryCategorization: | |
| def __init__( | |
| self, | |
| prompt_file: Path = PROMPT_FILE, | |
| model_name: str | None = None, | |
| ) -> None: | |
| model_name = model_name or os.getenv("GROQ_MODEL", os.getenv("LLM_MODEL", "qwen/qwen3.8-27b")) | |
| api_key = os.getenv("GROQ_API_KEY") | |
| if not api_key: | |
| raise EnvironmentError("GROQ_API_KEY is not set.") | |
| self.prompt_file = Path(prompt_file) | |
| if not self.prompt_file.exists(): | |
| raise FileNotFoundError( | |
| f"Prompt file not found: {self.prompt_file}" | |
| ) | |
| self.system_prompt = self._load_prompt() | |
| self.model_name = model_name | |
| self.model = ChatGroq( | |
| model=model_name, | |
| api_key=api_key, | |
| temperature=0, | |
| max_retries=2, | |
| ).with_structured_output(QueryClassification) | |
| def _load_prompt(self) -> str: | |
| try: | |
| data = json.loads( | |
| self.prompt_file.read_text(encoding="utf-8") | |
| ) | |
| except json.JSONDecodeError as exc: | |
| raise ValueError( | |
| f"Invalid JSON in {self.prompt_file}: {exc}" | |
| ) from exc | |
| if isinstance(data, dict): | |
| if isinstance(data.get("system"), str): | |
| return data["system"] | |
| query_classifier = data.get("query_classifier") | |
| if isinstance(query_classifier, dict): | |
| system_prompt = query_classifier.get("system") | |
| if isinstance(system_prompt, str): | |
| return system_prompt | |
| raise ValueError( | |
| "query_classification.json must contain either " | |
| "'system' or 'query_classifier.system'." | |
| ) | |
| def classify(self, query: str) -> QueryClassification: | |
| if not isinstance(query, str): | |
| raise TypeError("Query must be a string.") | |
| query = query.strip() | |
| if not query: | |
| raise ValueError("Query cannot be empty.") | |
| result = self.model.invoke( | |
| [ | |
| ("system", self.system_prompt), | |
| ("human", query) | |
| ], | |
| config={ | |
| "metadata": { | |
| "component": "query_categorization", | |
| "model": self.model_name | |
| }, | |
| "tags": [ | |
| "rag", | |
| "query-classification", | |
| "groq", | |
| "open-source" | |
| ] | |
| } | |
| ) | |
| if not isinstance(result, QueryClassification): | |
| raise TypeError( | |
| f"Unexpected output type: {type(result).__name__}" | |
| ) | |
| return result | |
| _classifier = None | |
| def get_classifier() -> QueryCategorization: | |
| global _classifier | |
| if _classifier is None: | |
| _classifier = QueryCategorization() | |
| return _classifier | |
| def classify_query(query: str) -> QueryClassification: | |
| return get_classifier().classify(query) | |
| if __name__ == "__main__": | |
| classifier = QueryCategorization() | |
| test_queries = [ | |
| "What was the company's revenue in 2025?", | |
| "Explain the complete employee onboarding process.", | |
| "Tell me about the company's AI strategy." | |
| ] | |
| for query in test_queries: | |
| try: | |
| result = classifier.classify(query) | |
| print(f"Query: {query}") | |
| print(f"Category: {result.category}") | |
| print(f"Confidence: {result.confidence:.2f}") | |
| print("-" * 60) | |
| except Exception as exc: | |
| print(f"Error: {exc}") | |