vindruid commited on
Commit
4e1b432
·
1 Parent(s): 7540aaa

remove retriever

Browse files
Files changed (1) hide show
  1. gaia_agent.py +24 -25
gaia_agent.py CHANGED
@@ -123,21 +123,21 @@ Your answer should only start with "FINAL ANSWER: ", then follows with the answe
123
  """)
124
 
125
  # build a retriever
126
- embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") # dim=768
127
- supabase: Client = create_client(
128
- os.environ.get("SUPABASE_URL"),
129
- os.environ.get("SUPABASE_SERVICE_KEY"))
130
- vector_store = SupabaseVectorStore(
131
- client=supabase,
132
- embedding= embeddings,
133
- table_name="documents",
134
- query_name="match_documents_langchain",
135
- )
136
- create_retriever_tool = create_retriever_tool(
137
- retriever=vector_store.as_retriever(),
138
- name="Question Search",
139
- description="A tool to retrieve similar questions from a vector store.",
140
- )
141
 
142
 
143
 
@@ -169,20 +169,20 @@ def build_graph(provider: str = "groq"):
169
  """Assistant node"""
170
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
171
 
172
- def retriever(state: MessagesState):
173
- """Retriever node"""
174
- similar_question = vector_store.similarity_search(state["messages"][0].content)
175
- example_msg = HumanMessage(
176
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
177
- )
178
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
179
 
180
  builder = StateGraph(MessagesState)
181
- builder.add_node("retriever", retriever)
182
  builder.add_node("assistant", assistant)
183
  builder.add_node("tools", ToolNode(tools))
184
  builder.add_edge(START, "retriever")
185
- builder.add_edge("retriever", "assistant")
186
  builder.add_conditional_edges(
187
  "assistant",
188
  tools_condition,
@@ -197,7 +197,6 @@ class GaiaAgent:
197
  self.graph = build_graph(provider=provider)
198
 
199
  def __call__(self, question:str) -> str:
200
- print(f"Agent received question (first 50 chars): {question[:50]}...")
201
  # Wrap the question in a HumanMessage from langchain_core
202
  messages = [HumanMessage(content=question)]
203
  messages = self.graph.invoke({"messages": messages})
 
123
  """)
124
 
125
  # build a retriever
126
+ # embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") # dim=768
127
+ # supabase: Client = create_client(
128
+ # os.environ.get("SUPABASE_URL"),
129
+ # os.environ.get("SUPABASE_SERVICE_KEY"))
130
+ # vector_store = SupabaseVectorStore(
131
+ # client=supabase,
132
+ # embedding= embeddings,
133
+ # table_name="documents",
134
+ # query_name="match_documents_langchain",
135
+ # )
136
+ # create_retriever_tool = create_retriever_tool(
137
+ # retriever=vector_store.as_retriever(),
138
+ # name="Question Search",
139
+ # description="A tool to retrieve similar questions from a vector store.",
140
+ # )
141
 
142
 
143
 
 
169
  """Assistant node"""
170
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
171
 
172
+ # def retriever(state: MessagesState):
173
+ # """Retriever node"""
174
+ # similar_question = vector_store.similarity_search(state["messages"][0].content)
175
+ # example_msg = HumanMessage(
176
+ # content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
177
+ # )
178
+ # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
179
 
180
  builder = StateGraph(MessagesState)
181
+ # builder.add_node("retriever", retriever)
182
  builder.add_node("assistant", assistant)
183
  builder.add_node("tools", ToolNode(tools))
184
  builder.add_edge(START, "retriever")
185
+ # builder.add_edge("retriever", "assistant")
186
  builder.add_conditional_edges(
187
  "assistant",
188
  tools_condition,
 
197
  self.graph = build_graph(provider=provider)
198
 
199
  def __call__(self, question:str) -> str:
 
200
  # Wrap the question in a HumanMessage from langchain_core
201
  messages = [HumanMessage(content=question)]
202
  messages = self.graph.invoke({"messages": messages})