Update src/state.py

src/state.py

@@ -34,7 +34,7 @@ def build_graph(rag_chain):
 # For agent_chain.invoke
 def build_graph_with_callable(call_fn):
     def call_model(state: State):
-        response = call_fn(state)
+        response = call_fn({"input": state["input"]})
         return {
             "chat_history": [
                 HumanMessage(state["input"]),