diff --git a/backend/app/api/endpoints/ai_chat.py b/backend/app/api/endpoints/ai_chat.py index e306ca8a..56ca3b88 100644 --- a/backend/app/api/endpoints/ai_chat.py +++ b/backend/app/api/endpoints/ai_chat.py @@ -172,6 +172,7 @@ async def post_message( session, data.content, db, action_type=data.action_type or "open_chat", focal_node_id=data.focal_node_id, + flow_context=data.flow_context, ) except Exception as e: logger.exception("AI chat message failed: %s", e) diff --git a/backend/app/core/ai_chat_service.py b/backend/app/core/ai_chat_service.py index c899cb54..c4c4ed75 100644 --- a/backend/app/core/ai_chat_service.py +++ b/backend/app/core/ai_chat_service.py @@ -533,13 +533,33 @@ async def send_message( db: AsyncSession, action_type: str = "open_chat", focal_node_id: str | None = None, + flow_context: dict | None = None, ) -> tuple[str, Optional[dict], Optional[str], Optional[dict]]: """Send a user message and get AI response. + Args: + flow_context: Live flow structure from the editor. Contains the current + tree_structure (troubleshooting) or steps + intake_form (procedural). + This gives the AI full awareness of the flow being edited. + Returns (ai_content, working_tree_update, new_phase, metadata_update). """ system_prompt = _build_system_prompt(session.flow_type) + # Inject live flow context so the AI can see current editor state + if flow_context: + context_json = json.dumps(flow_context, indent=2) + system_prompt += ( + f"\n\nCURRENT FLOW STATE (live from editor):\n{context_json}" + ) + if focal_node_id: + focal_node = _find_node_by_id(flow_context, focal_node_id) + if focal_node: + system_prompt += ( + f"\n\nFOCAL NODE/STEP (the item being acted on):\n" + f"{json.dumps(focal_node, indent=2)}" + ) + # Build messages array from conversation history now_iso = datetime.now(timezone.utc).isoformat() history = list(session.conversation_history) diff --git a/backend/app/schemas/ai_chat.py b/backend/app/schemas/ai_chat.py index 761f4064..01f40d88 100644 --- a/backend/app/schemas/ai_chat.py +++ b/backend/app/schemas/ai_chat.py @@ -43,6 +43,10 @@ class AIChatMessageRequest(BaseModel): default=None, description="ID of the node/step being acted on", ) + flow_context: Optional[dict[str, Any]] = Field( + default=None, + description="Live flow structure from the editor (tree structure, steps, intake form)", + ) class AIChatImportRequest(BaseModel): diff --git a/frontend/src/api/editorAI.ts b/frontend/src/api/editorAI.ts index a9d55800..8af9c752 100644 --- a/frontend/src/api/editorAI.ts +++ b/frontend/src/api/editorAI.ts @@ -6,6 +6,7 @@ export interface SendMessageParams { content: string actionType?: AIActionType focalNodeId?: string | null + flowContext?: Record | null } export const editorAIApi = { @@ -17,11 +18,12 @@ export const editorAIApi = { return data }, - sendMessage: async ({ sessionId, content, actionType, focalNodeId }: SendMessageParams) => { + sendMessage: async ({ sessionId, content, actionType, focalNodeId, flowContext }: SendMessageParams) => { const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/messages`, { content, action_type: actionType || 'open_chat', focal_node_id: focalNodeId, + flow_context: flowContext || undefined, }) return data }, diff --git a/frontend/src/hooks/useEditorAI.ts b/frontend/src/hooks/useEditorAI.ts index 57f6cf64..d5dc9fa1 100644 --- a/frontend/src/hooks/useEditorAI.ts +++ b/frontend/src/hooks/useEditorAI.ts @@ -10,9 +10,11 @@ import type { interface UseEditorAIOptions { flowType: 'troubleshooting' | 'procedural' treeId?: string | null + /** Returns the live flow structure from the editor for AI context */ + getFlowContext?: () => Record | null } -export function useEditorAI({ flowType, treeId }: UseEditorAIOptions) { +export function useEditorAI({ flowType, treeId, getFlowContext }: UseEditorAIOptions) { const [isOpen, setIsOpen] = useState(false) const [focalNodeId, setFocalNodeId] = useState(null) const [contextMenu, setContextMenu] = useState<{ @@ -95,6 +97,7 @@ export function useEditorAI({ flowType, treeId }: UseEditorAIOptions) { content: currentInput, actionType: currentAction, focalNodeId: currentFocalNodeId, + flowContext: getFlowContext?.() || null, }) setMessages((prev) => [ @@ -118,7 +121,7 @@ export function useEditorAI({ flowType, treeId }: UseEditorAIOptions) { setIsLoading(false) pendingActionRef.current = 'open_chat' } - }, [input, isLoading, ensureSession, focalNodeId]) + }, [input, isLoading, ensureSession, focalNodeId, getFlowContext]) const triggerAction = useCallback( (nodeId: string, actionType: AIActionType, prompt: string) => { diff --git a/frontend/src/pages/ProceduralEditorPage.tsx b/frontend/src/pages/ProceduralEditorPage.tsx index b807e08d..0874a1c4 100644 --- a/frontend/src/pages/ProceduralEditorPage.tsx +++ b/frontend/src/pages/ProceduralEditorPage.tsx @@ -53,6 +53,12 @@ export function ProceduralEditorPage() { const editorAI = useEditorAI({ flowType: 'procedural', treeId: id, + getFlowContext: useCallback(() => { + return { + steps: steps as unknown as Record[], + intake_form: intakeForm, + } + }, [steps, intakeForm]), }) const isMaintenance = treeType === 'maintenance' diff --git a/frontend/src/pages/TreeEditorPage.tsx b/frontend/src/pages/TreeEditorPage.tsx index 422413c9..b760b869 100644 --- a/frontend/src/pages/TreeEditorPage.tsx +++ b/frontend/src/pages/TreeEditorPage.tsx @@ -85,6 +85,10 @@ export function TreeEditorPage() { const editorAI = useEditorAI({ flowType: 'troubleshooting', treeId: id, + getFlowContext: useCallback(() => { + if (!treeStructure) return null + return treeStructure as unknown as Record + }, [treeStructure]), }) const previousEditingNodeRef = useRef(null)