SimplicityGuy
diff --git a/‎api/nlq/actions.py‎
Lines changed: 117 additions & 0 deletions b/‎api/nlq/actions.py‎
Lines changed: 117 additions & 0 deletions
diff --git a/‎api/nlq/engine.py‎
Lines changed: 55 additions & 10 deletions b/‎api/nlq/engine.py‎
Lines changed: 55 additions & 10 deletions
diff --git a/‎api/nlq/suggestions.py‎
Lines changed: 98 additions & 0 deletions b/‎api/nlq/suggestions.py‎
Lines changed: 98 additions & 0 deletions
@@ -0,0 +1,117 @@
+"""NLQ action schemas and validation."""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+import structlog
+
+
+logger = structlog.get_logger(__name__)
+
+_MAX_FIELD_LEN = 256
+
+EntityType = Literal["artist", "label", "genre", "style", "release"]
+PaneName = Literal["explore", "trends", "insights", "genres", "credits"]
+FilterDimension = Literal["year", "genre", "label"]
+
+
+class _SeedEntity(BaseModel):
+    name: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]
+    entity_type: EntityType
+
+
+class SeedGraphAction(BaseModel):
+    type: Literal["seed_graph"] = "seed_graph"
+    entities: list[_SeedEntity]
+    replace: bool = False
+
+
+class HighlightPathAction(BaseModel):
+    type: Literal["highlight_path"] = "highlight_path"
+    nodes: list[Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]]
+
+
+class FocusNodeAction(BaseModel):
+    type: Literal["focus_node"] = "focus_node"
+    name: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]
+    entity_type: EntityType
+
+
+class FilterGraphAction(BaseModel):
+    type: Literal["filter_graph"] = "filter_graph"
+    by: FilterDimension
+    value: Annotated[str | int | tuple[int, int], Field()]
+
+
+class FindPathAction(BaseModel):
+    type: Literal["find_path"] = "find_path"
+    from_name: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN, alias="from")]
+    to_name: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN, alias="to")]
+    from_type: EntityType
+    to_type: EntityType
+
+    model_config = {"populate_by_name": True}
+
+
+class ShowCreditsAction(BaseModel):
+    type: Literal["show_credits"] = "show_credits"
+    name: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]
+    entity_type: EntityType
+
+
+class SwitchPaneAction(BaseModel):
+    type: Literal["switch_pane"] = "switch_pane"
+    pane: PaneName
+
+
+class OpenInsightTileAction(BaseModel):
+    type: Literal["open_insight_tile"] = "open_insight_tile"
+    tile_id: Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]
+
+
+class SetTrendRangeAction(BaseModel):
+    type: Literal["set_trend_range"] = "set_trend_range"
+    from_year: Annotated[str, Field(min_length=4, max_length=10, alias="from")]
+    to_year: Annotated[str, Field(min_length=4, max_length=10, alias="to")]
+
+    model_config = {"populate_by_name": True}
+
+
+class SuggestFollowupsAction(BaseModel):
+    type: Literal["suggest_followups"] = "suggest_followups"
+    queries: list[Annotated[str, Field(min_length=1, max_length=_MAX_FIELD_LEN)]]
+
+
+Action = Annotated[
+    SeedGraphAction
+    | HighlightPathAction
+    | FocusNodeAction
+    | FilterGraphAction
+    | FindPathAction
+    | ShowCreditsAction
+    | SwitchPaneAction
+    | OpenInsightTileAction
+    | SetTrendRangeAction
+    | SuggestFollowupsAction,
+    Field(discriminator="type"),
+]
+
+_action_adapter: TypeAdapter[Action] = TypeAdapter(Action)
+
+
+def parse_action(raw: dict[str, Any]) -> Action:
+    """Parse and validate a single action. Raises ValidationError on failure."""
+    return _action_adapter.validate_python(raw)
+
+
+def parse_action_list(raw: list[dict[str, Any]]) -> list[Action]:
+    """Parse a list of raw action dicts, dropping malformed entries with a warning."""
+    parsed: list[Action] = []
+    for item in raw:
+        try:
+            parsed.append(parse_action(item))
+        except ValidationError as exc:
+            logger.warning("⚠️ dropping malformed NLQ action", item=item, errors=exc.errors())
+    return parsed
@@ -9,14 +9,16 @@
 
 from dataclasses import dataclass, field
 import json
