Update src/state.py

src/state.py

@@ -1 +1,29 @@
-# Placeholder for LangGraph state
+from langgraph.graph import START, StateGraph
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
+from typing import Sequence
+from typing_extensions import TypedDict, Annotated
+
+class State(TypedDict):
+    input: str
+    chat_history: Annotated[Sequence[BaseMessage], "add_messages"]
+    context: str
+    answer: str
+
+def build_graph(rag_chain):
+    def call_model(state: State):
+        response = rag_chain.invoke(state)
+        return {
+            "chat_history": [
+                HumanMessage(state["input"]),
+                AIMessage(response["answer"]),
+            ],
+            "context": response.get("context", ""),
+            "answer": response.get("answer", "")
+        }
+
+    workflow = StateGraph(state_schema=State)
+    workflow.add_node("model", call_model)
+    workflow.add_edge(START, "model")
+    memory = MemorySaver()
+    return workflow.compile(checkpointer=memory)