budivoy commited on
Commit
19ced36
·
verified ·
1 Parent(s): 1b4bea8

Remove retriever from Agent

Browse files
Files changed (1) hide show
  1. agents.py +0 -31
agents.py CHANGED
@@ -13,8 +13,6 @@ from langchain_community.document_loaders import ArxivLoader
13
  from langchain_community.vectorstores import SupabaseVectorStore
14
  from langchain_core.messages import SystemMessage, HumanMessage
15
  from langchain_core.tools import tool
16
- from langchain.tools.retriever import create_retriever_tool
17
- from supabase.client import Client, create_client
18
 
19
  load_dotenv()
20
 
@@ -133,24 +131,6 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
133
  # System message
134
  sys_msg = SystemMessage(content=system_prompt)
135
 
136
- # build a retriever
137
- embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") # dim=768
138
- supabase: Client = create_client(
139
- os.environ.get("SUPABASE_URL"),
140
- os.environ.get("SUPABASE_SERVICE_KEY"))
141
- vector_store = SupabaseVectorStore(
142
- client=supabase,
143
- embedding= embeddings,
144
- table_name="documents",
145
- query_name="match_documents_langchain",
146
- )
147
- create_retriever_tool = create_retriever_tool(
148
- retriever=vector_store.as_retriever(),
149
- name="Question Search",
150
- description="A tool to retrieve similar questions from a vector store.",
151
- )
152
-
153
-
154
 
155
  tools = [
156
  multiply,
@@ -190,21 +170,10 @@ def build_graph(provider: str = "groq"):
190
  def assistant(state: MessagesState):
191
  """Assistant node"""
192
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
193
-
194
- def retriever(state: MessagesState):
195
- """Retriever node"""
196
- similar_question = vector_store.similarity_search(state["messages"][0].content)
197
- example_msg = HumanMessage(
198
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
199
- )
200
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
201
 
202
  builder = StateGraph(MessagesState)
203
- builder.add_node("retriever", retriever)
204
  builder.add_node("assistant", assistant)
205
  builder.add_node("tools", ToolNode(tools))
206
- builder.add_edge(START, "retriever")
207
- builder.add_edge("retriever", "assistant")
208
  builder.add_conditional_edges(
209
  "assistant",
210
  tools_condition,
 
13
  from langchain_community.vectorstores import SupabaseVectorStore
14
  from langchain_core.messages import SystemMessage, HumanMessage
15
  from langchain_core.tools import tool
 
 
16
 
17
  load_dotenv()
18
 
 
131
  # System message
132
  sys_msg = SystemMessage(content=system_prompt)
133
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
134
 
135
  tools = [
136
  multiply,
 
170
  def assistant(state: MessagesState):
171
  """Assistant node"""
172
  return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
 
 
 
 
 
 
 
173
 
174
  builder = StateGraph(MessagesState)
 
175
  builder.add_node("assistant", assistant)
176
  builder.add_node("tools", ToolNode(tools))
 
 
177
  builder.add_conditional_edges(
178
  "assistant",
179
  tools_condition,