Aleix Hernandez
commited on
Commit
·
662bf8b
1
Parent(s):
5c9ba83
fix retriever function
Browse files
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 |
-
|
| 216 |
-
|
| 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)
|