Spaces:
Sleeping
Sleeping
Update agent.py
Browse files
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 |
-
|
| 89 |
-
|
| 90 |
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
|
| 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 |
|