Update agent.py

agent.py

@@ -44,7 +44,7 @@ from transformers import BertTokenizer, BertModel
 import torch
 import torch.nn.functional as F
 from langchain.agents import initialize_agent, AgentType
-from langchain_openai import ChatOpenAI
+from langchain_community.chat_models import ChatOpenAI
 from langchain_community.tools import Tool
 
 import time
@@ -104,17 +104,12 @@ def modulus(a: int, b: int) -> int:
 
 
 @tool
-def calculator(data: dict) -> float:
-    """
-    Perform a calculation between two numbers.
-
-    Args:
-        data: A dictionary with keys 'a', 'b', and 'operation'
-    """
-    a = data.get("a")
-    b = data.get("b")
-    operation = data.get("operation", "").lower()
-
+def calculator_tool(inputs: dict):
+    """Perform mathematical operations based on the operation provided."""
+    a = inputs.get("a")
+    b = inputs.get("b")
+    operation = inputs.get("operation")
+    
     if operation == "add":
         return a + b
     elif operation == "subtract":
@@ -123,12 +118,12 @@ def calculator(data: dict) -> float:
         return a * b
     elif operation == "divide":
         if b == 0:
-            raise ValueError("Cannot divide by zero.")
+            return "Error: Division by zero"
         return a / b
     elif operation == "modulus":
         return a % b
     else:
-        raise ValueError(f"Unsupported operation: {operation}")
+        return "Unknown operation"
 
 @tool
 def wiki_search(query: str) -> str:
@@ -419,33 +414,19 @@ docs = [
 vector_store = FAISS.from_documents(docs, embedding_model)
 vector_store.save_local("faiss_index")
 
-# -----------------------------
-# 5. Query & Filter Results (optional preview)
-# -----------------------------
-query = "How many albums did Mercedes Sosa release between 2000 and 2009?"
-results = vector_store.similarity_search_with_score(query, k=5)
-threshold = 0.75
-filtered = [doc for doc, score in results if score < threshold]
-
-
-print("\n📊 Retrieved Documents with Similarity Scores:")
-filtered = []
-for doc, score in results:
-    print(f"🔢 Score: {score:.4f}")
-    print(f"📄 Content: {doc.page_content}")
-    if score < threshold:
-        filtered.append(doc)
-        print("✅ Accepted")
-    else:
-        print("❌ Rejected")
-    print("-" * 80)
-
 
 # -----------------------------
 # 6. Create LangChain Retriever Tool
 # -----------------------------
 retriever = vector_store.as_retriever()
 
+question_retriever_tool = create_retriever_tool(
+    retriever=retriever,
+    name="Question_Search",
+    description="A tool to retrieve documents related to a user's question."
+)
+
+
 # -------------------------------
 # Step 6: Create LangChain Tools
 # -------------------------------
@@ -463,9 +444,6 @@ wikiq_tool = wikidata_query
 # -------------------------------
 # Step 7: Create the Planner-Agent Logic
 # -------------------------------
-# Define the agent tool set
-from langchain.chat_models import ChatOpenAI
-from langchain.agents import initialize_agent, AgentType
 
 # Define the tools (as you've already done)
 tools = [wiki_tool, calc_tool, file_tool, web_tool, arvix_tool, youtube_tool, video_tool, analyze_tool, wikiq_tool]
@@ -510,14 +488,14 @@ def process_question(question):
     # Step 5: Execute task (with error handling)
     try:
         if task_type == "wiki_search":
-            response = wiki_search_tool(question)
+            response = wiki_tool(question)
         elif task_type == "math":
-            response = calculator_tool(question)
+            response = calc_tool(question)
         else:
             response = "Default answer logic"
         
         # Step 6: Final response formatting
-        final_response = generate_final_answer(state, {'wiki_search': response})
+        final_response = final_answer_tool(state, {'wiki_search': response})
         return final_response
     
     except Exception as e:
@@ -527,16 +505,11 @@ def process_question(question):
 
 # Run the process
 question = "How many albums did Mercedes Sosa release between 2000 and 2009?"
-response = process_question(question)
+response = agent.run(question)
 print("Final Response:", response)
 
 
 
-question_retriever_tool = create_retriever_tool(
-    retriever=retriever,
-    name="Question_Search",
-    description="A tool to retrieve documents related to a user's question."
-)
 
 
 
@@ -562,39 +535,116 @@ def retriever(state: MessagesState):
 
 
 
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    wiki_search,
-    web_search,
-    arvix_search,
-]
-
-
+# ----------------------------------------------------------------
+# LLM Loader
+# ----------------------------------------------------------------
 def get_llm(provider: str, config: dict):
     if provider == "google":
