From 7b4060a4d1dc2b5dd85096f67ff8605b2da16e14 Mon Sep 17 00:00:00 2001 From: chihlasm Date: Tue, 24 Mar 2026 08:41:22 +0000 Subject: [PATCH] feat: integrate branching into chat service and step creation Add is_branching guard to unified_chat_service.send_chat_message() that routes messages through BranchAwarePromptBuilder when a session has active branching. Add branch_id to all AISessionStep constructor calls in flowpilot_engine.py via optional branch_id param on _create_step_from_parsed. Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/app/services/flowpilot_engine.py | 7 +++ backend/app/services/unified_chat_service.py | 48 ++++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/backend/app/services/flowpilot_engine.py b/backend/app/services/flowpilot_engine.py index 50aa1711..49e665fb 100644 --- a/backend/app/services/flowpilot_engine.py +++ b/backend/app/services/flowpilot_engine.py @@ -319,6 +319,7 @@ async def start_session( parsed=parsed, input_tokens=input_tokens, output_tokens=output_tokens, + branch_id=session.active_branch_id if session.is_branching else None, ) db.add(step) @@ -421,6 +422,7 @@ async def process_response( parsed=parsed, input_tokens=input_tokens, output_tokens=output_tokens, + branch_id=session.active_branch_id if session.is_branching else None, ) db.add(step) @@ -640,6 +642,7 @@ async def pickup_session( briefing_step = AISessionStep( id=uuid.uuid4(), session_id=session.id, + branch_id=session.active_branch_id if session.is_branching else None, step_order=session.step_count, step_type="action", content={ @@ -714,6 +717,7 @@ async def pickup_session( parsed=parsed, input_tokens=input_tokens, output_tokens=output_tokens, + branch_id=session.active_branch_id if session.is_branching else None, ) db.add(next_step) @@ -926,6 +930,7 @@ async def generate_status_update( step = AISessionStep( id=uuid.uuid4(), session_id=session.id, + branch_id=session.active_branch_id if session.is_branching else None, step_order=session.step_count, step_type="status_update", content={ @@ -1207,6 +1212,7 @@ def _create_step_from_parsed( parsed: dict[str, Any], input_tokens: int, output_tokens: int, + branch_id: UUID | None = None, ) -> AISessionStep: """Create an AISessionStep from parsed LLM output.""" step_type = parsed["type"] @@ -1244,6 +1250,7 @@ def _create_step_from_parsed( return AISessionStep( id=uuid.uuid4(), session_id=session_id, + branch_id=branch_id, step_order=step_order, step_type=step_type if parsed["type"] != "resolution_suggestion" else "action", content=content, diff --git a/backend/app/services/unified_chat_service.py b/backend/app/services/unified_chat_service.py index e04054fc..b13d0e7f 100644 --- a/backend/app/services/unified_chat_service.py +++ b/backend/app/services/unified_chat_service.py @@ -81,6 +81,54 @@ async def send_chat_message( if session.status not in ("active", "paused"): raise ValueError(f"Cannot send messages to a {session.status} session") + # If branching is active, route to branch message handler + if session.is_branching and session.active_branch_id: + from app.services.branch_manager import BranchManager + from app.services.branch_aware_prompt_builder import BranchAwarePromptBuilder + from app.models.session_branch import SessionBranch + + branch_result = await db.execute( + select(SessionBranch).where(SessionBranch.id == session.active_branch_id) + ) + branch = branch_result.scalar_one_or_none() + if branch: + manager = BranchManager(db) + sibling_ctx = await manager.build_cross_branch_context(branch.id) + + builder = BranchAwarePromptBuilder() + session_context = f"Problem: {session.problem_summary or 'Unknown'}. Domain: {session.problem_domain or 'Unknown'}." + prompt_args = builder.build( + branch_messages=branch.conversation_messages, + sibling_summaries=sibling_ctx, + session_context=session_context, + attachments=[], + new_message=message, + revival_context=branch.evidence_description if branch.status == "revived" else None, + ) + + # Override images from prompt_args with actual images if provided + if images: + prompt_args["images"] = images + ai_content, input_tokens, output_tokens = await _call_ai(**prompt_args) + + # Update branch conversation + msgs = list(branch.conversation_messages or []) + msgs.append({"role": "user", "content": message}) + msgs.append({"role": "assistant", "content": ai_content}) + branch.conversation_messages = msgs + + session.total_input_tokens += input_tokens + session.total_output_tokens += output_tokens + session.step_count += 2 + + if session.status == "paused": + session.status = "active" + + suggested_flows = extract_suggested_flows( + await rag_search(query=message, account_id=account_id, db=db, limit=8) + ) + return ai_content, suggested_flows, session + # Auto-title from first message if still default if session.step_count == 0 and message.strip(): session.title = _auto_title(message)