Separated guard and fallback logic into initial and final.

RAGManager/app/agents/graph.py

@@ -8,11 +8,15 @@
     fallback_final,
     fallback_inicial,
     generator,
-    guard,
+    guard_final,
+    guard_inicial,
     parafraseo,
     retriever,
 )
-from app.agents.routing import route_after_fallback_final, route_after_guard
+from app.agents.routing import (
+    route_after_guard_final,
+    route_after_guard_inicial,
+)
 from app.agents.state import AgentState
 
 
@@ -22,14 +26,15 @@ 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)
-    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)
+    2. agent_host -> guard_inicial (Nodo 2)
+    3. guard_inicial -> [conditional] -> fallback_inicial (if malicious) or parafraseo (if continue)
+    4. fallback_inicial -> END (error: breaks rules)
+    5. parafraseo -> retriever (Nodo 4)
+    6. retriever -> context_builder (Nodo 5)
+    7. context_builder -> generator (Nodo 6)
+    8. generator -> guard_final (Nodo 7)
+    9. guard_final -> [conditional] -> fallback_final (if risky) or END (if continue)
+    10. fallback_final -> END (error: classified info)
 
     Returns:
         Configured StateGraph instance ready for execution
@@ -39,33 +44,34 @@ def create_agent_graph() -> StateGraph:
 
     # Add nodes
     workflow.add_node("agent_host", agent_host)
-    workflow.add_node("guard", guard)
+    workflow.add_node("guard_inicial", guard_inicial)
     workflow.add_node("fallback_inicial", fallback_inicial)
     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("guard_final", guard_final)
     workflow.add_node("fallback_final", fallback_final)
 
     # Define edges
     # Start -> agent_host
     workflow.add_edge(START, "agent_host")
 
-    # agent_host -> guard
-    workflow.add_edge("agent_host", "guard")
+    # agent_host -> guard_inicial
+    workflow.add_edge("agent_host", "guard_inicial")
 
-    # guard -> conditional routing
+    # guard_inicial -> conditional routing
     workflow.add_conditional_edges(
-        "guard",
-        route_after_guard,
+        "guard_inicial",
+        route_after_guard_inicial,
         {
-            "malicious": END,  # End with error if malicious
-            "continue": "fallback_inicial",  # Continue to fallback_inicial if valid
+            "malicious": "fallback_inicial",  # Exception path: malicious content detected
+            "continue": "parafraseo",  # Normal path: continue processing
         },
     )
 
-    # fallback_inicial -> parafraseo
-    workflow.add_edge("fallback_inicial", "parafraseo")
+    # fallback_inicial -> END (stop flow with error message)
+    workflow.add_edge("fallback_inicial", END)
 
     # parafraseo -> retriever
     workflow.add_edge("parafraseo", "retriever")
@@ -77,19 +83,21 @@ def create_agent_graph() -> StateGraph:
     # Note: Primary LLM is called within context_builder node
     workflow.add_edge("context_builder", "generator")
 
-    # generator -> fallback_final
-    workflow.add_edge("generator", "fallback_final")
+    # generator -> guard_final
+    workflow.add_edge("generator", "guard_final")
 
-    # fallback_final -> conditional routing
+    # guard_final -> conditional routing
     workflow.add_conditional_edges(
-        "fallback_final",
-        route_after_fallback_final,
+        "guard_final",
+        route_after_guard_final,
         {
-            "risky": END,  # End with error if risky
-            "continue": END,  # End with final_response if valid
-            # Note: Final LLM is called within fallback_final node
+            "risky": "fallback_final",  # Exception path: risky content detected
+            "continue": END,  # Normal path: end successfully
         },
     )
 
+    # fallback_final -> END (stop flow with error message)
+    workflow.add_edge("fallback_final", END)
+
     # Compile the graph
     return workflow.compile()