Update agent.py

agent.py

@@ -10,8 +10,9 @@ from langchain_groq import ChatGroq
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
+from langchain.utilities import WikipediaAPIWrapper
 from langchain_community.document_loaders import ArxivLoader
-from langchain_community.vectorstores import SupabaseVectorStore
+#from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain.tools.retriever import create_retriever_tool
@@ -98,6 +99,29 @@ def modulus(a: int, b: int) -> int:
     """
     return a % b
 
+@tool
+def calc_tool(a: int, b: int, operation: str) -> float:
+    """
+    Perform a calculation between two numbers.
+
+    Args:
+        a: First number.
+        b: Second number.
+        operation: One of 'add', 'subtract', 'multiply', 'divide', 'modulus'.
+    """
+    operation = operation.lower()
+    if operation == "add":
+        return add(a, b)
+    elif operation == "subtract":
+        return subtract(a, b)
+    elif operation == "multiply":
+        return multiply(a, b)
+    elif operation == "divide":
+        return divide(a, b)
+    elif operation == "modulus":
+        return modulus(a, b)
+    else:
+        raise ValueError(f"Unsupported operation: {operation}")
 
 @tool
 def wiki_search(query: str) -> str:
@@ -159,10 +183,11 @@ def arvix_search(query: str) -> str:
 
 
 
+
 @tool
 def analyze_attachment(file_path: str) -> str:
     """
-    Analyzes attachments including PDF, TXT, DOCX, and XLSX files and returns text content.
+    Analyzes attachments including PY, PDF, TXT, DOCX, and XLSX files and returns text content.
 
     Args:
         file_path: Local path to the attachment.
@@ -170,32 +195,38 @@ def analyze_attachment(file_path: str) -> str:
     if not os.path.exists(file_path):
         return f"File not found: {file_path}"
 
-    if file_path.lower().endswith(".pdf"):
-        loader = PyMuPDFLoader(file_path)
-        documents = loader.load()
-        content = "\n\n".join([doc.page_content for doc in documents])
-
-    elif file_path.lower().endswith(".txt"):
-        loader = TextLoader(file_path)
-        documents = loader.load()
-        content = "\n\n".join([doc.page_content for doc in documents])
-
-    elif file_path.lower().endswith(".docx"):
-        doc = DocxDocument(file_path)
-        content = "\n".join([para.text for para in doc.paragraphs])
+    try:
+        ext = file_path.lower()
+
+        if ext.endswith(".pdf"):
+            loader = PyMuPDFLoader(file_path)
+            documents = loader.load()
+            content = "\n\n".join([doc.page_content for doc in documents])
+
+        elif ext.endswith(".txt") or ext.endswith(".py"):
+            # Both .txt and .py are plain text files
+            with open(file_path, "r", encoding="utf-8") as file:
+                content = file.read()
+
+        elif ext.endswith(".docx"):
+            doc = DocxDocument(file_path)
+            content = "\n".join([para.text for para in doc.paragraphs])
+
+        elif ext.endswith(".xlsx"):
+            wb = openpyxl.load_workbook(file_path, data_only=True)
+            content = ""
+            for sheet in wb:
+                content += f"Sheet: {sheet.title}\n"
+                for row in sheet.iter_rows(values_only=True):
+                    content += "\t".join([str(cell) if cell is not None else "" for cell in row]) + "\n"
 
-    elif file_path.lower().endswith(".xlsx"):
-        wb = openpyxl.load_workbook(file_path, data_only=True)
-        content = ""
-        for sheet in wb:
-            content += f"Sheet: {sheet.title}\n"
-            for row in sheet.iter_rows(values_only=True):
-                content += "\t".join([str(cell) if cell is not None else "" for cell in row]) + "\n"
+        else:
+            return "Unsupported file format. Please use PY, PDF, TXT, DOCX, or XLSX."
 
-    else:
-        return "Unsupported file format. Please use PDF, TXT, DOCX, or XLSX."
+        return content[:3000]  # Limit output size for readability
 
-    return content[:3000]  # Limit size for readability
+    except Exception as e:
+        return f"An error occurred while processing the file: {str(e)}"
 
 
 @tool
@@ -425,15 +456,57 @@ calc_tool = Tool(
 
 file_tool = Tool(
     name="File_Analysis",
-    func=your_file_analysis_function,  # Replace with a file analysis function
+    func=analyze_attachment().run,  # Replace with a file analysis function
     description="Analyze and extract data from attachments."
 )
 
+
+web_tool = Tool(
+    name="web_search",
+    func=web_search().run,
+    description="Perform web search."
+)
+
+
+arvix_tool = Tool(
+    name="arvix_search",
+    func=arvix_search().run,
+    description="Perform arvix search."
+)
+
+youtube_tool = Tool(
+    name="get_youtube_transcript",
+    func=get_youtube_transcript().run,
+    description="get youtube transcript."
+)
+
+
+video_tool = Tool(
+    name="extract_video_id",
+    func=extract_video_id().run,
+    description="extract_video_id."
+)
+
+
+analyze_tool = Tool(
+    name="analyze_attachment",
+    func=analyze_attachment().run,
+    description="analyze attachment."
+)
+
+
+wikiq_tool = Tool(
+    name="wikidata_query",
+    func=wikidata_query().run,
+    description="wikidata query."
+)
+
+
 # -------------------------------
 # Step 7: Create the Planner-Agent Logic
 # -------------------------------
 # Define the agent tool set
-tools = [wiki_tool, calc_tool, file_tool]
+tools = [wiki_tool, calc_tool, file_tool, web_tool, arvix_tool, youtube_tool, video_tool, analyze_tool, wikiq_tool]
 
 # Create an agent using the planner, task classifier, and decision logic
 agent = initialize_agent(