ORromu commited on
Commit
0f1bf73
·
verified ·
1 Parent(s): 519b3cd

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +35 -34
agent.py CHANGED
@@ -35,28 +35,28 @@ SUPABASE_SERVICE_ROLE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
35
  with open("prompt.txt", "r", encoding="utf-8") as f:
36
  system_prompt = f.read()
37
 
38
- # build a retriever
39
- embeddings = HuggingFaceEmbeddings(
40
- model_name="sentence-transformers/all-mpnet-base-v2"
41
- ) # dim=768
42
- supabase: Client = create_client(
43
- SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY
44
- )
45
- vector_store = SupabaseVectorStore(
46
- client=supabase,
47
- embedding=embeddings,
48
- table_name="documents2",
49
- query_name="match_documents_2",
50
- )
51
- create_retriever_tool = create_retriever_tool(
52
- retriever=vector_store.as_retriever(),
53
- name="Question Search",
54
- description="A tool to retrieve similar questions from a vector store.",
55
- )
56
-
57
  # System message
58
  sys_msg = SystemMessage(content=system_prompt)
59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
60
  # Loading the assistant
61
  chat = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
62
 
@@ -82,31 +82,32 @@ def simple_graph():
82
  ## Defining our nodes
83
  def assistant(state: MessagesState):
84
  """Assistant node"""
85
- return {"messages": [chat_with_tools.invoke(state["messages"])]}
86
 
87
- def retriever(state: MessagesState):
88
- """Retriever node"""
89
- similar_question = vector_store.similarity_search(state["messages"][0].content)
90
 
91
- if similar_question: # Check if the list is not empty
92
- example_msg = HumanMessage(
93
- content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
94
- )
95
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
96
- else:
97
- # Handle the case when no similar questions are found
98
- return {"messages": [sys_msg] + state["messages"]}
99
 
100
 
101
  # Build graph / nodes
102
  builder = StateGraph(MessagesState)
103
- builder.add_node("retriever", retriever) # Retriever
104
  builder.add_node("assistant", assistant) # Assistant
105
  builder.add_node("tools", ToolNode(tools)) # Tools
106
 
107
  # Logic / edges
108
- builder.add_edge(START, "retriever")
109
- builder.add_edge("retriever", "assistant")
 
110
  builder.add_conditional_edges("assistant", tools_condition)
111
  builder.add_edge("tools", "assistant")
112
 
 
35
  with open("prompt.txt", "r", encoding="utf-8") as f:
36
  system_prompt = f.read()
37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  # System message
39
  sys_msg = SystemMessage(content=system_prompt)
40
 
41
+ # # build a retriever
42
+ # embeddings = HuggingFaceEmbeddings(
43
+ # model_name="sentence-transformers/all-mpnet-base-v2"
44
+ # ) # dim=768
45
+ # supabase: Client = create_client(
46
+ # SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY
47
+ # )
48
+ # vector_store = SupabaseVectorStore(
49
+ # client=supabase,
50
+ # embedding=embeddings,
51
+ # table_name="documents2",
52
+ # query_name="match_documents_2",
53
+ # )
54
+ # create_retriever_tool = create_retriever_tool(
55
+ # retriever=vector_store.as_retriever(),
56
+ # name="Question Search",
57
+ # description="A tool to retrieve similar questions from a vector store.",
58
+ # )
59
+
60
  # Loading the assistant
61
  chat = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
62
 
 
82
  ## Defining our nodes
83
  def assistant(state: MessagesState):
84
  """Assistant node"""
85
+ return {"messages": [chat_with_tools.invoke([sys_msg] + state["messages"])]}
86
 
87
+ # def retriever(state: MessagesState):
88
+ # """Retriever node"""
89
+ # similar_question = vector_store.similarity_search(state["messages"][0].content)
90
 
91
+ # if similar_question: # Check if the list is not empty
92
+ # example_msg = HumanMessage(
93
+ # content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
94
+ # )
95
+ # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
96
+ # else:
97
+ # # Handle the case when no similar questions are found
98
+ # return {"messages": [sys_msg] + state["messages"]}
99
 
100
 
101
  # Build graph / nodes
102
  builder = StateGraph(MessagesState)
103
+ #builder.add_node("retriever", retriever) # Retriever
104
  builder.add_node("assistant", assistant) # Assistant
105
  builder.add_node("tools", ToolNode(tools)) # Tools
106
 
107
  # Logic / edges
108
+ # builder.add_edge(START, "retriever")
109
+ # builder.add_edge("retriever", "assistant")
110
+ builder.add_edge(START, "assistant")
111
  builder.add_conditional_edges("assistant", tools_condition)
112
  builder.add_edge("tools", "assistant")
113