Create agent_tools.py

agent_tools.py

@@ -0,0 +1,112 @@
+from langchain_community.document_loaders import TextLoader, PyPDFLoader, Docx2txtLoader, CSVLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.retrievers import BM25Retriever
+
+@tool
+def addition_tool(a:float, b:float)->float:
+  """
+  Takes two floating point numbers and returns the results of their addition.
+  Args:
+    a: First number of the addition operation
+    b: Second number of the addition operation
+  """
+  return f"The result of addition between {a} and {b} is {a+b}"
+
+@tool
+def subtraction_tool(a:float, b:float)->float:
+  """
+  Takes two floating point numbers and returns the results of their subtraction.
+  Args:
+    a: First number of the subtraction operation
+    b: Second number of the subtraction operation
+  """
+  return f"The result of subtraction between {a} and {b} is {a-b}"
+
+@tool
+def multiplication_tool(a:float, b:float)->float:
+  """
+  Takes two floating point numbers and returns the results of their addition.
+  Args:
+    a: First number of the multiplication operation
+    b: Second number of the multiplication operation
+  """
+  return f"The result of multiplication between {a} and {b} is {a*b}"
+
+@tool
+def division_tool(a:float, b:float)->float:
+  """
+  Takes two floating point numbers and returns the results of their addition.
+  Args:
+    a: The dividend of the division operation
+    b: The divisor of the division operation
+  """
+  return f"The result of division between {a} and {b} is {a/b}"
+
+@tool
+def exponent_tool(a:float, b:int)->float:
+  """
+  Takes a floating point number and an integer expoent and returns the results of their exponential operation.
+  Args:
+    a: Base number of the addition operation
+    b: Exponent number of the addition operation
+  """
+  return f"The result of exponential operation between {a} and {b} is {a**b}"
+
+
+@tool
+def modulus_tool(a:int, b:int)->float:
+  """Get the modulus of two numbers.
+    Args:
+        a: first int
+        b: second int
+  """
+  return f"The result of modulus operation between {a} and {b} is {a%b}"
+
+
+class RagTool(Tool):
+    name = "document_retrieval_tool"
+    description = "Loads documents and retrieves important information from them."
+    inputs = {
+        "doc_path": {
+            "type": "string",
+            "description": "The local document path"
+        },
+        "query": {
+            "type": "string",
+            "description": "The information that needs to be retrieved from the document."
+        },
+    }
+    output_type = "string"
+
+    class GeneralLoader:
+        def __init__(self, doc_path: str):
+            self.doc_path = doc_path
+
+        def get_doc_format(self):
+            doc_format = self.doc_path.split(".")[-1]
+            return "docx" if doc_format == "doc" else doc_format
+
+        def load_doc(self, path, format):
+            loaders = {
+                "txt": TextLoader,
+                "pdf": PyPDFLoader,
+                "docx": Docx2txtLoader,
+                "csv": CSVLoader
+            }
+            doc_loader_cls = loaders.get(format)
+            if not doc_loader_cls:
+                raise ValueError(f"Unsupported document format: {format}")
+            loader = doc_loader_cls(path)
+            splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=30)
+            docs = loader.load_and_split(text_splitter=splitter)
+            return docs
+
+        def get_docs(self):
+            format = self.get_doc_format()
+            return self.load_doc(self.doc_path, format)
+
+    def forward(self, doc_path: str, query: str):
+        docs = self.GeneralLoader(doc_path).get_docs()
+        retriever = BM25Retriever.from_documents(docs)
+        results = retriever.get_relevant_documents(query)
+        return "\n\n".join([doc.page_content for doc in results[:3]]) if results else "No matching information found"