From 270c20912e60da2140ad6997af3f9d742fea422d Mon Sep 17 00:00:00 2001 From: chihlasm Date: Sat, 7 Mar 2026 00:17:39 -0500 Subject: [PATCH] feat: add delta response parsing and action-type prompt dispatch Co-Authored-By: Claude Opus 4.6 --- backend/app/core/ai_chat_service.py | 86 ++++++++++++++++++ backend/tests/test_ai_delta_response.py | 116 ++++++++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 backend/tests/test_ai_delta_response.py diff --git a/backend/app/core/ai_chat_service.py b/backend/app/core/ai_chat_service.py index 7c28382f..c899cb54 100644 --- a/backend/app/core/ai_chat_service.py +++ b/backend/app/core/ai_chat_service.py @@ -284,6 +284,92 @@ def _strip_markdown_fences(text: str) -> str: return text +def _parse_delta(response: str) -> dict | None: + """Extract [DELTA]...[/DELTA] JSON from AI response.""" + match = re.search(r'\[DELTA\](.*?)\[/DELTA\]', response, re.DOTALL) + if not match: + return None + raw = _strip_markdown_fences(match.group(1).strip()) + try: + return json.loads(raw) + except json.JSONDecodeError: + return None + + +def _find_node_by_id(tree: dict, node_id: str) -> dict | None: + """Find a node by ID in a tree structure (recursive).""" + if tree.get("id") == node_id: + return tree + for child in tree.get("children", []): + found = _find_node_by_id(child, node_id) + if found: + return found + for step in tree.get("steps", []): + if step.get("id") == node_id: + return step + return None + + +def _build_action_prompt( + action_type: str, + focal_node_id: str | None, + tree_structure: dict, + flow_type: str, +) -> str: + """Build action-specific system prompt supplement.""" + tree_json = json.dumps(tree_structure, indent=2) + + focal_context = "" + if focal_node_id: + focal_node = _find_node_by_id(tree_structure, focal_node_id) + if focal_node: + focal_context = f"\n\nFOCAL NODE (the node being acted on):\n{json.dumps(focal_node, indent=2)}" + + prompts = { + "generate_branch": ( + f"Generate a complete branch of child nodes for the focal node. " + f"Return the new nodes wrapped in [DELTA]...[/DELTA] markers as JSON with " + f"action='add', target_node_id='{focal_node_id}', and nodes array." + f"{focal_context}" + ), + "modify_node": ( + f"Modify the focal node based on the user's instruction. " + f"Return the updated node in [DELTA]...[/DELTA] markers with action='modify'." + f"{focal_context}" + ), + "add_steps": ( + f"Generate new procedural steps to insert after the focal step. " + f"Return them in [DELTA]...[/DELTA] markers with action='add'." + f"{focal_context}" + ), + "quick_action": ( + f"Respond to the user's quick action request about the focal node. " + f"If the action modifies the node, return changes in [DELTA]...[/DELTA] markers. " + f"If it's informational (e.g. explain), just respond in text." + f"{focal_context}" + ), + "open_chat": ( + "Have a helpful conversation about the flow. If the user asks for changes, " + "return them in [DELTA]...[/DELTA] markers. Otherwise respond in text." + ), + "generate_full": ( + "Generate a complete flow structure based on the user's description." + ), + "variable_inference": ( + "Analyze the procedural steps for implicit variables. Look for references to " + "specific servers, clients, credentials, or other values that should be captured " + "in an intake form. Return suggestions as JSON." + ), + } + + action_prompt = prompts.get(action_type, prompts["open_chat"]) + + return ( + f"CURRENT FLOW STRUCTURE ({flow_type}):\n{tree_json}\n\n" + f"ACTION: {action_type}\n{action_prompt}" + ) + + def _parse_ai_response(raw_response: str) -> dict[str, Any]: """Parse structured markers from AI response. diff --git a/backend/tests/test_ai_delta_response.py b/backend/tests/test_ai_delta_response.py new file mode 100644 index 00000000..d0fea499 --- /dev/null +++ b/backend/tests/test_ai_delta_response.py @@ -0,0 +1,116 @@ +"""Tests for AI delta response parsing and action-type prompt dispatch.""" +from app.core.ai_chat_service import _parse_delta, _build_action_prompt, _find_node_by_id + + +def test_parse_delta_from_response(): + """Service extracts [DELTA] markers from AI responses.""" + response = '''Here's a new branch for that node. + +[DELTA] +{"action": "add", "target_node_id": "check-dns", "nodes": [{"id": "verify-dns-server", "type": "decision", "question": "Is the DNS server responding?"}], "explanation": "Added DNS verification branch"} +[/DELTA] + +Let me know if you'd like to adjust this.''' + + parsed = _parse_delta(response) + assert parsed is not None + assert parsed["action"] == "add" + assert parsed["target_node_id"] == "check-dns" + assert len(parsed["nodes"]) == 1 + + +def test_parse_delta_none_when_absent(): + """Returns None when no delta marker present.""" + response = "Sure, I can explain that node. It checks connectivity." + parsed = _parse_delta(response) + assert parsed is None + + +def test_parse_delta_with_markdown_fences(): + """Handles delta JSON wrapped in markdown code fences.""" + response = '''[DELTA] +```json +{"action": "modify", "target_node_id": "node-1", "nodes": [{"id": "node-1", "type": "action", "title": "Updated"}], "explanation": "Modified title"} +``` +[/DELTA]''' + parsed = _parse_delta(response) + assert parsed is not None + assert parsed["action"] == "modify" + + +def test_parse_delta_invalid_json(): + """Returns None for invalid JSON inside delta markers.""" + response = "[DELTA]not valid json[/DELTA]" + parsed = _parse_delta(response) + assert parsed is None + + +def test_build_action_prompt_generate_branch(): + """Generate branch action includes focal node context.""" + tree = { + "id": "root", + "type": "decision", + "question": "Is the server up?", + "children": [], + "options": [], + } + prompt = _build_action_prompt( + action_type="generate_branch", + focal_node_id="root", + tree_structure=tree, + flow_type="troubleshooting", + ) + assert "root" in prompt + assert "generate" in prompt.lower() or "branch" in prompt.lower() + + +def test_build_action_prompt_open_chat(): + """Open chat action is general conversation.""" + prompt = _build_action_prompt( + action_type="open_chat", + focal_node_id=None, + tree_structure={"id": "root", "type": "decision"}, + flow_type="troubleshooting", + ) + assert isinstance(prompt, str) + assert len(prompt) > 0 + + +def test_find_node_by_id_root(): + """Finds root node.""" + tree = {"id": "root", "type": "decision", "children": []} + assert _find_node_by_id(tree, "root") is not None + + +def test_find_node_by_id_nested(): + """Finds nested child node.""" + tree = { + "id": "root", + "type": "decision", + "children": [ + {"id": "child-1", "type": "action", "children": []}, + {"id": "child-2", "type": "solution", "children": []}, + ], + } + found = _find_node_by_id(tree, "child-2") + assert found is not None + assert found["id"] == "child-2" + + +def test_find_node_by_id_not_found(): + """Returns None for non-existent node.""" + tree = {"id": "root", "type": "decision", "children": []} + assert _find_node_by_id(tree, "nonexistent") is None + + +def test_find_node_by_id_in_steps(): + """Finds node in procedural steps array.""" + tree = { + "steps": [ + {"id": "step-1", "type": "procedure_step"}, + {"id": "step-2", "type": "procedure_step"}, + ] + } + found = _find_node_by_id(tree, "step-2") + assert found is not None + assert found["id"] == "step-2"