File size: 6,112 Bytes
b2b6341 fd1e711 b2b6341 fd1e711 b2b6341 fd1e711 b2b6341 fd1e711 b2b6341 fd1e711 b2b6341 fd1e711 b2b6341 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | 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
|