| """QueryPlannerService — single LLM call: question + catalog → JSON IR. |
| |
| Prompt: src/config/prompts/query_planner.md (system) + the human content |
| built by `prompt.build_planner_prompt(...)`. |
| |
| Output: a QueryIR ready for the IRValidator. Validation + retry are the |
| caller's concern (`QueryService` orchestrates that loop). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| from langchain_core.messages import SystemMessage |
| from langchain_core.prompts import ChatPromptTemplate |
| from langchain_core.runnables import Runnable |
| from langchain_openai import AzureChatOpenAI |
|
|
| from src.middlewares.logging import get_logger |
|
|
| from ...catalog.models import Catalog |
| from ..ir.models import QueryIR |
| from .prompt import build_planner_prompt |
|
|
| logger = get_logger("query_planner") |
|
|
| _PROMPT_PATH = ( |
| Path(__file__).resolve().parent.parent.parent |
| / "config" |
| / "prompts" |
| / "query_planner.md" |
| ) |
|
|
|
|
| def _load_prompt_text() -> str: |
| return _PROMPT_PATH.read_text(encoding="utf-8") |
|
|
|
|
| def _build_default_chain() -> Runnable: |
| from src.config.settings import settings |
|
|
| llm = AzureChatOpenAI( |
| azure_deployment=settings.azureai_deployment_name_54m, |
| openai_api_version=settings.azureai_api_version_54m, |
| azure_endpoint=settings.azureai_endpoint_url_54m, |
| api_key=settings.azureai_api_key_54m, |
| temperature=0, |
| ) |
| prompt = ChatPromptTemplate.from_messages( |
| [ |
| SystemMessage(content=_load_prompt_text()), |
| ("human", "{human_content}"), |
| ] |
| ) |
| return prompt | llm.with_structured_output(QueryIR) |
|
|
|
|
| _default_chain: Runnable | None = None |
|
|
|
|
| def _get_default_chain() -> Runnable: |
| global _default_chain |
| if _default_chain is None: |
| _default_chain = _build_default_chain() |
| return _default_chain |
|
|
|
|
| class QueryPlannerService: |
| """Wraps the LLM call with structured-output parsing into QueryIR. |
| |
| Inject `structured_chain` for tests. The planner prompt is composed |
| by `build_planner_prompt(question, catalog, previous_error)` so retry |
| callers can append the prior error context to nudge the LLM. |
| """ |
|
|
| def __init__(self, structured_chain: Runnable | None = None) -> None: |
| self._chain = structured_chain |
|
|
| def _ensure_chain(self) -> Runnable: |
| if self._chain is None: |
| self._chain = _get_default_chain() |
| return self._chain |
|
|
| async def plan( |
| self, |
| question: str, |
| catalog: Catalog, |
| previous_error: str | None = None, |
| ) -> QueryIR: |
| human_content = build_planner_prompt(question, catalog, previous_error) |
| chain = self._ensure_chain() |
| ir: QueryIR = await chain.ainvoke({"human_content": human_content}) |
| logger.info( |
| "query planned", |
| source_id=ir.source_id, |
| table_id=ir.table_id, |
| select_n=len(ir.select), |
| filters_n=len(ir.filters), |
| retry=previous_error is not None, |
| ) |
| return ir |
|
|
