Spaces:
Sleeping
Sleeping
| """ | |
| Query Understanding Module - Parses and analyzes research queries. | |
| """ | |
| import logging | |
| from typing import Optional, Dict, Any, List | |
| from ..models import QueryAnalysis, Entity, QueryComplexity | |
| from ..llm_client import llm_client | |
| from ..prompts.query_prompts import QUERY_PROMPTS | |
| logger = logging.getLogger(__name__) | |
| class QueryUnderstanding: | |
| """ | |
| Query Understanding module for analyzing and decomposing research queries. | |
| Implements FR-1: Query Understanding requirements. | |
| """ | |
| def __init__(self): | |
| self.llm = llm_client | |
| async def analyze_query(self, query: str) -> QueryAnalysis: | |
| """ | |
| Analyze a research query to understand intent, entities, and structure. | |
| Args: | |
| query: Raw natural language research query | |
| Returns: | |
| QueryAnalysis object with parsed query information | |
| """ | |
| logger.info(f"Analyzing query: {query[:100]}...") | |
| # First, validate the query | |
| validation = await self._validate_query(query) | |
| if not validation.get("proceed", True): | |
| logger.warning(f"Query validation failed: {validation}") | |
| # Create minimal analysis for invalid query | |
| analysis = QueryAnalysis(raw_query=query) | |
| analysis.sub_queries = [] | |
| return analysis | |
| # Analyze the query | |
| analysis_result = await self._analyze(query) | |
| # Extract entities | |
| entities = await self._extract_entities(query) | |
| # Classify intent | |
| intent_result = await self._classify_intent(query) | |
| # Build QueryAnalysis object | |
| analysis = self._build_analysis( | |
| query=query, | |
| analysis_result=analysis_result, | |
| entities=entities, | |
| intent_result=intent_result | |
| ) | |
| # Decompose into sub-queries if complex | |
| if analysis.complexity in [QueryComplexity.MEDIUM, QueryComplexity.COMPLEX]: | |
| sub_queries = await self._decompose_query(query, analysis_result) | |
| analysis.sub_queries = sub_queries | |
| else: | |
| analysis.sub_queries = [query] | |
| logger.info(f"Query analysis complete. Complexity: {analysis.complexity.value}, " | |
| f"Sub-queries: {len(analysis.sub_queries)}") | |
| return analysis | |
| async def _validate_query(self, query: str) -> Dict[str, Any]: | |
| """Validate if the query is researchable and appropriate.""" | |
| prompt = QUERY_PROMPTS["validation"].format(query=query) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| return result | |
| except Exception as e: | |
| logger.error(f"Query validation failed: {e}") | |
| return {"proceed": True, "is_valid": True} | |
| async def _analyze(self, query: str) -> Dict[str, Any]: | |
| """Perform main query analysis.""" | |
| prompt = QUERY_PROMPTS["analysis"].format(query=query) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| return result | |
| except Exception as e: | |
| logger.error(f"Query analysis failed: {e}") | |
| return { | |
| "intent": "unknown", | |
| "domain": "general", | |
| "entities": [], | |
| "complexity": "medium", | |
| "output_type": "report" | |
| } | |
| async def _extract_entities(self, query: str) -> List[Entity]: | |
| """Extract named entities from the query.""" | |
| prompt = QUERY_PROMPTS["entity_extraction"].format(query=query) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| entities = [] | |
| for entity_data in result.get("entities", []): | |
| entity = Entity( | |
| text=entity_data.get("text", ""), | |
| type=entity_data.get("type", "CONCEPT"), | |
| relevance=entity_data.get("relevance", "secondary"), | |
| context=entity_data.get("context") | |
| ) | |
| entities.append(entity) | |
| return entities | |
| except Exception as e: | |
| logger.error(f"Entity extraction failed: {e}") | |
| return [] | |
| async def _classify_intent(self, query: str) -> Dict[str, Any]: | |
| """Classify the query intent.""" | |
| prompt = QUERY_PROMPTS["intent_classification"].format(query=query) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| return result | |
| except Exception as e: | |
| logger.error(f"Intent classification failed: {e}") | |
| return { | |
| "primary_intent": "EXPLORATORY", | |
| "confidence": 0.5, | |
| "research_approach": "general" | |
| } | |
| async def _decompose_query( | |
| self, | |
| query: str, | |
| analysis: Dict[str, Any] | |
| ) -> List[str]: | |
| """Decompose a complex query into sub-queries.""" | |
| import json | |
| prompt = QUERY_PROMPTS["decomposition"].format( | |
| query=query, | |
| query_analysis=json.dumps(analysis, indent=2) | |
| ) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| sub_queries = [] | |
| for sq in result.get("sub_queries", []): | |
| sub_queries.append(sq.get("query", "")) | |
| # Ensure we have at least the original query | |
| if not sub_queries: | |
| sub_queries = [query] | |
| return sub_queries | |
| except Exception as e: | |
| logger.error(f"Query decomposition failed: {e}") | |
| return [query] | |
| async def check_clarity(self, query: str) -> Dict[str, Any]: | |
| """Check if the query needs clarification.""" | |
| prompt = QUERY_PROMPTS["clarification"].format(query=query) | |
| try: | |
| result = await self.llm.generate_json(prompt) | |
| return result | |
| except Exception as e: | |
| logger.error(f"Clarity check failed: {e}") | |
| return {"is_clear": True, "ambiguities": []} | |
| def _build_analysis( | |
| self, | |
| query: str, | |
| analysis_result: Dict[str, Any], | |
| entities: List[Entity], | |
| intent_result: Dict[str, Any] | |
| ) -> QueryAnalysis: | |
| """Build a QueryAnalysis object from component results.""" | |
| # Map complexity string to enum | |
| complexity_map = { | |
| "simple": QueryComplexity.SIMPLE, | |
| "medium": QueryComplexity.MEDIUM, | |
| "complex": QueryComplexity.COMPLEX | |
| } | |
| complexity_str = analysis_result.get("complexity", "medium").lower() | |
| complexity = complexity_map.get(complexity_str, QueryComplexity.MEDIUM) | |
| # Combine entities from analysis and extraction | |
| all_entities = entities.copy() | |
| for entity_data in analysis_result.get("entities", []): | |
| if isinstance(entity_data, dict): | |
| entity = Entity( | |
| text=entity_data.get("text", ""), | |
| type=entity_data.get("type", "CONCEPT"), | |
| relevance=entity_data.get("relevance", "secondary") | |
| ) | |
| # Avoid duplicates | |
| if not any(e.text == entity.text for e in all_entities): | |
| all_entities.append(entity) | |
| return QueryAnalysis( | |
| raw_query=query, | |
| intent=intent_result.get("primary_intent", analysis_result.get("intent", "")), | |
| domain=analysis_result.get("domain", "general"), | |
| entities=all_entities, | |
| temporal_scope=analysis_result.get("temporal_scope"), | |
| geographic_scope=analysis_result.get("geographic_scope"), | |
| complexity=complexity, | |
| output_type=analysis_result.get("output_type", "report") | |
| ) | |
| # Module instance | |
| query_understanding = QueryUnderstanding() | |