Spaces:
Runtime error
Runtime error
File size: 3,080 Bytes
f085180 | 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 | # agent/query_analyzer.py
from typing import Dict, List, Union
import google.generativeai as genai
from .config import GEMINI_API_KEY
import logging
import time
import ast
# ─── Logging Config ─────────────────────────────────────────────
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class QueryAnalyzer:
def __init__(self):
self.model = genai.GenerativeModel("gemini-2.0-pro-exp-02-05")
def analyze_query(self, query: str) -> Dict[str, Union[str, List[List[str]]]]:
"""
Uses Gemini to extract structured web search intent and keyword clusters.
Returns:
Dict with Intent, Info Types, Time Range, KeywordChunks
"""
prompt = f"""
You are a web research planning assistant. Break down the user query into:
1. Intent: Describe what the user is trying to learn
2. Info Types: Type of data/sources needed
3. Time Range: Temporal scope (e.g., recent, past year)
4. KeywordChunks: Group 2–5 related search phrases into Python-style list of lists
User Query:
\"\"\"{query}\"\"\"
Format:
Intent: ...
Info Types: ...
Time Range: ...
KeywordChunks:
[["term1", "term2"], ["group2a", "group2b"]]
"""
try:
response = self.model.generate_content(prompt)
logger.info("Gemini query analysis succeeded")
time.sleep(5)
return self._parse_response(response.text.strip(), query)
except Exception as e:
logger.error(f"Query analysis failed: {e}")
return {
"Intent": "unknown",
"Info Types": "general",
"Time Range": "none",
"KeywordChunks": [[query]]
}
def _parse_response(self, response_text: str, fallback_query: str) -> Dict[str, Union[str, List[List[str]]]]:
result = {
"Intent": "unknown",
"Info Types": "general",
"Time Range": "none",
"KeywordChunks": [[fallback_query]]
}
lines = response_text.splitlines()
chunk_lines = []
inside_chunk = False
for line in lines:
if "KeywordChunks:" in line:
inside_chunk = True
continue
if inside_chunk:
chunk_lines.append(line)
elif ':' in line:
key, value = line.split(":", 1)
result[key.strip()] = value.strip()
if chunk_lines:
try:
chunk_text = "\n".join(chunk_lines).strip()
result["KeywordChunks"] = ast.literal_eval(chunk_text)
except Exception as e:
logger.warning(f"Failed to parse keyword chunks: {e}")
result["KeywordChunks"] = [[fallback_query]]
return result
|