Add llama-index rag (#697)

Author: tastelikefeet

ms_agent/agent/llm_agent.py

@@ -10,7 +10,7 @@
 from ms_agent.callbacks import Callback, callbacks_mapping
 from ms_agent.llm.llm import LLM
 from ms_agent.llm.utils import Message, Tool
-from ms_agent.rag.base import Rag
+from ms_agent.rag.base import RAG
 from ms_agent.rag.utils import rag_mapping
 from ms_agent.tools import ToolManager
 from ms_agent.utils import async_retry
@@ -62,7 +62,7 @@ def __init__(self,
         self.tool_manager: Optional[ToolManager] = None
         self.memory_tools: List[Memory] = []
         self.planer: Optional[Planer] = None
-        self.rag: Optional[Rag] = None
+        self.rag: Optional[RAG] = None
         self.llm: Optional[LLM] = None
         self.runtime: Optional[Runtime] = None
         self.max_chat_round: int = 0
@@ -220,7 +220,7 @@ async def _prepare_messages(
             Message(role='user', content=inputs or query),
         ]
         if self.rag is not None:
-            messages = await self.rag.run(messages)
+            messages = await self.rag.query(messages[1].content)
         return messages
 
     async def _prepare_memory(self):
@@ -251,7 +251,7 @@ async def _prepare_rag(self):
                 assert rag.name in rag_mapping, (
                     f'{rag.name} not in rag_mapping, '
                     f'which supports: {list(rag_mapping.keys())}')
-                self.rag: Rag = rag_mapping(rag.name)(self.config)
+                self.rag: RAG = rag_mapping(rag.name)(self.config)
 
     async def _refine_memory(self, messages: List[Message]) -> List[Message]:
         """

ms_agent/llm/openai_llm.py

@@ -50,7 +50,7 @@ def __init__(
             base_url=base_url,
         )
         self.args: Dict = OmegaConf.to_container(
-            getattr(config, 'generation_config', {}))
+            getattr(config, 'generation_config', DictConfig({})))
 
     def format_tools(self,
                      tools: Optional[List[Tool]] = None

ms_agent/rag/base.py

@@ -1,61 +1,52 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from typing import Any, List
 
-from ms_agent.llm import Message
 
-
-class Rag:
+class RAG(ABC):
     """The base class for rags"""
 
     def __init__(self, config):
         self.config = config
 
     @abstractmethod
-    async def add_document(self, url: str, content: str, **metadata) -> bool:
+    async def add_documents(self, documents: List[str]) -> bool:
         """Add document to Rag
 
         Args:
-            url(`str`): The url of the document
-            content(`str`): The content of the document
-            **metadata: Metadata information
+            documents(`List[str]`): The content of the document
 
         Returns:
             success or not
         """
         pass
 
     @abstractmethod
-    async def search_documents(self,
-                               query: str,
-                               limit: int = 5,
-                               score_threshold: float = 0.7,
-                               **filters) -> List[Any]:
-        """Search documents in Rag
+    async def query(self, query: str) -> str:
+        """Search documents
 
         Args:
             query(`str`): The query to search for
-            limit(`int`): The number of documents to return
-            score_threshold(`float`): The score threshold
-            **filters: Any extra filters
-
         Returns:
-            List of documents
+            The query result
         """
         pass
 
     @abstractmethod
-    async def delete_document(self, url: str) -> bool:
-        """Delete document from Rag
+    async def retrieve(self,
+                       query: str,
+                       limit: int = 5,
+                       score_threshold: float = 0.7,
+                       **filters) -> List[Any]:
+        """Retrieve documents
 
         Args:
-            url(`str`): The url of the document
+            query(`str`): The query to search for
+            limit(`int`): The number of documents to return
+            score_threshold(`float`): The score threshold
+            **filters: Any extra filters
 
         Returns:
-            bool: True if the document was successfully deleted
+            List of documents
         """
         pass
-
-    @abstractmethod
-    async def run(self, inputs: List[Message]) -> List[Message]:
-        pass