ZeroTimo commited on
Commit
7ba0f28
·
verified ·
1 Parent(s): 25036c7

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +28 -4
agent.py CHANGED
@@ -180,12 +180,36 @@ def build_graph(provider: str = "google"):
180
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
181
 
182
  def retriever(state: MessagesState):
183
- """Retriever node"""
184
- similar_question = vector_store.similarity_search(state["messages"][0].content)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
185
  example_msg = HumanMessage(
186
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
 
 
 
 
 
 
 
187
  )
188
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
 
 
189
 
190
  builder = StateGraph(MessagesState)
191
  builder.add_node("retriever", retriever)
 
180
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
181
 
182
  def retriever(state: MessagesState):
183
+ current_question = state["messages"][-1].content
184
+ similar_doc = vector_store.similarity_search(current_question, k=1)[0]
185
+ similar_question = similar_doc.metadata["Question"]
186
+ similar_answer = similar_doc.page_content
187
+
188
+ # Wenn gleich: direkt beantworten
189
+ if current_question.strip() == similar_question.strip():
190
+ return {
191
+ "messages": state["messages"] + [
192
+ sys_msg,
193
+ HumanMessage(content=f"FINAL ANSWER: {similar_answer}")
194
+ ],
195
+ "should_end": True # überspringt Tool-Nutzung etc.
196
+ }
197
+
198
+ # Andernfalls Beispiel einbauen
199
  example_msg = HumanMessage(
200
+ content=f"""
201
+ Here I provide a similar question and answer for reference:
202
+
203
+ Question: {similar_question}
204
+ Final answer: {similar_answer}
205
+
206
+ If the current question is the same, return the same final answer.
207
+ """
208
  )
209
+ return {
210
+ "messages": [sys_msg] + state["messages"] + [example_msg],
211
+ "should_end": False # weitere Bearbeitung notwendig
212
+ }
213
 
214
  builder = StateGraph(MessagesState)
215
  builder.add_node("retriever", retriever)