From d8312c24a58023998d3875801ad47fe9dee8b749 Mon Sep 17 00:00:00 2001 From: chihlasm Date: Tue, 24 Mar 2026 08:36:13 +0000 Subject: [PATCH] feat: add BranchAwarePromptBuilder with unit tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure function that assembles system prompt, cross-branch context, history, and images for _call_ai — no DB access, no LLM calls. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../services/branch_aware_prompt_builder.py | 58 +++++++++++++ .../tests/test_branch_aware_prompt_builder.py | 85 +++++++++++++++++++ 2 files changed, 143 insertions(+) create mode 100644 backend/app/services/branch_aware_prompt_builder.py create mode 100644 backend/tests/test_branch_aware_prompt_builder.py diff --git a/backend/app/services/branch_aware_prompt_builder.py b/backend/app/services/branch_aware_prompt_builder.py new file mode 100644 index 00000000..055b834f --- /dev/null +++ b/backend/app/services/branch_aware_prompt_builder.py @@ -0,0 +1,58 @@ +"""Branch-aware prompt builder — assembles AI context with cross-branch awareness. + +Pure function: takes data, returns dict matching _call_ai parameter names. +No DB access, no LLM calls. The caller pre-fetches all data. + +Return keys: system_base, rag_context, history, new_message, images +- system_base: stable system prompt (cached by Anthropic) +- rag_context: cross-branch summaries + attachment descriptions (NOT cached) +""" +from typing import Any + +from app.services.assistant_chat_service import ASSISTANT_SYSTEM_PROMPT + + +class BranchAwarePromptBuilder: + """Assembles prompt components for branch-aware AI calls.""" + + def build( + self, + branch_messages: list[dict[str, Any]], + sibling_summaries: str, + session_context: str, + attachments: list[dict[str, Any]], + new_message: str, + revival_context: str | None = None, + token_budget: int = 100_000, + ) -> dict[str, Any]: + """Build prompt components for _call_ai. + + Returns dict with keys: system_base, rag_context, history, new_message, images. + """ + # 1. system_base — stable, cached across turns + system_base = ASSISTANT_SYSTEM_PROMPT + "\n\n## Session Context\n" + session_context + + # 2. rag_context — changes per query, NOT cached + rag_parts = [] + if revival_context: + rag_parts.append(f"\n## Branch Revival\n{revival_context}") + if sibling_summaries: + rag_parts.append(sibling_summaries) + rag_context = "\n".join(rag_parts) + + # 3. history — branch messages filtered to user/assistant + history = [] + for msg in branch_messages: + if msg.get("role") in ("user", "assistant"): + history.append({"role": msg["role"], "content": msg["content"]}) + + # 4. images + images = attachments if attachments else None + + return { + "system_base": system_base, + "rag_context": rag_context, + "history": history, + "new_message": new_message, + "images": images, + } diff --git a/backend/tests/test_branch_aware_prompt_builder.py b/backend/tests/test_branch_aware_prompt_builder.py new file mode 100644 index 00000000..6b0541d7 --- /dev/null +++ b/backend/tests/test_branch_aware_prompt_builder.py @@ -0,0 +1,85 @@ +"""Unit tests for BranchAwarePromptBuilder — pure function, no DB needed.""" +import pytest +from app.services.branch_aware_prompt_builder import BranchAwarePromptBuilder + + +def test_build_basic(): + """Basic build with no cross-branch context.""" + builder = BranchAwarePromptBuilder() + result = builder.build( + branch_messages=[ + {"role": "user", "content": "DNS not resolving"}, + {"role": "assistant", "content": "Let's check DNS config"}, + ], + sibling_summaries="", + session_context="Problem: DNS resolution failure. Domain: networking.", + attachments=[], + new_message="I ran nslookup and got timeout", + ) + assert "system_base" in result + assert "rag_context" in result + assert "history" in result + assert "new_message" in result + assert "images" in result + assert result["new_message"] == "I ran nslookup and got timeout" + assert len(result["history"]) == 2 + + +def test_build_with_cross_branch_context(): + """Cross-branch summaries go into rag_context, not system_base.""" + builder = BranchAwarePromptBuilder() + sibling_ctx = "\n## Cross-Branch Context\n- **Network** [dead_end]: Network was fine." + result = builder.build( + branch_messages=[], + sibling_summaries=sibling_ctx, + session_context="Problem: test", + attachments=[], + new_message="test message", + ) + assert "Cross-Branch Context" in result["rag_context"] + assert "Cross-Branch Context" not in result["system_base"] + + +def test_build_with_images(): + """Image attachments are passed through.""" + builder = BranchAwarePromptBuilder() + result = builder.build( + branch_messages=[], + sibling_summaries="", + session_context="Problem: test", + attachments=[{"media_type": "image/png", "data": "base64data"}], + new_message="check this screenshot", + ) + assert len(result["images"]) == 1 + assert result["images"][0]["media_type"] == "image/png" + + +def test_build_with_revival_context(): + """Revival context is prepended to rag_context.""" + builder = BranchAwarePromptBuilder() + result = builder.build( + branch_messages=[], + sibling_summaries="", + session_context="Problem: test", + attachments=[], + new_message="test", + revival_context="New evidence: the error appears when VPN is active", + ) + assert "New evidence" in result["rag_context"] + + +def test_history_filters_to_user_assistant(): + """Only user and assistant messages appear in history.""" + builder = BranchAwarePromptBuilder() + result = builder.build( + branch_messages=[ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + {"role": "system", "content": "should be excluded"}, + ], + sibling_summaries="", + session_context="Problem: test", + attachments=[], + new_message="third", + ) + assert len(result["history"]) == 2