Upload veryfinal.py

veryfinal.py

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