Spaces:
Sleeping
Sleeping
File size: 1,220 Bytes
1e732dd 696f787 1e732dd | 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 | """
MediGuard AI — Rewrite Query Node
Reformulates the user query to improve retrieval recall.
"""
from __future__ import annotations
import logging
from typing import Any
from src.services.agents.prompts import REWRITE_SYSTEM
logger = logging.getLogger(__name__)
def rewrite_query_node(state: dict, *, context: Any) -> dict:
"""Rewrite the original query for better retrieval."""
original = state.get("query", "")
patient_context = state.get("patient_context", "")
if context.tracer:
context.tracer.trace(name="rewrite_query_node", metadata={"query": original})
user_msg = f"Original query: {original}"
if patient_context:
user_msg += f"\n\nPatient context: {patient_context}"
try:
response = context.llm.invoke(
[
{"role": "system", "content": REWRITE_SYSTEM},
{"role": "user", "content": user_msg},
]
)
rewritten = response.content.strip()
if not rewritten:
rewritten = original
except Exception as exc:
logger.warning("Rewrite LLM failed: %s — keeping original query", exc)
rewritten = original
return {"rewritten_query": rewritten}
|