+        from langchain_google_genai import ChatGoogleGenerativeAI
         return ChatGoogleGenerativeAI(model=config["model"], temperature=config["temperature"])
+
     elif provider == "groq":
+        from langchain_groq import ChatGroq
         return ChatGroq(model=config["model"], temperature=config["temperature"])
+
     elif provider == "huggingface":
+        from langchain_huggingface import ChatHuggingFace
+        from langchain_huggingface import HuggingFaceEndpoint
         return ChatHuggingFace(
             llm=HuggingFaceEndpoint(url=config["url"], temperature=config["temperature"])
         )
+
     else:
         raise ValueError(f"Invalid provider: {provider}")
 
 
-def generate_final_answer(state, tools_results):
-    final_answer = ""
-    
-    # Concatenate results from each tool (wiki_search, calculator, etc.)
-    for tool_name, result in tools_results.items():
-        final_answer += f"{tool_name} result: {result}\n"
-    
-    return final_answer
+
+
+# ----------------------------------------------------------------
+# Planning & Execution Logic
+# ----------------------------------------------------------------
+def planner(question: str) -> list:
+    if "calculate" in question or any(op in question for op in ["add", "subtract", "multiply", "divide", "modulus"]):
+        return ["math"]
+    elif "wiki" in question or "who is" in question.lower():
+        return ["wiki_search"]
+    else:
+        return ["default"]
+
+
+def task_classifier(question: str) -> str:
+    if any(op in question.lower() for op in ["add", "subtract", "multiply", "divide", "modulus"]):
+        return "math"
+    elif "who" in question.lower() or "what is" in question.lower():
+        return "wiki_search"
+    else:
+        return "default"
+
+# Function to extract math operation from the question
+def extract_math_from_question(question: str):
+    """Extract numbers and operator from a math question."""
+    match = re.search(r'(\d+)\s*(\+|\-|\*|\/|\%)\s*(\d+)', question)
+    if match:
+        num1 = int(match.group(1))
+        operator = match.group(2)
+        num2 = int(match.group(3))
+        return num1, operator, num2
+    else:
+        return None
+
+def decide_task(state: dict) -> str:
+    return planner(state["question"])[0]
+
+
+def node_skipper(state: dict) -> bool:
+    return False
+
+
+def generate_final_answer(state: dict, task_results: dict) -> str:
+    if "wiki_search" in task_results:
+        return f"📚 Wiki Summary:\n{task_results['wiki_search']}"
+    elif "math" in task_results:
+        return f"🧮 Math Result: {task_results['math']}"
+    else:
+        return "🤖 Unable to generate a specific answer."
+
+
+# ----------------------------------------------------------------
+# Process Function (Main Agent Runner)
+# ----------------------------------------------------------------
+def process_question(question: str):
+    tasks = planner(question)
+    print(f"Tasks to perform: {tasks}")
+
+    task_type = task_classifier(question)
+    print(f"Task type: {task_type}")
+
+    state = {"question": question, "last_response": "", "messages": [HumanMessage(content=question)]}
+    next_task = decide_task(state)
+    print(f"Next task: {next_task}")
+
+    if node_skipper(state):
+        print(f"Skipping task: {next_task}")
+        return "Task skipped."
+
+    try:
+        if task_type == "wiki_search":
+            response = wiki_tool.run(question)
+        elif task_type == "math":
+            # You should dynamically parse these inputs in real use
+            response = calc_tool.run(question)
+        elif task_type == "retriever":
+            retrieval_result = retriever(state)
+            response = retrieval_result["messages"][-1].content
+        else:
+            response = "Default fallback answer."
+
+        return generate_final_answer(state, {task_type: response})
+
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        return "Sorry, I encountered an error processing your request."
+