veryfinal.py

import os, json
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Imports
from langchain_nvidia_ai_endpoints import ChatNVIDIA, 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?")