graph

app.py

@@ -87,98 +87,91 @@ _ = vector_store.add_documents(documents=training_documents)
 
 retriever = vector_store.as_retriever(search_kwargs={"k": 6})
 
+class AgentState(TypedDict):
+    messages: Annotated[list, "add_messages"]
+    question: str
+    context: List[Document]  # Para el RAG
+
+# ----------------- RAG Components -----------------
 def retrieve(state):
-  retrieved_docs = retriever.invoke(state["question"])
-  return {"context" : retrieved_docs}
+    retrieved_docs = retriever.invoke(state["question"])
+    return {"context": retrieved_docs}
 
 RAG_PROMPT = """\
 You are a helpful AI-powered Flu & Respiratory Illness Consultant. Your job is to help users determine whether they have the flu, a cold, RSV, or allergies based on their symptoms. 
 Provide recommendations based on the context provided. If symptoms are severe, advise the user to seek medical attention. 
 Avoid giving definitive diagnoses or prescriptions—always encourage users to consult a healthcare professional for serious cases.
+
 ### Question
 {question}
+
 ### Context
 {context}
 """
 
 rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
-
 llm = ChatOpenAI(model="gpt-4o")
 
 def generate(state):
-  docs_content = "\n\n".join(doc.page_content for doc in state["context"])
-  messages = rag_prompt.format_messages(question=state["question"], context=docs_content)
-  response = llm.invoke(messages)
-  return {"response" : response.content}
-
-from langgraph.graph import START, StateGraph
-from typing_extensions import List, TypedDict
-from langchain_core.documents import Document
-
-class State(TypedDict):
-  question: str
-  context: List[Document]
-  response: str
-
-graph_builder = StateGraph(State).add_sequence([retrieve, generate])
-graph_builder.add_edge(START, "retrieve")
-graph = graph_builder.compile()
-
+    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
+    messages = rag_prompt.format_messages(question=state["question"], context=docs_content)
+    response = llm.invoke(messages)
+    return {"messages": [response]}  
+# ----------------- Tools & Agent -----------------
 tavily_tool = TavilySearchResults(max_results=5)
-
-tool_belt = [
-    tavily_tool
-]
-
-model = ChatOpenAI(model="gpt-4o", temperature=0)
-model = model.bind_tools(tool_belt)
-
-
-class AgentState(TypedDict):
-  messages: Annotated[list, add_messages]
-  context: List[Document]
-
+tool_belt = [tavily_tool]
+model = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tool_belt)
 tool_node = ToolNode(tool_belt)
 
-uncompiled_graph = StateGraph(AgentState)
-
 def call_model(state):
-  messages = state["messages"]
-  response = model.invoke(messages)
-  return {
+    """Llama al modelo base para generar respuestas."""
+    messages = state["messages"]
+    response = model.invoke(messages)
+    return {
         "messages": [response],
+        "question": state["question"], 
         "context": state.get("context", [])
     }
 
-uncompiled_graph.add_node("agent", call_model)
+# ----------------- Create graph -----------------
+uncompiled_graph = StateGraph(AgentState)
+
+uncompiled_graph.add_node("retrieve", retrieve)
+uncompiled_graph.add_node("generate", generate)
 uncompiled_graph.add_node("action", tool_node)
 
-uncompiled_graph.set_entry_point("agent")
+uncompiled_graph.set_entry_point("retrieve")
 
+# ----------------- Logic -----------------
 def should_continue(state):
-  last_message = state["messages"][-1]
+    """Decide si usar herramientas después de `generate`."""
+    last_message = state["messages"][-1]
 
-  if last_message.tool_calls:
-    return "action"
+    if last_message.tool_calls:
+        return "action"  
 
-  return END
+    return END 
 
-uncompiled_graph.add_conditional_edges(
-    "agent",
-    should_continue
-)
 
-uncompiled_graph.add_edge("action", "agent")
+uncompiled_graph.add_edge("retrieve", "generate")
+uncompiled_graph.add_conditional_edges("generate", should_continue) 
+uncompiled_graph.add_edge("action", "generate") 
 
 compiled_graph = uncompiled_graph.compile()
 
+# ----------------- Chainlit Integration -----------------
 @cl.on_chat_start
 async def start():
-  cl.user_session.set("graph", compiled_graph)
+    cl.user_session.set("graph", compiled_graph)
 
 @cl.on_message
 async def handle(message: cl.Message):
-  graph = cl.user_session.get("graph")
-  state = {"messages" : [HumanMessage(content=message.content)]}
-  response = await graph.ainvoke(state)
-  await cl.Message(content=response["messages"][-1].content).send()
+    graph = cl.user_session.get("graph")
+    state = {
+        "messages": [HumanMessage(content=message.content)],
+        "question": message.content, 
+        "context": []  
+    }
+
+    response = await graph.ainvoke(state)
+    await cl.Message(content=response["messages"][-1].content).send()