Spaces:

mlops-pltf
/

hf-agents-final-assignment

Sleeping

Saumalya Sarkar

🚧 Trying to achieve 30% using supabase retriever node

6865997 over 1 year ago

7.72 kB

from typing import TypedDict, Annotated, Optional
from tools import multiply, sum, substract, divide\
    , wiki_search, web_search, arvix_search, gaia_database_search
    # , read_python, read_excel, transcribe_audio
from llm import chat_llm
from vector_store import get_supabase_client, get_gaia_data_vector_store, get_gaia_data_retriever_tool
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

tools = [
    multiply,
    sum,
    substract,
    divide,
    wiki_search,
    web_search,
    arvix_search,
    # read_python,
    # read_excel,
    # transcribe_audio,
    # gaia_database_search
]
llm_with_tools = chat_llm.bind_tools(
    tools
    , parallel_tool_calls=False
)

textual_description_of_tool="""
divide(a: int, b: int) -> float:
    Divide a and b
    Args:
        a: The numerator (int).
        b: The denominator (int).
    Returns:
        The result of the division (float).
multiply(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the multiplication (float).
sum(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the summation (float).
substract(a: float, b: float) -> float:
    multiply a and b
    Args:
        a: Number one.
        b: Number two.
    Returns:
        The result of the substraction (float).
wiki_search(query: str) -> str:
    Search Wikipedia for a query and return maximum 2 results.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
web_search(query: str) -> str:
    Search Tavily for a query and return maximum 3 results.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
arvix_search(query: str) -> str:
    Search Arxiv for a query and return maximum 3 result.
    Args:
        query: The Search Query.
    Returns:
        Formatted search result.
gaia_database_search(query: str) -> str:
    Search the Supabase vector store for documents relevant to the query.
    Args:
        query: The Search Query.
    Returns:
        Formatted string with the top results
"""
# SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
#     You can search internet, arvix and wikipedia and run computations with provided tools:\n{textual_description_of_tool} \n.
#     Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
#     - ALWAYS PASS THE HUMAN PROVIDED QUERY TO THE TOOLS, DO NOT SUMMARIZE THE QUESTION BEFORE SENDING TO TOOLS.
#     - For any kind of academic and educational question, use the tool for Arvix search.
#     - Always prioritize internet search tools to get an answer
#     - If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
#     Report your thoughts, and finish your answer with the following template:
#     FINAL ANSWER: [YOUR FINAL ANSWER]. 
#     YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
#     If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
#     If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
#     If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
#     Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
# """)
SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
    You can search internet, arvix and wikipedia and run computations with binded tools.
    Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
    - IT IS ABSOLUTELY MANDATORY TO CALL TOOLS WITH USER'S FULL  MESSAGE AS THE `QUERY` ARGUMENT.
    - If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
    Report your thoughts, and finish your answer with the following template:
    FINAL ANSWER: [YOUR FINAL ANSWER]. 
    YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
    If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
    If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
    If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
    Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
