Update agent.py

agent.py

@@ -42,7 +42,7 @@ from io import StringIO
 from transformers import BertTokenizer, BertModel
 import torch
 import torch.nn.functional as F
-
+import time
 
 load_dotenv()
 
@@ -333,81 +333,8 @@ for task in tasks:
 
 
 # -------------------------------
-# Step 1: Define the planner function
-# -------------------------------
-def planner(question: str):
-    """Break down the question into smaller tasks"""
-    if "how many" in question and "albums" in question:
-        return ["Retrieve album list", "Filter by date", "Count albums"]
-    elif "who" in question:
-        return ["Retrieve biography", "Find related works"]
-    return ["Default task"]
-
-# -------------------------------
-# Step 2: Task Classifier (decides the best tool to use)
-# -------------------------------
-def task_classifier(question: str):
-    """Classify the question to select the best tool"""
-    if "calculate" in question or any(op in question for op in ["+", "-", "*", "/"]):
-        return "math"
-    elif "album" in question or "music" in question:
-        return "wiki_search"
-    elif "file" in question or "attachment" in question:
-        return "file_analysis"
-    return "default_tool"
-
-# -------------------------------
-# Step 3: Decide Task Function
-# -------------------------------
-def decide_task(state):
-    """Logic to decide what to do based on prior actions and results"""
-    if "no relevant documents" in state.get("last_response", ""):
-        return "web_search"
-    if "not found" in state.get("last_response", "").lower():
-        return "wiki_search"
-    return "final_answer"
-
-# -------------------------------
-# Step 4: Node Skipper (Skip unnecessary nodes)
+# Step 4: Set up BERT Embeddings and FAISS VectorStore
 # -------------------------------
-def node_skipper(state):
-    """Skip unnecessary nodes based on context"""
-    if "just generate" in state.get("question", "").lower():
-        return "answer_generation"  # Skip all tools and just generate the answer.
-    return None  # Continue to the next tool or node
-
-
-
-# -------------------------------
-# Step 4: Set up HuggingFace Embeddings and FAISS VectorStore
-# -------------------------------
-# Initialize HuggingFace Embedding model
-#embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-#embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
-
-
-
-class BERTEmbeddings(Embeddings):
-    def __init__(self, model_name='bert-base-uncased'):
-        self.tokenizer = BertTokenizer.from_pretrained(model_name)
-        self.model = BertModel.from_pretrained(model_name)
-        self.model.eval()  # Set to evaluation mode
-
-    def embed_documents(self, texts: List[str]) -> List[List[float]]:
-        return self._embed(texts)
-
-    def embed_query(self, text: str) -> List[float]:
-        return self._embed([text])[0]
-
-    def _embed(self, texts: List[str]) -> List[List[float]]:
-        inputs = self.tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
-        return embeddings.cpu().numpy().tolist()
-
-# Example usage of BERTEmbedding with LangChain
-
 
 # -----------------------------
 # 1. Define Custom BERT Embedding Model
@@ -517,6 +444,7 @@ agent = initialize_agent(
 # -------------------------------
 # Step 8: Use the Planner, Classifier, and Decision Logic
 # -------------------------------
+
 def process_question(question):
     # Step 1: Planner generates the task sequence
     tasks = planner(question)
@@ -537,18 +465,31 @@ def process_question(question):
         print(f"Skipping to {skip}")
         return skip  # Or move directly to generating answer
 
-    # Execute the task via the agent
-    response = agent.run(question)
-    return response
+    # Step 5: Execute task (with error handling)
+    try:
+        if task_type == "wiki_search":
+            response = wiki_search_tool(question)
+        elif task_type == "math":
+            response = calculator_tool(question)
+        else:
+            response = "Default answer logic"
+        
+        # Step 6: Final response formatting
+        final_response = generate_final_answer(state, {'wiki_search': response})
+        return final_response
+    
+    except Exception as e:
+        print(f"Error executing task: {e}")
+        return "Sorry, I encountered an error processing your request."
 
-# -------------------------------
-# Step 9: Run the Planner-Agent Workflow
-# -------------------------------
+
+# Run the process
 question = "How many albums did Mercedes Sosa release between 2000 and 2009?"
 response = process_question(question)
 print("Final Response:", response)
 
 
+
 question_retriever_tool = create_retriever_tool(
     retriever=retriever,
     name="Question_Search",
@@ -561,12 +502,12 @@ def retriever(state: MessagesState):
     """Retriever node using similarity scores for filtering"""
     query = state["messages"][0].content
     results = vector_store.similarity_search_with_score(query, k=4)  # top 4 matches
-
+    
     # Dynamically adjust threshold based on query complexity
     threshold = 0.75 if "who" in query else 0.8
-
     filtered = [doc for doc, score in results if score < threshold]
-    
+
+    # Provide a default message if no documents found
     if not filtered:
         example_msg = HumanMessage(content="No relevant documents found.")
     else:
@@ -605,9 +546,12 @@ def get_llm(provider: str, config: dict):
 
 
 def generate_final_answer(state, tools_results):
-    # Combine results from all tools
-    # For example, if both the calculator and the Wikipedia tool were used:
-    final_answer = f"Answer: {tools_results['wiki_search']}\nAdditional Info: {tools_results['calculator']}"
+    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