Update veryfinal.py

veryfinal.py

@@ -1,242 +1,242 @@
-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?")
+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?")