feat: added RAG agent tools

src/rag_agent.py

@@ -1,6 +1,7 @@
 from llama_index.core import SimpleDirectoryReader, Settings, VectorStoreIndex
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.tools import FunctionTool, QueryEngineTool, RetrieverTool
 
 
 class RAGAgent:
@@ -9,7 +10,6 @@ class RAGAgent:
         self.streaming = streaming
         self.llm = llm
         self.index = None
-        self.query_engine = None
 
     def load_and_index_folder(self, folder_path, embedding_model="BAAI/bge-small-en-v1.5"):
         """
@@ -32,27 +32,37 @@ class RAGAgent:
         Settings.embed_model = HuggingFaceEmbedding(model_name=embedding_model)
         index = VectorStoreIndex.from_documents(nodes)
         
-        # create query engine
+        # create index
         Settings.llm = self.llm
         self.index = index
-        self.query_engine = index.as_query_engine(streaming=self.streaming)
-
-    def query_index(self, query_text):
+    
+    def get_tools(self):
         """
-        Query the indexed documents.
-
-        Args:
-            query_text (str): The query string.
+        Get tools for the RAG agent.
 
         Returns:
-            str: The RAG agent's response.
+            list: List of FunctionTool instances.
         """
         if not hasattr(self, 'index'):
             raise ValueError("Index not found. Please load and index a folder first.")
 
-        response = self.query_engine.query(query_text)
+        retriever = self.index.as_retriever()
+        query_engine = self.index.as_query_engine()
+
+        tools = [
+            RetrieverTool.from_defaults(
+                retriever,
+                name="RAG_Document_Retrieval",
+                description="Useful for retrieving relevant documents based on user queries."
+            ),
+            QueryEngineTool.from_defaults(
+                query_engine,
+                name="RAG_Document_Query",
+                description="Useful for answering questions by querying the indexed data directly."
+            )
+        ]
 
-        return str(response)
+        return tools