| import json |
| import re |
|
|
| from app.core.logging import logger |
| from app.infrastructure.llm.base import LLMClientBase |
| from app.workflow.prompts import QUERY_ANALYSIS_PROMPT |
| from app.workflow.state import RAGState |
|
|
|
|
| def query_analysis_node(llm: LLMClientBase): |
| """Factory that creates query analysis node function.""" |
| async def node(state: RAGState) -> dict: |
| logger.info("Executing Node 1: Query Analysis...") |
| prompt = QUERY_ANALYSIS_PROMPT.format(question=state["question"]) |
| |
| try: |
| response = await llm.ainvoke([{"role": "user", "content": prompt}]) |
| |
| |
| cleaned_response = response.strip() |
| if cleaned_response.startswith("```"): |
| |
| cleaned_response = re.sub(r"^```(?:json)?\n", "", cleaned_response) |
| cleaned_response = re.sub(r"\n```$", "", cleaned_response) |
| cleaned_response = cleaned_response.strip() |
| |
| parsed = json.loads(cleaned_response) |
| rewritten = parsed.get("rewritten_query", state["question"]) |
| query_type = parsed.get("query_type", "conceptual") |
| except Exception as e: |
| logger.warning(f"Failed to parse query analysis output: {e}. Falling back to original question.") |
| rewritten = state["question"] |
| query_type = "conceptual" |
| |
| logger.info(f"Query Analysis complete: rewritten='{rewritten}', type='{query_type}'") |
| return { |
| "rewritten_query": rewritten, |
| "query_type": query_type |
| } |
| return node |
|
|
