Merge remote-tracking branch 'origin/main' into structured-outputs

vintasoftware · Sep 13, 2024 · d8326c7 · d8326c7
2 parents 108327e + 2a2efc4
commit d8326c7
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 1,584 deletions.
diff --git a/django_ai_assistant/helpers/assistants.py b/django_ai_assistant/helpers/assistants.py
@@ -482,12 +482,17 @@ def setup(state: AgentState):
 
             return {"messages": messages}
 
+        def history(state: AgentState):
+            history = message_history.messages if message_history else []
+            return {"messages": [*history, HumanMessage(content=state["input"])]}
+
         def retriever(state: AgentState):
             if not self.has_rag:
                 return
 
             retriever = self.get_history_aware_retriever()
-            docs = retriever.invoke({"input": state["input"], "history": state["messages"]})
+            messages_without_input = state["messages"][:-1]
+            docs = retriever.invoke({"input": state["input"], "history": messages_without_input})
 
             document_separator = self.get_document_separator()
             document_prompt = self.get_document_prompt()
@@ -501,10 +506,6 @@ def retriever(state: AgentState):
                 f"\n\n---START OF CONTEXT---\n{formatted_docs}---END OF CONTEXT---\n\n"
             )
 
-        def history(state: AgentState):
-            history = message_history.messages if message_history else []
-            return {"messages": [*history, HumanMessage(content=state["input"])]}
-
         def agent(state: AgentState):
             response = llm_with_tools.invoke(state["messages"])
 
@@ -537,16 +538,16 @@ def record_response(state: AgentState):
         workflow = StateGraph(AgentState)
 
         workflow.add_node("setup", setup)
-        workflow.add_node("retriever", retriever)
         workflow.add_node("history", history)
+        workflow.add_node("retriever", retriever)
         workflow.add_node("agent", agent)
         workflow.add_node("tools", ToolNode(tools))
         workflow.add_node("respond", record_response)
 
         workflow.set_entry_point("setup")
-        workflow.add_edge("setup", "retriever")
-        workflow.add_edge("retriever", "history")
-        workflow.add_edge("history", "agent")
+        workflow.add_edge("setup", "history")
+        workflow.add_edge("history", "retriever")
+        workflow.add_edge("retriever", "agent")
         workflow.add_conditional_edges(
             "agent",
             tool_selector,