import os, json from dotenv import load_dotenv # Load environment variables load_dotenv() # Imports from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings from langchain_google_genai import ChatGoogleGenerativeAI from langchain_community.tools.tavily_search import TavilySearchResults from langchain_community.document_loaders import WikipediaLoader from langchain_community.document_loaders import ArxivLoader from langchain_community.vectorstores import FAISS from langchain_core.messages import SystemMessage, HumanMessage from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.tools import tool from langchain.tools.retriever import create_retriever_tool from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_community.document_loaders import JSONLoader from langgraph.prebuilt import create_react_agent from langgraph.checkpoint.memory import MemorySaver # Define all tools @tool def multiply(a: int | float, b: int | float) -> int | float: """Multiply two numbers. Args: a: first int | float b: second int | float """ return a * b @tool def add(a: int | float, b: int | float) -> int | float: """Add two numbers. Args: a: first int | float b: second int | float """ return a + b @tool def subtract(a: int | float , b: int | float) -> int | float: """Subtract two numbers. Args: a: first int | float b: second int | float """ return a - b @tool def divide(a: int | float, b: int | float) -> int | float: """Divide two numbers. Args: a: first int | float b: second int | float """ if b == 0: raise ValueError("Cannot divide by zero.") return a / b @tool def modulus(a: int | float, b: int | float) -> int | float: """Get the modulus of two numbers. Args: a: first int | float b: second int | float """ return a % b @tool def wiki_search(query: str) -> str: """Search the wikipedia for a query and return the first paragraph args: query: the query to search for """ loader = WikipediaLoader(query=query, load_max_docs=1) data = loader.load() formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.page_content}\n' for doc in data ]) return formatted_search_docs @tool def web_search(query: str) -> str: """Search Tavily for a query and return maximum 3 results. Args: query: The search query. """ search_docs = TavilySearchResults(max_results=3).invoke(query=query) formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.get("content", "")}\n' for doc in search_docs ]) return formatted_search_docs @tool def arxiv_search(query: str) -> str: """Search Arxiv for a query and return maximum 3 result. Args: query: The search query. """ search_docs = ArxivLoader(query=query, load_max_docs=3).load() formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.page_content[:1000]}\n' for doc in search_docs ]) return formatted_search_docs # Load and process your JSONL data jq_schema = """ { page_content: .Question, metadata: { task_id: .task_id, Level: .Level, Final_answer: ."Final answer", file_name: .file_name, Steps: .["Annotator Metadata"].Steps, Number_of_steps: .["Annotator Metadata"]["Number of steps"], How_long: .["Annotator Metadata"]["How long did this take?"], Tools: .["Annotator Metadata"].Tools, Number_of_tools: .["Annotator Metadata"]["Number of tools"] } } """ # Load documents and create vector database json_loader = JSONLoader(file_path="metadata.jsonl", jq_schema=jq_schema, json_lines=True, text_content=False) json_docs = json_loader.load() # Split documents text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=200) json_chunks = text_splitter.split_documents(json_docs) # Create vector database database = FAISS.from_documents(json_chunks, NVIDIAEmbeddings()) # Initialize LLM llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0) # Create retriever and retriever tool retriever = database.as_retriever(search_type="similarity", search_kwargs={"k": 3}) retriever_tool = create_retriever_tool( retriever=retriever, name="question_search", description="Search for similar questions and their solutions from the knowledge base." ) # Combine all tools tools = [ multiply, add, subtract, divide, modulus, wiki_search, web_search, arxiv_search, retriever_tool ] # Create memory for conversation memory = MemorySaver() # Create the agent agent_executor = create_react_agent( model=llm, tools=tools, checkpointer=memory ) # Function to run the agent def run_agent(query, thread_id="conversation_1"): """Run the agent with a query""" config = {"configurable": {"thread_id": thread_id}} system_msg = SystemMessage(content='''You are a helpful assistant tasked with answering questions using a set of tools. Now, I will ask you a question. 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.''') user_msg = HumanMessage(content=query) print(f"User: {query}") print("\nAgent:") for step in agent_executor.stream( {"messages": [system_msg, user_msg]}, config, stream_mode="values" ): step["messages"][-1].pretty_print() # Function to run agent with error handling def robust_agent_run(query, thread_id="robust_conversation"): """Run agent with error handling""" config = {"configurable": {"thread_id": thread_id}} try: system_msg = SystemMessage(content='''You are a helpful assistant tasked with answering questions using a set of tools. Now, I will ask you a question. 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.''') user_msg = HumanMessage(content=query) result = [] for step in agent_executor.stream( {"messages": [system_msg, user_msg]}, config, stream_mode="values" ): result = step["messages"] return result[-1].content if result else "No response generated" except Exception as e: return f"Error occurred: {str(e)}" # Main function def main(query: str) -> str: """Main function to run the agent""" return(robust_agent_run(query)) # Or use the interactive version # run_agent("What is 25 * 4 + 10?")