fix al workflow (#26)

* fix al workflow: nodo condicional guard te lleva a fallback_inicial si malicious sino a parafraseo y el segundo guard si malicious te lleva a fallback_final sino termina

* Logica de fallback final y fallback inicial juntada en un unico fallback, no hace falta separarlos en 2

* Estado cambiado a MessageState para usar el historial de mensajes

* Copilot suggestions.

---------

Co-authored-by: JPAmorin <juanpabloamorinjusto@gmail.com>

Author: NicoWagner2005

RAGManager/app/agents/graph.py

@@ -5,16 +5,13 @@
 from app.agents.nodes import (
     agent_host,
     context_builder,
-    fallback_final,
-    fallback_inicial,
-    generator,
+    fallback,
     guard,
     parafraseo,
     retriever,
 )
-from app.agents.routing import route_after_fallback_final, route_after_guard
 from app.agents.state import AgentState
-
+from app.agents.routing import route_after_guard
 
 def create_agent_graph() -> StateGraph:
     """
@@ -23,13 +20,13 @@ def create_agent_graph() -> StateGraph:
     The graph implements the following flow:
     1. START -> agent_host (Nodo 1)
     2. agent_host -> guard (Nodo 2)
-    3. guard -> [conditional] -> fallback_inicial (Nodo 3) or END
-    4. fallback_inicial -> parafraseo (Nodo 4)
+    3. guard -> [conditional] -> fallback (Nodo 3) or END
+    4. fallback -> parafraseo (Nodo 4)
     5. parafraseo -> retriever (Nodo 5)
     6. retriever -> context_builder (Nodo 6)
     7. context_builder -> generator (Nodo 7)
-    8. generator -> fallback_final (Nodo 8)
-    9. fallback_final -> [conditional] -> END (with final_response) or END (with error)
+    8. generator -> fallback (Nodo 8)
+    9. fallback -> [conditional] -> END (with final_response) or END (with error)
 
     Returns:
         Configured StateGraph instance ready for execution
@@ -40,12 +37,10 @@ def create_agent_graph() -> StateGraph:
     # Add nodes
     workflow.add_node("agent_host", agent_host)
     workflow.add_node("guard", guard)
-    workflow.add_node("fallback_inicial", fallback_inicial)
+    workflow.add_node("fallback", fallback)
     workflow.add_node("parafraseo", parafraseo)
     workflow.add_node("retriever", retriever)
     workflow.add_node("context_builder", context_builder)
-    workflow.add_node("generator", generator)
-    workflow.add_node("fallback_final", fallback_final)
 
     # Define edges
     # Start -> agent_host
@@ -59,37 +54,29 @@ def create_agent_graph() -> StateGraph:
         "guard",
         route_after_guard,
         {
-            "malicious": END,  # End with error if malicious
-            "continue": "fallback_inicial",  # Continue to fallback_inicial if valid
+            "malicious": "fallback",  # go to fallback if malicious
+            "continue": "parafraseo",  # Continue to parafraseo if valid
         },
     )
 
-    # fallback_inicial -> parafraseo
-    workflow.add_edge("fallback_inicial", "parafraseo")
-
     # parafraseo -> retriever
     workflow.add_edge("parafraseo", "retriever")
 
     # retriever -> context_builder
     workflow.add_edge("retriever", "context_builder")
 
-    # context_builder -> generator
-    # Note: Primary LLM is called within context_builder node
-    workflow.add_edge("context_builder", "generator")
+    # context_builder -> guard
+    workflow.add_edge("context_builder", "guard")
 
-    # generator -> fallback_final
-    workflow.add_edge("generator", "fallback_final")
-
-    # fallback_final -> conditional routing
+    # guard -> conditional routing
     workflow.add_conditional_edges(
-        "fallback_final",
-        route_after_fallback_final,
+        "guard",
+        route_after_guard,
         {
-            "risky": END,  # End with error if risky
-            "continue": END,  # End with final_response if valid
-            # Note: Final LLM is called within fallback_final node
+            "malicious": "fallback",  # go to fallback if malicious
+            "continue": END,  # if there's no error ends
         },
     )
-
+    workflow.add_edge("fallback", END)
     # Compile the graph
     return workflow.compile()

RAGManager/app/agents/nodes/__init__.py

@@ -2,20 +2,16 @@
 
 from app.agents.nodes.agent_host import agent_host
 from app.agents.nodes.context_builder import context_builder
-from app.agents.nodes.fallback_final import fallback_final
-from app.agents.nodes.fallback_inicial import fallback_inicial
-from app.agents.nodes.generator import generator
+from app.agents.nodes.fallback import fallback
 from app.agents.nodes.guard import guard
 from app.agents.nodes.parafraseo import parafraseo
 from app.agents.nodes.retriever import retriever
 
 __all__ = [
     "agent_host",
     "guard",
-    "fallback_inicial",
+    "fallback",
     "parafraseo",
     "retriever",
     "context_builder",
-    "generator",
-    "fallback_final",
 ]

RAGManager/app/agents/nodes/agent_host.py

@@ -29,7 +29,8 @@ def agent_host(state: AgentState) -> AgentState:
         Updated state with chat_session_id, chat_messages, and initial_context set
     """
     updated_state = state.copy()
-    prompt = state.get("prompt", "")
+    
+    prompt = state["messages"][-1]
     chat_session_id = state.get("chat_session_id")
     user_id = state.get("user_id")
 

RAGManager/app/agents/nodes/context_builder.py

@@ -1,6 +1,10 @@
 """Nodo 6: Context Builder - Enriches query with retrieved context."""
 
 from app.agents.state import AgentState
+from langchain_core.messages import SystemMessage
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(model="gpt-5-nano")
 
 
 def context_builder(state: AgentState) -> AgentState:
@@ -31,13 +35,18 @@ def context_builder(state: AgentState) -> AgentState:
     paraphrased = state.get("paraphrased_text", "")
     chunks = state.get("relevant_chunks", [])
 
-    # Build enriched query
-    context_section = "\n\n".join(chunks) if chunks else ""
-    enriched_query = f"{paraphrased}\n\nContext:\n{context_section}" if context_section else paraphrased
-    updated_state["enriched_query"] = enriched_query
+    # Build enriched query with context
+    context_section = "\n\n".join(chunks) if chunks else "No relevant context found."
+    
+    system_content = f"""You are a helpful assistant. Use the following context to answer the user's question.
+If the answer is not in the context, say you don't know.
+
+Context:
+{context_section}"""
+
+    messages = [SystemMessage(content=system_content)] + state["messages"]
 
-    # TODO: Call Primary LLM here
-    # updated_state["primary_response"] = call_primary_llm(enriched_query)
-    updated_state["primary_response"] = None
+    # Call Primary LLM
+    response = llm.invoke(messages)
 
-    return updated_state
+    return {"messages": [response]}

RAGManager/app/agents/nodes/fallback.py

@@ -0,0 +1,42 @@
+"""Nodo 3: Fallback - Handles fallback processing from multiple workflow points."""
+
+import logging
+
+from app.agents.state import AgentState
+from langchain_core.messages import SystemMessage
+from langchain_openai import ChatOpenAI
+
+logger = logging.getLogger(__name__)
+
+llm = ChatOpenAI(
+    model="gpt-5-nano",
+)
+
+# TO DO: implementar clase nodo fallback y inicializar el llm en el init
+def fallback(state: AgentState) -> AgentState:
+    """
+    Fallback node - Performs fallback processing.
+
+    This node:
+    1. Alerts about malicious prompt
+    2. Generates an error_message from llm to show the user
+
+    Args:
+        state: Agent state containing the prompt or initial context
+
+    Returns:
+        error_message
+    """
+
+    logger.warning(
+        "Defensive check triggered: Malicious prompt detected" 
+    )
+    
+    messages = [
+        SystemMessage(
+            content="Your job is to generate an error message in user's language for the user explaining the database doesn't have the information to respond what the user asked"
+        )
+    ] + state["messages"]
+    error_message = llm.invoke(messages)
+    return {"messages": [error_message]}
+

RAGManager/app/agents/nodes/fallback_final.py

@@ -37,7 +37,9 @@ def guard(state: AgentState) -> AgentState:
         Updated state with is_malicious and error_message set
     """
     updated_state = state.copy()
-    prompt = state.get("prompt", "")
+    messages = state.get("messages", [])
+    last_message = messages[-1] if messages else None
+    prompt = last_message.content if last_message else ""
 
     if not prompt:
         # Empty prompt is considered safe
@@ -62,7 +64,7 @@ def guard(state: AgentState) -> AgentState:
             updated_state["error_message"] = (
                 "Jailbreak attempt detected. Your request contains content that violates security policies."
             )
-            logger.warning(f"Jailbreak attempt detected in prompt: {prompt[:100]}...")
+            logger.warning("Jailbreak attempt detected in prompt (len=%d)", len(prompt))
 
     except Exception as e:
         # If validation fails due to error, log it but don't block the request

RAGManager/app/agents/nodes/fallback_inicial.py

@@ -1,6 +1,10 @@
 """Nodo 4: Parafraseo - Paraphrases user input."""
 
 from app.agents.state import AgentState
+from langchain_core.messages import SystemMessage
+from langchain_openai import ChatOpenAI
+
+llm = ChatOpenAI(model="gpt-5-nano")
 
 
 def parafraseo(state: AgentState) -> AgentState:
@@ -24,9 +28,16 @@ def parafraseo(state: AgentState) -> AgentState:
     # 2. Improve clarity, adjust tone, or format as needed
     # 3. Set paraphrased_text with the result
 
-    # Placeholder: For now, we'll use the adjusted_text as-is
-    updated_state = state.copy()
-    text_to_paraphrase = state.get("adjusted_text") or state.get("prompt", "")
-    updated_state["paraphrased_text"] = text_to_paraphrase
+    # Paraphrase the last message using history
+    
+    system_instruction = """You are an expert at paraphrasing user questions to be standalone and clear, given the conversation history.
+Reformulate the last user message to be a self-contained query that includes necessary context from previous messages.
+Do not answer the question, just rewrite it."""
+
+    messages = [SystemMessage(content=system_instruction)] + state["messages"]
+    
+    response = llm.invoke(messages)
+    updated_state = state.copy()  # Create a copy of the state to update
+    updated_state["paraphrased_text"] = response.content
 
     return updated_state