Update agent.py

agent.py

@@ -331,6 +331,53 @@ 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)
+# -------------------------------
+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
 # -------------------------------
@@ -433,6 +480,75 @@ for doc, score in results:
 # -----------------------------
 retriever = vector_store.as_retriever()
 
+# -------------------------------
+# Step 6: Create LangChain Tools
+# -------------------------------
+wiki_tool = Tool(
+    name="Wiki_Search",
+    func=WikipediaAPIWrapper().run,  # Assuming WikipediaAPIWrapper is implemented
+    description="Search Wikipedia for related information."
+)
+
+calc_tool = Tool(
+    name="Calculator",
+    func=Calculator().run,
+    description="Perform mathematical calculations."
+)
+
+file_tool = Tool(
+    name="File_Analysis",
+    func=your_file_analysis_function,  # Replace with a file analysis function
+    description="Analyze and extract data from attachments."
+)
+
+# -------------------------------
+# Step 7: Create the Planner-Agent Logic
+# -------------------------------
+# Define the agent tool set
+tools = [wiki_tool, calc_tool, file_tool]
+
+# Create an agent using the planner, task classifier, and decision logic
+agent = initialize_agent(
+    tools=tools,
+    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # Or another agent type like CHAT_ZERO_SHOT
+    verbose=True
+)
+
+# -------------------------------
+# Step 8: Use the Planner, Classifier, and Decision Logic
+# -------------------------------
+def process_question(question):
+    # Step 1: Planner generates the task sequence
+    tasks = planner(question)
+    print(f"Tasks to perform: {tasks}")
+
+    # Step 2: Classify the task (based on question)
+    task_type = task_classifier(question)
+    print(f"Task type: {task_type}")
+
+    # Step 3: Use the classifier and planner to decide on the next task or node
+    state = {"question": question, "last_response": ""}
+    next_task = decide_task(state)
+    print(f"Next task: {next_task}")
+
+    # Step 4: Use node skipper logic (skip if needed)
+    skip = node_skipper(state)
+    if skip:
+        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 9: Run the Planner-Agent Workflow
+# -------------------------------
+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",
@@ -446,10 +562,11 @@ def retriever(state: MessagesState):
     query = state["messages"][0].content
     results = vector_store.similarity_search_with_score(query, k=4)  # top 4 matches
 
-    # Filter by score (lower is more similar; adjust threshold as needed)
-    threshold = 0.8
-    filtered = [doc for doc, score in results if score < threshold]
+    # 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]
+    
     if not filtered:
         example_msg = HumanMessage(content="No relevant documents found.")
     else:
@@ -487,6 +604,14 @@ def get_llm(provider: str, config: dict):
         raise ValueError(f"Invalid provider: {provider}")
 
 
+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']}"
+    return final_answer
+
+
+
 # Build graph function
 def build_graph():
     """Build the graph based on provider"""