Aleix Hernandez commited on
Commit
662bf8b
·
1 Parent(s): 5c9ba83

fix retriever function

Browse files
Files changed (1) hide show
  1. agent.py +11 -3
agent.py CHANGED
@@ -212,11 +212,19 @@ def build_graph(provider: str = "groq"):
212
  def retriever(state: MessagesState):
213
  """Retriever node"""
214
  similar_question = vector_store.similarity_search(state["messages"][0].content)
215
- example_msg = HumanMessage(
216
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
217
- )
 
 
 
 
 
 
 
218
  return {"messages": [sys_msg] + state["messages"] + [example_msg]}
219
 
 
220
  builder = StateGraph(MessagesState)
221
  builder.add_node("retriever", retriever)
222
  builder.add_node("assistant", assistant)
 
212
  def retriever(state: MessagesState):
213
  """Retriever node"""
214
  similar_question = vector_store.similarity_search(state["messages"][0].content)
215
+
216
+ if not similar_question:
217
+ example_msg = HumanMessage(
218
+ content="I couldn't find any similar question or answer in the database."
219
+ )
220
+ else:
221
+ example_msg = HumanMessage(
222
+ content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}"
223
+ )
224
+
225
  return {"messages": [sys_msg] + state["messages"] + [example_msg]}
226
 
227
+
228
  builder = StateGraph(MessagesState)
229
  builder.add_node("retriever", retriever)
230
  builder.add_node("assistant", assistant)