+import re
 from typing import TYPE_CHECKING, Any
 
 import structlog
 
+from api.nlq.actions import Action, parse_action_list
 from api.nlq.tools import get_authenticated_tool_schemas, get_public_tool_schemas
 
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from collections.abc import Awaitable, Callable
 
     from api.nlq.config import NLQConfig
@@ -49,6 +51,20 @@
 member biographies, music theory, concert schedules). Your knowledge comes \
 exclusively from the tools provided."""
 
+_ACTIONS_INSTRUCTION = """
+
+When your answer should mutate the UI, append a machine-readable actions block \
+at the end of your response, formatted exactly as:
+
+    <!--actions:[{"type":"seed_graph","entities":[{"name":"Kraftwerk","entity_type":"artist"}]}]-->
+
+The marker is invisible to the user; the UI strips it out. Supported action \
+types: seed_graph, highlight_path, focus_node, filter_graph, find_path, \
+show_credits, switch_pane, open_insight_tile, set_trend_range, suggest_followups. \
+Only emit actions that directly follow from the user's question."""
+
+_SYSTEM_PROMPT = _SYSTEM_PROMPT + _ACTIONS_INSTRUCTION
+
 _AUTH_ADDENDUM = """
 
 The user is logged in and has a Discogs collection. You can access their \
@@ -80,6 +96,7 @@ class NLQResult:
     summary: str
     entities: list[dict[str, Any]] = field(default_factory=list)
     tools_used: list[str] = field(default_factory=list)
+    actions: list[Action] = field(default_factory=list)
 
 
 # ── Engine ───────────────────────────────────────────────────────────────
@@ -162,26 +179,27 @@ async def run(
             messages.append({"role": "assistant", "content": response.content})
             if not tool_results:
                 # Non-tool stop (e.g., max_tokens) — treat as final response.
-                # Return directly without off-topic guardrail since the response
-                # was truncated, not intentionally off-topic.
                 summary = _extract_text(response)
-                deduped = _deduplicate_entities(entities)
-                return NLQResult(summary=summary, entities=deduped, tools_used=tools_used)
+                return self._build_result(summary, entities, tools_used)
             messages.append({"role": "user", "content": tool_results})
 
         # Max iterations reached — return whatever we have
         logger.warning("⚠️ NLQ engine reached max iterations", max_iterations=self._config.max_iterations)
         summary = _extract_text(response) if response else ""
         return self._build_result(summary, entities, tools_used)
 
-    def _build_result(self, summary: str, entities: list[dict[str, Any]], tools_used: list[str]) -> NLQResult:
-        """Build an NLQResult with guardrails and deduplication applied."""
-        # Off-topic guardrail: if no tools were used, check the response
+    def _build_result(
+        self,
+        summary: str,
+        entities: list[dict[str, Any]],
+        tools_used: list[str],
+    ) -> NLQResult:
+        """Build an NLQResult with guardrails, deduplication, and action extraction."""
         if not tools_used:
             summary = _apply_off_topic_guardrail(summary)
-
+        cleaned_summary, actions = _extract_actions(summary)
         deduped = _deduplicate_entities(entities)
-        return NLQResult(summary=summary, entities=deduped, tools_used=tools_used)
+        return NLQResult(summary=cleaned_summary, entities=deduped, tools_used=tools_used, actions=actions)
 
 
 # ── Helpers ──────────────────────────────────────────────────────────────
@@ -215,6 +233,33 @@ def _apply_off_topic_guardrail(summary: str) -> str:
     return summary
 
 
+_ACTIONS_MARKER_RE = re.compile(r"<!--actions:(\[.*?\])-->", re.DOTALL)
+
+
+def _extract_actions(text: str) -> tuple[str, list[Action]]:
+    """Strip an optional ``<!--actions:[...]-->`` marker from the agent's text.
+
+    The system prompt instructs the agent to append its structured UI actions
+    inside this marker at the end of its response. We strip the marker from the
+    user-visible summary and parse the JSON list.
+    """
+    match = _ACTIONS_MARKER_RE.search(text)
+    if not match:
+        return text, []
+    raw_json = match.group(1)
+    cleaned = _ACTIONS_MARKER_RE.sub("", text).strip()
+    try:
+        import json as _json  # noqa: PLC0415
+
+        raw = _json.loads(raw_json)
+        if not isinstance(raw, list):  # pragma: no cover — regex guarantees list-shaped JSON
+            return cleaned, []
+        return cleaned, parse_action_list(raw)
+    except (json.JSONDecodeError, ValueError):
+        logger.warning("⚠️ NLQ actions marker contained invalid JSON")
+        return cleaned, []
+
+
 def _deduplicate_entities(entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
     """Deduplicate entities by (id, type) tuple, preserving order."""
     seen: set[tuple[str, str]] = set()
 
@@ -0,0 +1,98 @@
+"""Template-based suggestion engine for the NLQ Ask pill."""
+
+from __future__ import annotations
+
+
+_MAX_FOCUS_LEN = 120
+_MAX_SUGGESTION_LEN = 256
+
+_DEFAULT_EXPLORE = [
+    "How are Kraftwerk and Afrika Bambaataa connected?",
+    "What genres emerged in the 1990s?",
+    "Most prolific electronic label",
+    "Show the shortest path from David Bowie to Daft Punk",
+]
+
+_DEFAULT_TRENDS = [
+    "Which labels grew the most in 2024?",
+    "Show the trend of techno releases over the last decade",
+    "Peak year for Detroit techno",
+    "Which genres are declining since 2020?",
+]
+
+_DEFAULT_INSIGHTS = [
+    "Biggest labels of 2024",
+    "Most connected artists overall",
+    "Top collaborators in electronic music",
+    "Rarest releases on Warp Records",
+]
+
+_DEFAULT_GENRES = [
+    "What genres split off from house in the 1990s?",
+    "Parent genre of jungle",
+    "Sub-genres of ambient",
+    "Genres that combine jazz and electronic",
+]
+
+_DEFAULT_CREDITS = [
+    "Who produced 'Computer World'?",
+    "Engineers credited on Kraftwerk releases",
+    "Writers who collaborated with Brian Eno",
+    "Vocalists credited on Massive Attack releases",
+]
+
+_PANE_DEFAULTS: dict[str, list[str]] = {
+    "explore": _DEFAULT_EXPLORE,
+    "trends": _DEFAULT_TRENDS,
+    "insights": _DEFAULT_INSIGHTS,
+    "genres": _DEFAULT_GENRES,
+    "credits": _DEFAULT_CREDITS,
+}
+
+_ARTIST_TEMPLATES = [
+    "Who influenced {focus}?",
+    "What labels has {focus} released on?",
+    "{focus}'s collaborators in the 70s",
+    "Most prolific decade for {focus}",
+    "How are {focus} and Kraftwerk connected?",
+]
+
+_LABEL_TEMPLATES = [
+    "Biggest artists on {focus}",
+    "Genres most associated with {focus}",
+    "Peak year for {focus}",
+    "Artists who moved from {focus} to a rival label",
+]
+
+_GENRE_TEMPLATES = [
+    "Who are the pioneers of {focus}?",
+    "Sub-genres of {focus}",
+    "Labels most associated with {focus}",
+    "How did {focus} evolve between 1990 and 2010?",
+]
+
+
+def build_suggestions(
+    *,
+    pane: str,
+    focus: str | None,
+    focus_type: str | None,
+) -> list[str]:
+    """Return 4-6 suggested queries for the given context."""
+    if focus is None or focus_type is None:
+        return _PANE_DEFAULTS.get(pane, _DEFAULT_EXPLORE)[:6]
+
+    focus_trimmed = focus.strip()[:_MAX_FOCUS_LEN]
+    if not focus_trimmed:
+        return _PANE_DEFAULTS.get(pane, _DEFAULT_EXPLORE)[:6]
+
+    templates = {
+        "artist": _ARTIST_TEMPLATES,
+        "label": _LABEL_TEMPLATES,
+        "genre": _GENRE_TEMPLATES,
+        "style": _GENRE_TEMPLATES,
+    }.get(focus_type, _ARTIST_TEMPLATES)
+
+    rendered = [t.format(focus=focus_trimmed) for t in templates]
+    capped = [q[:_MAX_SUGGESTION_LEN] for q in rendered]
+    return capped[:6]