Update veryfinal.py

veryfinal.py

@@ -7,7 +7,6 @@ load_dotenv()
 # Imports
 from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
 from langchain_groq import ChatGroq
-from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_nvidia_ai_endpoints import ChatNVIDIA
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
@@ -42,7 +41,57 @@ nvidia_rate_limiter = InMemoryRateLimiter(
     max_bucket_size=10
 )
 
-# Define all tools
+# Initialize individual LLMs
+groq_llm = ChatGroq(
+    model="llama-3.3-70b-versatile",
+    temperature=0,
+    api_key=os.getenv("GROQ_API_KEY"),
+    rate_limiter=groq_rate_limiter,
+    max_retries=2,
+    request_timeout=60
+)
+
+nvidia_llm = ChatNVIDIA(
+    model="meta/llama-3.1-405b-instruct",
+    temperature=0,
+    api_key=os.getenv("NVIDIA_API_KEY"),
+    rate_limiter=nvidia_rate_limiter,
+    max_retries=2
+)
+
+# Create LLM tools that can be selected by the agent
+@tool
+def groq_reasoning_tool(query: str) -> str:
+    """Use Groq's Llama model for fast reasoning, mathematical calculations, and logical problems.
+    Best for: Math problems, logical reasoning, quick calculations, code generation.
+    
+    Args:
+        query: The question or problem to solve
+    """
+    try:
+        time.sleep(random.uniform(1, 2))  # Rate limiting
+        response = groq_llm.invoke([HumanMessage(content=query)])
+        return f"Groq Response: {response.content}"
+    except Exception as e:
+        return f"Groq tool failed: {str(e)}"
+
+
+@tool
+def nvidia_specialist_tool(query: str) -> str:
+    """Use NVIDIA's large model for specialized tasks, technical questions, and domain expertise.
+    Best for: Technical questions, specialized domains, scientific problems, detailed analysis.
+    
+    Args:
+        query: The specialized question or technical problem
+    """
+    try:
+        time.sleep(random.uniform(2, 4))  # Rate limiting
+        response = nvidia_llm.invoke([HumanMessage(content=query)])
+        return f"NVIDIA Response: {response.content}"
+    except Exception as e:
+        return f"NVIDIA tool failed: {str(e)}"
+
+# Define calculation tools
 @tool
 def multiply(a: int | float, b: int | float) -> int | float:
     """Multiply two numbers.
@@ -94,6 +143,7 @@ def modulus(a: int | float, b: int | float) -> int | float:
     """
     return a % b
 
+# Define search tools
 @tool
 def wiki_search(query: str) -> str:
     """Search the wikipedia for a query and return the first paragraph
@@ -120,7 +170,6 @@ def web_search(query: str) -> str:
         query: The search query.
     """
     try:
