Spaces:
Sleeping
Sleeping
Download scripts/lib/chat/query_planner.py from defenestrate/asa: direct link, hf CLI and curl.
- Browser
- Download file 5.94 kB
-
https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/chat/query_planner.py
- Command line
-
hf download hf://spaces/defenestrate/asa/scripts/lib/chat/query_planner.py
-
curl -L -o query_planner.py https://huggingface.co/spaces/defenestrate/asa/resolve/main/scripts/lib/chat/query_planner.py
5.94 kB
| """QueryPlanner: rewrites follow-ups into standalone retrieval queries + | |
| classifies intent. Uses Haiku 4.5 for cost. Falls back to verbatim user | |
| message on any failure (never a hard blocker). | |
| First-turn short-circuit: when history is empty, the user message is treated | |
| as already self-contained and we skip the Haiku call entirely. Saves cost + | |
| latency on what's typically the majority of opening turns. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| import re | |
| from dataclasses import dataclass | |
| from typing import Protocol | |
| logger = logging.getLogger(__name__) | |
| class LLMClient(Protocol): | |
| def complete(self, system: str, prompt: str) -> str: ... | |
| VALID_INTENTS = {"explain", "practice", "debug", "compare", "summarize"} | |
| # Practice-intent signals. Substring match (case-insensitive) on the user | |
| # message. If any matches, intent=practice regardless of what Haiku says | |
| # (and we can short-circuit Haiku entirely on the first turn). | |
| _PRACTICE_SIGNALS = ( | |
| "test me", | |
| "test my knowledge", | |
| "test my understanding", | |
| "quiz me", | |
| "challenge me", | |
| "ask me", | |
| "give me a problem", | |
| "give me a question", | |
| "give me an example", | |
| "let me practice", | |
| "let me try", | |
| "practice with", | |
| "practice on", | |
| "drill me", | |
| "work through", | |
| ) | |
| def _looks_like_practice(user_message: str) -> bool: | |
| msg = user_message.lower() | |
| return any(sig in msg for sig in _PRACTICE_SIGNALS) | |
| def _strip_code_fence(text: str) -> str: | |
| """Strip ```json ... ``` or ``` ... ``` wrapping if present.""" | |
| text = text.strip() | |
| if text.startswith("```"): | |
| text = re.sub(r"^```(?:json)?\s*", "", text) | |
| text = re.sub(r"\s*```\s*$", "", text) | |
| return text.strip() | |
| class PlannedQuery: | |
| rewritten: str | |
| intent: str | |
| is_follow_up: bool | |
| _PLANNER_SYSTEM_BASE = """You rewrite student chat messages into standalone retrieval queries for a course tutoring system. | |
| Course records describe concepts in declarative form (e.g., "Elasticity of demand measures the responsiveness of quantity demanded to price changes.") and use the course's term vocabulary drawn from its syntopicon. Whenever you rewrite, prefer that declarative shape and prefer the course's terms over student-vernacular synonyms — that's the form retrieval scores best against. | |
| Given the prior conversation and the latest user message, output a JSON object with three fields: | |
| - "intent": one of "explain", "practice", "debug", "compare", "summarize" | |
| - "is_follow_up": boolean (true if the message refers back to prior context) | |
| - "rewritten": a standalone retrieval query (1-2 sentences) that would surface the right course materials | |
| Rules: | |
| - If the latest message is self-contained, set is_follow_up=false and rewritten=verbatim. | |
| - If it's a follow-up like "why?" or "another example", expand it using the prior topic. | |
| - "explain" = conceptual question; "practice" = student wants to work a problem; "debug" = student stuck on something specific; "compare" = relating two concepts; "summarize" = give an overview. | |
| Return ONLY the JSON object, no commentary.""" | |
| def _build_system_prompt(term_names: list[str]) -> str: | |
| """Compose the planner system prompt; append the course-term registry | |
| when terms are available.""" | |
| if not term_names: | |
| return _PLANNER_SYSTEM_BASE | |
| terms_block = ( | |
| "Course terms (use these in rewrites when they cover the topic):\n" | |
| + ", ".join(term_names) | |
| ) | |
| return _PLANNER_SYSTEM_BASE + "\n\n" + terms_block | |
| class QueryPlanner: | |
| def __init__( | |
| self, | |
| llm: LLMClient, | |
| *, | |
| term_names: list[str] | None = None, | |
| ) -> None: | |
| self._llm = llm | |
| self._term_names = sorted(term_names) if term_names else [] | |
| self._system_prompt = _build_system_prompt(self._term_names) | |
| def plan(self, history: list[dict], user_message: str) -> PlannedQuery: | |
| """Rewrite + classify. Falls back to verbatim on any failure.""" | |
| # First-turn short-circuit: no history means nothing to rewrite from. | |
| # Use the user message as-is and default intent to "explain". Saves | |
| # a Haiku call + roundtrip on what's typically the majority of | |
| # opening turns. | |
| if not history: | |
| return PlannedQuery( | |
| rewritten=user_message, | |
| intent="practice" if _looks_like_practice(user_message) else "explain", | |
| is_follow_up=False, | |
| ) | |
| try: | |
| prompt = self._build_prompt(history, user_message) | |
| raw = self._llm.complete(self._system_prompt, prompt) | |
| parsed = json.loads(_strip_code_fence(raw)) | |
| intent = parsed.get("intent", "explain") | |
| if intent not in VALID_INTENTS: | |
| intent = "explain" | |
| # Heuristic override: explicit practice signals win over Haiku's call. | |
| if _looks_like_practice(user_message): | |
| intent = "practice" | |
| return PlannedQuery( | |
| rewritten=parsed.get("rewritten", user_message) or user_message, | |
| intent=intent, | |
| is_follow_up=bool(parsed.get("is_follow_up", False)), | |
| ) | |
| except Exception as e: | |
| logger.warning("QueryPlanner fallback (verbatim): %s", e) | |
| return PlannedQuery( | |
| rewritten=user_message, | |
| intent="practice" if _looks_like_practice(user_message) else "explain", | |
| is_follow_up=False, | |
| ) | |
| def _build_prompt(history: list[dict], user_message: str) -> str: | |
| lines = ["Prior conversation:"] | |
| for turn in history[-6:]: # last 3 user/assistant pairs | |
| role = turn.get("role", "?") | |
| content = turn.get("content", "") | |
| lines.append(f" {role}: {content}") | |
| lines.append(f"\nLatest user message: {user_message!r}") | |
| return "\n".join(lines) | |