""" nodes.py Node definitions for the WalkXR LangGraph workflow. """ from __future__ import annotations import json import logging from typing import Any, Callable from walkxr_ai.core.state import AgentOutput, PromptContext, WalkState from walkxr_ai.rag.user_memory_store import UserMemoryStore RECENT_MESSAGE_WINDOW = 1 MAX_MEMORIES_PER_RUN = 5 logger = logging.getLogger(__name__) # Node: Retrieve Context (RAG) # ---------------------------- def build_retrieve_context_node(retrieval_engine, user_memory_store) -> Callable[[WalkState], dict]: """ Creates the LangGraph node for retrieving external context. This function "builds" the node by injecting the retrieval engine dependency into it. Why this exists: - LangGraph nodes must take only `state` as input - But retrieval requires an external dependency (retrieval_engine) - So we pre-configure the node here and return a ready-to-use function Returns: A function `retrieve_context(state)` that: - Reads `user_input` - Calls the retrieval engine - Writes to state and updates `phase` """ def retrieve_context(state: WalkState) -> dict: """ Retrieves relevant context for the current turn. Reads from state: - user_input - user_id - turn_index Writes to state: - retrieval_query - retrieved_context - phase """ query = state["user_input"] user_id = state.get("user_id", "").strip() is_new_session = state.get("turn_index", 0) == 0 # Retrieve general non-memory context from the shared RAG store knowledge_context = retrieval_engine.retrieve(query) # Only retrieve memory context at the start of a new session memory_context = [] if is_new_session and user_id: memory_results = user_memory_store.query_memories(user_id=user_id, query=query) # Convert user-memory retrieval results into the same dict-based shape used by retrieved_context elsewhere in the graph for result in memory_results: memory_context.append( { "text": result.node.get_text(), "source": "user_memory", "score": result.score, "metadata": result.node.metadata, } ) # Combine non-memory context with memory context to create the full retrieved context for this turn retrieved_context = knowledge_context + memory_context return { "retrieval_query": query, "retrieved_context": retrieved_context, "phase": "context_retrieved", } return retrieve_context # Node: Construct Prompt # ---------------------- def construct_prompt(state: WalkState) -> dict: """ Builds the structured prompt payload for the current turn. Reads from state: - user_input - conversation_history - retrieved_context - memory - current_step Writes to state: - prompt_context - phase """ recent_messages = state["conversation_history"][-RECENT_MESSAGE_WINDOW:] prompt_context = PromptContext( user_input=state["user_input"], recent_messages=recent_messages, memory=state["memory"], retrieved_context=state["retrieved_context"], current_step=state["current_step"], ) return { "prompt_context": prompt_context, "phase": "prompt_constructed", } ## Node: Call LLM ## -------------- def build_call_llm_node(responder: Any) -> Callable[[WalkState], dict]: """ Creates the LangGraph node for LLM invocation. The builder exists so we can inject an external responder dependency while still returning a standard LangGraph node that only accepts `state`. """ def call_llm(state: WalkState) -> dict: """ Calls the injected responder using the prepared prompt context. Reads from state: - prompt_context - bad_response_test (for testing purposes) - debug_print (for testing purposes) Writes to state on success: - agent_output - phase - bad_response_test (for testing purposes) Writes to state on failure: - error - phase """ try: result = responder.respond(state["prompt_context"]) if state["bad_response_test"] is not None: # Change response_text to custom response for testing result["response_text"] = state["bad_response_test"] agent_output = AgentOutput( response_text=result["response_text"], source_chunks=result.get("source_chunks", []), ) if state["debug_print"]: print("--- call_llm node ---") print("Prompt context:") print(json.dumps(state["prompt_context"], indent=2, default=str)) print("AI Response:") print(result["response_text"]) if state["bad_response_test"] is not None: return { "agent_output": agent_output, "phase": "response_generated", # Reset to None to allow for LLM response in the next loop "bad_response_test": None, } return { "agent_output": agent_output, "phase": "response_generated", } except Exception as exc: return { "error": str(exc), "phase": "error", } return call_llm # Node: Update Conversation History # --------------------------------- def update_conversation_history(state: WalkState) -> dict: """ Updates the conversation history with the latest user input and agent response. Reads from state: - conversation_history - user_input - agent_output Writes to state: - conversation_history - phase """ new_messages = [ {"role": "user", "content": state["user_input"]}, {"role": "assistant", "content": state["agent_output"]["response_text"]} ] updated_history = state["conversation_history"] + new_messages return { "conversation_history": updated_history, "phase": "history_updated", } # Node: Self-Correction # --------------------- def build_call_supervisor_node(responder: Any) -> Callable[[WalkState], dict]: """ Creates the LangGraph node for LLM supervision / reflection. The builder exists so we can inject an external responder dependency while still returning a standard LangGraph node that only accepts `state`. """ def reflect_on_response(state: WalkState) -> dict: """ Instructs an LLM of choice to act as a "supervisor" and evaluate the candidate response (given by the call_llm node) against the agent's constitution, checking for tone, safety, and helpfulness. Reads from state: - user_input - agent_output -> response_text - supervisor_retries - debug_print (for testing purposes) Writes to state on success: - supervisor_score - supervisor_retries - prompt_context -> feedback Writes to state on failure: - error - phase """ try: # Retrieve last user input and last AI response USER_INPUT = state["user_input"] CANDIDATE_RESPONSE = state["agent_output"]["response_text"] # Generate supervisor feedback supervisor_feedback_json = responder.reflect(USER_INPUT, CANDIDATE_RESPONSE) # Return new WalkState retries = state["supervisor_retries"] if supervisor_feedback_json["score"] != "good": retries += 1 if state["debug_print"]: print("--- call_supervisor node ---") print("Score:") print(supervisor_feedback_json["score"]) print("Feedback:") print(supervisor_feedback_json["rationale"]) return { "supervisor_score": supervisor_feedback_json["score"], "supervisor_retries": retries, "prompt_context": { **state["prompt_context"], # preserve existing keys "feedback": supervisor_feedback_json["rationale"] } } except Exception as exc: return { "error": str(exc), "phase": "error" } return reflect_on_response def fail_safe_response(state: WalkState) -> dict: """ Present a default fail-safe response in the event the LLM fails to generate a response that upholds the Agent Constitution after state["max_retries_allowed"] times. Writes to state: - agent_output -> response_text - phase - supervisor_score - prompt_context -> feedback """ safe_output = state["safe_response"] return { "agent_output": { **state["agent_output"], # preserve existing keys "response_text": safe_output }, "phase": "response_generated", "supervisor_score": "good", "prompt_context": { **state["prompt_context"], # preserve existing keys "feedback": "This is a default fail-safe response" } } # Node: Update Memory # ------------------- def build_update_memory_node(memory_extractor: Any) -> Callable[[WalkState], dict]: """ Creates the LangGraph node for extracting durable user memories and storing them in the persistent memory store. """ def update_memory(state: WalkState) -> dict: """ Extracts a few user memories from recent conversation history and stores them in the persistent user memory vector store. This node is meant to run at the end of a session (or for now, we will run it periodically every 10 turns). It should save only stable, cross-session information such as preferences, recurring facts, or meaningful ongoing insights. Reads from state: - user_id - conversation_history - walk_id - current_step Writes to state: - phase Writes to state on failure: - error - phase """ # Pull user_id from state and remove surrounding whitespace user_id = state.get("user_id", "").strip() # If user_id is missing, skip and exit the node gracefully instead of crashing if not user_id: logger.info("Skipping memory update: missing user_id in state.") return {"phase": "memory_updated"} # Pull conversation history from state conversation_history = state.get("conversation_history", []) # If there is no conversation history, there is nothing to extract memory from if not conversation_history: logger.info("Skipping memory update for user_id=%s: no conversation history.", user_id) return {"phase": "memory_updated"} # Limit memory extraction to the most recent messages to control token usage recent_messages = conversation_history[-RECENT_MESSAGE_WINDOW:] # Build a plain-text transcript of recent conversation messages transcript_lines = [] for message in recent_messages: role = message.get("role", "") content = message.get("content", "").strip() if not content: continue transcript_lines.append(f"{role.capitalize()}: {content}") if not transcript_lines: logger.info("Skipping memory update for user_id=%s: no usable conversation history.", user_id) return {"phase": "memory_updated"} transcript = "\n".join(transcript_lines) # Tell the LLM what kind of memories to extract system_prompt = ( "You extract only durable, useful, privacy-conscious cross-session memories. " "Keep memory extraction conservative. Avoid transient details, greetings, " "one-off plans, and sensitive personal information." ) prompt = f""" Review the recent conversation and extract up to {MAX_MEMORIES_PER_RUN} durable memories worth saving for future sessions. Only keep stable user preferences, recurring facts, or meaningful ongoing insights. Do not include filler, short-lived plans, or ephemeral details. Return ONLY valid JSON in this format: {{ "memories": [ "memory 1", "memory 2" ] }} If there is nothing worth saving, return: {{"memories": []}} Conversation: {transcript} """.strip() try: # Call the LLM to extract memories and parse the output raw_output = memory_extractor.extract_memories(system_prompt, prompt) parsed_output = json.loads(raw_output) memories = parsed_output.get("memories", []) # If "memories" key is missing or not a list, log a warning and exit gracefully if not isinstance(memories, list): logger.warning("update_memory received invalid memories payload.") return {"phase": "memory_updated"} # Filter extracted memories: keep only strings, remove empty/low-value items, deduplicate, cap the total number stored seen_memories = set() cleaned_memories = [] for memory in memories: if not isinstance(memory, str): continue memory_text = memory.strip() normalized = memory_text.lower() if not memory_text: continue # Very short fragments are usually low-value or too vague if len(memory_text.split()) < 3: continue # Avoid storing duplicate memories if normalized in seen_memories: continue seen_memories.add(normalized) cleaned_memories.append(memory_text) # Stop once we reach the configured cap if len(cleaned_memories) >= MAX_MEMORIES_PER_RUN: break # If nothing useful survives the cleaning step, skip storing if not cleaned_memories: logger.info("No durable memories extracted for user_id=%s.", user_id) return {"phase": "memory_updated"} # Create the persistent user memory store and save the cleaned memories with metadata memory_store = UserMemoryStore() metadata = {"memory_type": "extracted_memory",} if state.get("walk_id"): metadata["walk_id"] = state["walk_id"] if state.get("current_step"): metadata["current_step"] = state["current_step"] # Store each extracted memory as its own memory item in the vector store for memory in cleaned_memories: memory_store.add_memory(user_id=user_id, memory_text=memory, metadata=metadata) # Log how many memories were stored for observability/debugging logger.info( "Stored %s extracted memories for user_id=%s.", len(cleaned_memories), user_id, ) return {"phase": "memory_updated"} # If any step of the process fails, catch the exception, log it, and exit gracefully without crashing the entire workflow # The memory update is a "nice-to-have" that should not block the main user experience if it encounters issues except json.JSONDecodeError: logger.warning("update_memory received non-JSON memory extraction output.") return {"phase": "memory_updated"} except Exception as exc: logger.exception("Failed to update memory for user_id=%s", user_id) return {"error": str(exc), "phase": "error"} return update_memory # Node: Tool Use # -------------- # Placeholder for EPIC 4 (T004.2.2) # Node name: tool_node