""")


class AgentState(TypedDict):
    # The document provided
    input_file: Optional[str]  # Contains file path (PDF/PNG)
    messages: Annotated[list[AnyMessage], add_messages]


def assistant(state: AgentState):
    # System message

    image=state["input_file"]

    response = llm_with_tools.invoke([SYS_MSG] + state["messages"])
    return {
        "messages": [response],
        "input_file": state["input_file"]
    }

def controlled_gaia_tool_node(state: AgentState):
    # Find the last Human/User message
    for msg in reversed(state["messages"]):
        if isinstance(msg, HumanMessage):
            raw_user_message = msg.content
            break
        # Or if using HumanMessage directly:
        # if isinstance(msg, HumanMessage):
        #     raw_user_message = msg.content
        #     break
    else:
        raw_user_message = ""
    # Now, call the tool directly
    vector_store = get_gaia_data_vector_store(get_supabase_client())
    similar_question = vector_store.similarity_search(raw_user_message)
    # Format result for LLM
    example_msg = HumanMessage(
        content=f"Here I provide the most relevant question and answer from GAIA database: \n\n{similar_question[0].page_content}",
    )
    return {
        "messages": state["messages"] + [example_msg],   # Keep previous + append
        "input_file": state["input_file"]
    }


# # The graph
# builder = StateGraph(AgentState)

# # Define nodes: these do the work
# builder.add_node("assistant", assistant)
# builder.add_node("tools", ToolNode(tools))

# # Define edges: these determine how the control flow moves
# builder.add_edge(START, "assistant")
# builder.add_conditional_edges(
#     "assistant",
#     # If the latest message requires a tool, route to tools
#     # Otherwise, provide a direct response
#     tools_condition,
# )
# builder.add_edge("tools", "assistant")
# react_graph = builder.compile()

# Show the butler's thought process
# display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))

# Graph with retriever
builder = StateGraph(AgentState)
builder.add_node("retriever", controlled_gaia_tool_node)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "retriever")
builder.add_edge("retriever", "assistant")
builder.add_conditional_edges(
    "assistant",
    tools_condition,
)
builder.add_edge("tools", "assistant")

# Compile graph
react_graph = builder.compile()

	from typing import TypedDict, Annotated, Optional
	from tools import multiply, sum, substract, divide\
	, wiki_search, web_search, arvix_search, gaia_database_search
	# , read_python, read_excel, transcribe_audio
	from llm import chat_llm
	from vector_store import get_supabase_client, get_gaia_data_vector_store, get_gaia_data_retriever_tool
	from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
	from langgraph.graph.message import add_messages
	from langgraph.graph import START, StateGraph
	from langgraph.prebuilt import ToolNode, tools_condition

	tools = [
	multiply,
	sum,
	substract,
	divide,
	wiki_search,
	web_search,
	arvix_search,
	# read_python,
	# read_excel,
	# transcribe_audio,
	# gaia_database_search
	]
	llm_with_tools = chat_llm.bind_tools(
	tools
	, parallel_tool_calls=False
	)

	textual_description_of_tool="""
	divide(a: int, b: int) -> float:
	Divide a and b
	Args:
	a: The numerator (int).
	b: The denominator (int).
	Returns:
	The result of the division (float).
	multiply(a: float, b: float) -> float:
	multiply a and b
	Args:
	a: Number one.
	b: Number two.
	Returns:
	The result of the multiplication (float).
	sum(a: float, b: float) -> float:
	multiply a and b
	Args:
	a: Number one.
	b: Number two.
	Returns:
	The result of the summation (float).
	substract(a: float, b: float) -> float:
	multiply a and b
	Args:
	a: Number one.
	b: Number two.
	Returns:
	The result of the substraction (float).
	wiki_search(query: str) -> str:
	Search Wikipedia for a query and return maximum 2 results.
	Args:
	query: The Search Query.
	Returns:
	Formatted search result.
	web_search(query: str) -> str:
	Search Tavily for a query and return maximum 3 results.
	Args:
	query: The Search Query.
	Returns:
	Formatted search result.
	arvix_search(query: str) -> str:
	Search Arxiv for a query and return maximum 3 result.
	Args:
	query: The Search Query.
	Returns:
	Formatted search result.
	gaia_database_search(query: str) -> str:
	Search the Supabase vector store for documents relevant to the query.
	Args:
	query: The Search Query.
	Returns:
	Formatted string with the top results
	"""
	# SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
	# You can search internet, arvix and wikipedia and run computations with provided tools:\n{textual_description_of_tool} \n.
	# Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
	# - ALWAYS PASS THE HUMAN PROVIDED QUERY TO THE TOOLS, DO NOT SUMMARIZE THE QUESTION BEFORE SENDING TO TOOLS.
	# - For any kind of academic and educational question, use the tool for Arvix search.
	# - Always prioritize internet search tools to get an answer
	# - If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
	# Report your thoughts, and finish your answer with the following template:
	# FINAL ANSWER: [YOUR FINAL ANSWER].
	# YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
	# If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
	# If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
	# If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
	# Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
	# """)
	SYS_MSG = SystemMessage(content=f"""You are a helpful butler named Jarvis that serves Mr. Sarkar.
	You can search internet, arvix and wikipedia and run computations with binded tools.
	Now, I will ask you a question. Always prioritize arvix search over web serach tool and then web search over wiki search whenever possible.
	- IT IS ABSOLUTELY MANDATORY TO CALL TOOLS WITH USER'S FULL MESSAGE AS THE `QUERY` ARGUMENT.
	- If answer is not found via internet search tools, use GAIA database search tool to look into provided knowledge
	Report your thoughts, and finish your answer with the following template:
	FINAL ANSWER: [YOUR FINAL ANSWER].
	YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
	If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
	If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
	If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
	Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
	""")


	class AgentState(TypedDict):
	# The document provided
	input_file: Optional[str] # Contains file path (PDF/PNG)
	messages: Annotated[list[AnyMessage], add_messages]


	def assistant(state: AgentState):
	# System message

	image=state["input_file"]

	response = llm_with_tools.invoke([SYS_MSG] + state["messages"])
	return {
	"messages": [response],
	"input_file": state["input_file"]
	}

	def controlled_gaia_tool_node(state: AgentState):
	# Find the last Human/User message
	for msg in reversed(state["messages"]):
	if isinstance(msg, HumanMessage):
	raw_user_message = msg.content
	break
	# Or if using HumanMessage directly:
	# if isinstance(msg, HumanMessage):
	# raw_user_message = msg.content
	# break
	else:
	raw_user_message = ""
	# Now, call the tool directly
	vector_store = get_gaia_data_vector_store(get_supabase_client())
	similar_question = vector_store.similarity_search(raw_user_message)
	# Format result for LLM
	example_msg = HumanMessage(
	content=f"Here I provide the most relevant question and answer from GAIA database: \n\n{similar_question[0].page_content}",
	)
	return {
	"messages": state["messages"] + [example_msg], # Keep previous + append
	"input_file": state["input_file"]
	}


	# # The graph
	# builder = StateGraph(AgentState)

	# # Define nodes: these do the work
	# builder.add_node("assistant", assistant)
	# builder.add_node("tools", ToolNode(tools))

	# # Define edges: these determine how the control flow moves
	# builder.add_edge(START, "assistant")
	# builder.add_conditional_edges(
	# "assistant",
	# # If the latest message requires a tool, route to tools
	# # Otherwise, provide a direct response
	# tools_condition,
	# )
	# builder.add_edge("tools", "assistant")
	# react_graph = builder.compile()

	# Show the butler's thought process
	# display(Image(react_graph.get_graph(xray=True).draw_mermaid_png()))

	# Graph with retriever
	builder = StateGraph(AgentState)
	builder.add_node("retriever", controlled_gaia_tool_node)
	builder.add_node("assistant", assistant)
	builder.add_node("tools", ToolNode(tools))
	builder.add_edge(START, "retriever")
	builder.add_edge("retriever", "assistant")
	builder.add_conditional_edges(
	"assistant",
	tools_condition,
	)
	builder.add_edge("tools", "assistant")

	# Compile graph
	react_graph = builder.compile()