backend/rag/query_classifier.py

import json
import re
from dataclasses import dataclass
from groq import Groq
from backend.config import GROQ_API_KEY, GROQ_MODEL, GROQ_CLASSIFIER_MODEL
from backend.database.connection import get_connection

@dataclass
class QueryAnalysis:
    intent          : str      # "single_source" | "comparison" | "synthesis"
    source_types    : list[str]  # ["legal_statute", "legal_judgment", "web", "youtube", "any"]
    topics          : list[str]  # extracted topics like ["murder", "302", "IPC"]
    ipc_sections    : list[str]  # ["302", "304"] if mentioned
    time_filter     : str | None # "2024" if year mentioned
    language_hint   : str        # "en" | "hi" — detected from query language
    requires_compare: bool       # True if "compare", "difference", "vs" in query
    requires_summary: bool       # True if "common", "all documents", "across" in query
    source_names    : list[str]  # ["IPC", "CrPC"] if explicitly named
    temporal_markers: list[int]  # [180, 360] for 3min, 6min mentioned in video context
    route           : str = "rag" # "chat" or "rag"

def classify_query(question: str) -> QueryAnalysis:
    client = Groq(api_key=GROQ_API_KEY)
    
    system_prompt = """
    Analyze the user's query. Determine if it is a general chitchat message, casual greeting ("hi", "hello", "how are you"), bot capability question ("what can you do?", "who are you?"), or simple polite message ("thanks", "thank you") that should be routed to standard chitchat response ("chat") OR if it is a research, lookup, summarization, comparison, or file-specific query that requires searching/retrieving facts from the knowledge base/documents/videos ("rag").
    
    Return ONLY valid JSON with these exact keys:
    {
      "route": "chat" | "rag",
      "intent": "single_source" | "comparison" | "synthesis",
      "source_types": list of "legal_statute"|"legal_judgment"|"web"|"youtube"|"any",
      "topics": list of topic strings,
      "ipc_sections": list of section numbers as strings,
      "time_filter": year string or null,
      "language_hint": "en" or "hi",
      "requires_compare": boolean,
      "requires_summary": boolean,
      "source_names": list of named sources,
      "temporal_markers": list of integers (seconds) for any time marks mentioned like "3 minutes", "10:30", "at 45 seconds". 
                          Convert all to total seconds.
    }
    """

    try:
        response = client.chat.completions.create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Query: {question}"}
            ],
            model=GROQ_CLASSIFIER_MODEL,
            max_tokens=300,
            stream=False,
            response_format={"type": "json_object"}
        )
        
        data = json.loads(response.choices[0].message.content)
        
        return QueryAnalysis(
            intent=data.get("intent", "single_source"),
            source_types=data.get("source_types", ["any"]),
            topics=data.get("topics", []),
            ipc_sections=data.get("ipc_sections", []),
            time_filter=data.get("time_filter"),
            language_hint=data.get("language_hint", "en"),
            requires_compare=data.get("requires_compare", False),
            requires_summary=data.get("requires_summary", False),
            source_names=data.get("source_names", []),
            temporal_markers=data.get("temporal_markers", []),
            route=data.get("route", "rag")
        )

    except Exception as e:
        print(f"[QueryClassifier] Error: {e}")
        return QueryAnalysis(
            intent="single_source", 
            source_types=["any"], 
            topics=[], 
            ipc_sections=[], 
            time_filter=None, 
            language_hint="en", 
            requires_compare=False, 
            requires_summary=False, 
            source_names=[],
            temporal_markers=[],
            route="rag"
        )



def extract_source_filter(analysis: QueryAnalysis, available_source_ids: list[str] | None = None) -> list[str] | None:
    source_ids = []
    
    # 1. If explicitly named sources are found (e.g., "IPC", "Constitution")
    if analysis.source_names:
        try:
            with get_connection() as conn:
                cursor = conn.cursor(dictionary=True)
                for name in analysis.source_names:
                    query = "SELECT id FROM sources WHERE title LIKE %s"
                    cursor.execute(query, (f"%{name}%",))
                    rows = cursor.fetchall()
                    source_ids.extend([row['id'] for row in rows])
        except Exception as e:
            print(f"[QueryClassifier] DB Error in filter extraction (names): {e}")

    # 2. If source types are specified (e.g., "youtube", "web")
    if analysis.source_types and "any" not in analysis.source_types:
        try:
            type_map = {
                "youtube": "youtube",
                "web": "url",
                "legal_statute": "pdf",
                "legal_judgment": "pdf"
            }
            with get_connection() as conn:
                cursor = conn.cursor(dictionary=True)
                for stype in analysis.source_types:
                    db_type = type_map.get(stype)
                    if db_type:
                        cursor.execute("SELECT id FROM sources WHERE type = %s", (db_type,))
                        rows = cursor.fetchall()
                        source_ids.extend([row['id'] for row in rows])
        except Exception as e:
            print(f"[QueryClassifier] DB Error in filter extraction (types): {e}")

    # 3. Deduplicate and filter by available_source_ids if provided
    source_ids = list(set(source_ids))
    if available_source_ids and source_ids:
        source_ids = [sid for sid in source_ids if sid in available_source_ids]
        
    # If no specific filters found and "any" is allowed, return None for global search
    if not source_ids and "any" in analysis.source_types:
        return None
        
    return source_ids if source_ids else None