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) <noreply@anthropic.com>
This commit is contained in:
@@ -319,6 +319,7 @@ async def start_session(
|
|||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
input_tokens=input_tokens,
|
input_tokens=input_tokens,
|
||||||
output_tokens=output_tokens,
|
output_tokens=output_tokens,
|
||||||
|
branch_id=session.active_branch_id if session.is_branching else None,
|
||||||
)
|
)
|
||||||
db.add(step)
|
db.add(step)
|
||||||
|
|
||||||
@@ -421,6 +422,7 @@ async def process_response(
|
|||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
input_tokens=input_tokens,
|
input_tokens=input_tokens,
|
||||||
output_tokens=output_tokens,
|
output_tokens=output_tokens,
|
||||||
|
branch_id=session.active_branch_id if session.is_branching else None,
|
||||||
)
|
)
|
||||||
db.add(step)
|
db.add(step)
|
||||||
|
|
||||||
@@ -640,6 +642,7 @@ async def pickup_session(
|
|||||||
briefing_step = AISessionStep(
|
briefing_step = AISessionStep(
|
||||||
id=uuid.uuid4(),
|
id=uuid.uuid4(),
|
||||||
session_id=session.id,
|
session_id=session.id,
|
||||||
|
branch_id=session.active_branch_id if session.is_branching else None,
|
||||||
step_order=session.step_count,
|
step_order=session.step_count,
|
||||||
step_type="action",
|
step_type="action",
|
||||||
content={
|
content={
|
||||||
@@ -714,6 +717,7 @@ async def pickup_session(
|
|||||||
parsed=parsed,
|
parsed=parsed,
|
||||||
input_tokens=input_tokens,
|
input_tokens=input_tokens,
|
||||||
output_tokens=output_tokens,
|
output_tokens=output_tokens,
|
||||||
|
branch_id=session.active_branch_id if session.is_branching else None,
|
||||||
)
|
)
|
||||||
db.add(next_step)
|
db.add(next_step)
|
||||||
|
|
||||||
@@ -926,6 +930,7 @@ async def generate_status_update(
|
|||||||
step = AISessionStep(
|
step = AISessionStep(
|
||||||
id=uuid.uuid4(),
|
id=uuid.uuid4(),
|
||||||
session_id=session.id,
|
session_id=session.id,
|
||||||
|
branch_id=session.active_branch_id if session.is_branching else None,
|
||||||
step_order=session.step_count,
|
step_order=session.step_count,
|
||||||
step_type="status_update",
|
step_type="status_update",
|
||||||
content={
|
content={
|
||||||
@@ -1207,6 +1212,7 @@ def _create_step_from_parsed(
|
|||||||
parsed: dict[str, Any],
|
parsed: dict[str, Any],
|
||||||
input_tokens: int,
|
input_tokens: int,
|
||||||
output_tokens: int,
|
output_tokens: int,
|
||||||
|
branch_id: UUID | None = None,
|
||||||
) -> AISessionStep:
|
) -> AISessionStep:
|
||||||
"""Create an AISessionStep from parsed LLM output."""
|
"""Create an AISessionStep from parsed LLM output."""
|
||||||
step_type = parsed["type"]
|
step_type = parsed["type"]
|
||||||
@@ -1244,6 +1250,7 @@ def _create_step_from_parsed(
|
|||||||
return AISessionStep(
|
return AISessionStep(
|
||||||
id=uuid.uuid4(),
|
id=uuid.uuid4(),
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
|
branch_id=branch_id,
|
||||||
step_order=step_order,
|
step_order=step_order,
|
||||||
step_type=step_type if parsed["type"] != "resolution_suggestion" else "action",
|
step_type=step_type if parsed["type"] != "resolution_suggestion" else "action",
|
||||||
content=content,
|
content=content,
|
||||||
|
|||||||
@@ -81,6 +81,54 @@ async def send_chat_message(
|
|||||||
if session.status not in ("active", "paused"):
|
if session.status not in ("active", "paused"):
|
||||||
raise ValueError(f"Cannot send messages to a {session.status} session")
|
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
|
# Auto-title from first message if still default
|
||||||
if session.step_count == 0 and message.strip():
|
if session.step_count == 0 and message.strip():
|
||||||
session.title = _auto_title(message)
|
session.title = _auto_title(message)
|
||||||
|
|||||||
Reference in New Issue
Block a user