Prasanthkumar commited on
Commit
b9b310b
·
verified ·
1 Parent(s): 92c2a3f

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +21 -6
model.py CHANGED
@@ -183,13 +183,28 @@ def build_graph(provider: str = "groq"):
183
  """Assistant node"""
184
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
185
 
 
 
 
 
 
 
 
 
186
  def retriever(state: MessagesState):
187
- """Retriever node"""
188
- similar_question = vector_store.similarity_search(state["messages"][0].content)
189
- example_msg = HumanMessage(
190
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
191
- )
192
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
 
 
 
 
 
 
193
 
194
  builder = StateGraph(MessagesState)
195
  builder.add_node("retriever", retriever)
 
183
  """Assistant node"""
184
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
185
 
186
+ # def retriever(state: MessagesState):
187
+ # """Retriever node"""
188
+ # similar_question = vector_store.similarity_search(state["messages"][0].content)
189
+ # example_msg = HumanMessage(
190
+ # content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
191
+ # )
192
+ # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
193
+
194
  def retriever(state: MessagesState):
195
+ """Retriever node"""
196
+ query = state["messages"][0].content
197
+ similar_question = vector_store.similarity_search(query)
198
+
199
+ if not similar_question:
200
+ print("No similar documents found.")
201
+ return {"messages": [sys_msg] + state["messages"]}
202
+
203
+ example_msg = HumanMessage(
204
+ content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
205
+ )
206
+ return {"messages": [sys_msg] + state["messages"] + [example_msg]}
207
+
208
 
209
  builder = StateGraph(MessagesState)
210
  builder.add_node("retriever", retriever)