-        # Add delay to prevent rate limiting
         time.sleep(random.uniform(1, 3))
         search_docs = TavilySearchResults(max_results=3).invoke(query=query)
         formatted_search_docs = "\n\n---\n\n".join(
@@ -179,58 +228,6 @@ json_chunks = text_splitter.split_documents(json_docs)
 # Create vector database
 database = FAISS.from_documents(json_chunks, NVIDIAEmbeddings())
 
-# Initialize LLMs with rate limiting
-def create_rate_limited_llm(provider="groq"):
-    """Create rate-limited LLM based on provider"""
-    
-    if provider == "groq":
-        return ChatGroq(
-            model="llama-3.3-70b-versatile",
-            temperature=0,
-            api_key=os.getenv("GROQ_API_KEY"),
-            rate_limiter=groq_rate_limiter,
-            max_retries=2,
-            request_timeout=60
-        )
-    elif provider == "google":
-        return ChatGoogleGenerativeAI(
-            model="gemini-2.0-flash-exp",
-            temperature=0,
-            api_key=os.getenv("GOOGLE_API_KEY"),
-            rate_limiter=google_rate_limiter,
-            max_retries=2,
-            timeout=60
-        )
-    elif provider == "nvidia":
-        return ChatNVIDIA(
-            model="meta/llama-3.1-405b-instruct",
-            temperature=0,
-            api_key=os.getenv("NVIDIA_API_KEY"),
-            rate_limiter=nvidia_rate_limiter,
-            max_retries=2
-        )
-
-# Create fallback chain with exponential backoff
-def create_llm_with_smart_fallbacks():
-    """Create LLM with intelligent fallback and rate limiting"""
-    
-    # Primary: Groq (fastest)
-    primary_llm = create_rate_limited_llm("groq")
-    
-    # Fallback 1: Google (most capable)
-    fallback_1 = create_rate_limited_llm("google")
-    
-    # Fallback 2: NVIDIA (reliable)
-    fallback_2 = create_rate_limited_llm("nvidia")
-    
-    # Create fallback chain
-    llm_with_fallbacks = primary_llm.with_fallbacks([fallback_1, fallback_2])
-    
-    return llm_with_fallbacks
-
-# Initialize LLM with smart fallbacks
-llm = create_llm_with_smart_fallbacks()
-
 # Create retriever and retriever tool
 retriever = database.as_retriever(search_type="similarity", search_kwargs={"k": 3})
 
@@ -240,47 +237,74 @@ retriever_tool = create_retriever_tool(
     description="Search for similar questions and their solutions from the knowledge base."
 )
 
-# Combine all tools
+# Combine all tools including LLM tools
 tools = [
+    # Math tools
     multiply,
     add,
     subtract,
     divide,
     modulus,
+    
+    # Search tools
     wiki_search,
     web_search,
     arxiv_search,
-    retriever_tool
+    retriever_tool,
+    
+    # LLM tools - agent can choose which LLM to use
+    groq_reasoning_tool,
+    nvidia_specialist_tool
 ]
 
+# Use a lightweight coordinator LLM (Groq for speed)
+coordinator_llm = ChatGroq(
+    model="llama-3.3-70b-versatile",
+    temperature=0,
+    api_key=os.getenv("GROQ_API_KEY"),
+    rate_limiter=groq_rate_limiter
+)
+
 # Create memory for conversation
 memory = MemorySaver()
 
-# Create the agent
+# Create the agent with coordinator LLM
 agent_executor = create_react_agent(
-    model=llm,
+    model=coordinator_llm,
     tools=tools,
     checkpointer=memory
 )
 
-# Enhanced robust agent run with exponential backoff
+# Enhanced robust agent run
 def robust_agent_run(query, thread_id="robust_conversation", max_retries=3):
-    """Run agent with error handling, rate limiting, and exponential backoff"""
+    """Run agent with error handling, rate limiting, and LLM tool selection"""
     
     for attempt in range(max_retries):
         try:
             config = {"configurable": {"thread_id": f"{thread_id}_{attempt}"}}
             
-            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.''')
+            system_msg = SystemMessage(content='''You are a helpful assistant with access to multiple specialized LLM tools and other utilities. 
+
+AVAILABLE LLM TOOLS:
+- groq_reasoning_tool: Fast reasoning, math, calculations, code (use for quick logical problems)
+- google_analysis_tool: Complex analysis, creative tasks, detailed explanations (use for comprehensive analysis)
+- nvidia_specialist_tool: Technical questions, specialized domains, scientific problems (use for expert-level tasks)
+
+TOOL SELECTION STRATEGY:
+- For math/calculations: Use basic math tools (add, multiply, etc.) OR groq_reasoning_tool for complex math
+- For factual questions: Use web_search, wiki_search, or arxiv_search first
+- For analysis/reasoning: Choose the most appropriate LLM tool based on complexity
+- For technical/scientific: Use nvidia_specialist_tool
+- For creative/comprehensive: Use google_analysis_tool
+- For quick logical problems: Use groq_reasoning_tool
+
+Always finish with: FINAL ANSWER: [YOUR FINAL ANSWER]
+Your answer should be a number OR few words OR comma separated list as appropriate.''')
             
             user_msg = HumanMessage(content=query)
             result = []
             
-            print(f"Attempt {attempt + 1}: Processing query...")
+            print(f"Attempt {attempt + 1}: Processing query with multi-LLM agent...")
             
             for step in agent_executor.stream(
                 {"messages": [system_msg, user_msg]}, 
@@ -296,9 +320,8 @@ def robust_agent_run(query, thread_id="robust_conversation", max_retries=3):
         except Exception as e:
             error_msg = str(e).lower()
             
-            # Check for rate limit errors
             if any(keyword in error_msg for keyword in ['rate limit', 'too many requests', '429', 'quota exceeded']):
-                wait_time = (2 ** attempt) + random.uniform(1, 3)  # Exponential backoff with jitter
+                wait_time = (2 ** attempt) + random.uniform(1, 3)
                 print(f"Rate limit hit on attempt {attempt + 1}. Waiting {wait_time:.2f} seconds...")
                 time.sleep(wait_time)
                 
@@ -306,7 +329,6 @@ def robust_agent_run(query, thread_id="robust_conversation", max_retries=3):
                     return f"Rate limit exceeded after {max_retries} attempts: {str(e)}"
                 continue
                 
-            # Check for other API errors
             elif any(keyword in error_msg for keyword in ['api', 'connection', 'timeout', 'service unavailable']):
                 wait_time = (2 ** attempt) + random.uniform(0.5, 1.5)
                 print(f"API error on attempt {attempt + 1}. Retrying in {wait_time:.2f} seconds...")
@@ -317,7 +339,6 @@ def robust_agent_run(query, thread_id="robust_conversation", max_retries=3):
                 continue
                 
             else:
-                # Non-recoverable error
                 return f"Error occurred: {str(e)}"
     
     return "Maximum retries exceeded"
@@ -327,7 +348,7 @@ request_count = 0
 last_request_time = time.time()
 
 def main(query: str) -> str:
-    """Main function to run the agent with request tracking"""
+    """Main function to run the multi-LLM agent"""
     global request_count, last_request_time
     
     current_time = time.time()
@@ -338,15 +359,15 @@ def main(query: str) -> str:
         last_request_time = current_time
     
     request_count += 1
-    print(f"Processing request #{request_count}")
+    print(f"Processing request #{request_count} with multi-LLM agent")
     
-    # Add small delay between requests to prevent overwhelming APIs
+    # Add delay between requests
     if request_count > 1:
         time.sleep(random.uniform(2, 5))
     
     return robust_agent_run(query)
 
 if __name__ == "__main__":
-    # Test the agent
+    # Test the multi-LLM agent
     result = main("What are the names of the US presidents who were assassinated?")
     print(result)