Remove retriever from Agent
Browse files
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,
|