"""
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