From 317d6530637d1e979341aa939092aede2d2a341d Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 13 Mar 2024 14:32:10 +0800 Subject: [PATCH 01/21] initial design for RAG module --- src/agentscope/rag/__init__.py | 11 +++++ src/agentscope/rag/langchain_rag.py | 67 +++++++++++++++++++++++++++ src/agentscope/rag/llama_index_rag.py | 52 +++++++++++++++++++++ src/agentscope/rag/rag.py | 66 ++++++++++++++++++++++++++ 4 files changed, 196 insertions(+) create mode 100644 src/agentscope/rag/__init__.py create mode 100644 src/agentscope/rag/langchain_rag.py create mode 100644 src/agentscope/rag/llama_index_rag.py create mode 100644 src/agentscope/rag/rag.py diff --git a/src/agentscope/rag/__init__.py b/src/agentscope/rag/__init__.py new file mode 100644 index 000000000..2a1d3de61 --- /dev/null +++ b/src/agentscope/rag/__init__.py @@ -0,0 +1,11 @@ +# -*- coding: utf-8 -*- +""" Import all pipeline related modules in the package. """ +from .rag import RAGBase +from .llama_index_rag import LlamaIndexRAG +from .langchain_rag import LangChainRAG + +__all__ = [ + "RAGBase", + "LlamaIndexRAG", + "LangChainRAG", +] diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py new file mode 100644 index 000000000..3fbea2a15 --- /dev/null +++ b/src/agentscope/rag/langchain_rag.py @@ -0,0 +1,67 @@ +# -*- coding: utf-8 -*- +""" +This module is integrate the LangChain RAG model into our AgentScope package +""" + + +from typing import Type, Any, Union +from pathlib import Path + +from langchain_community.document_loaders.base import BaseLoader +from langchain_core.vectorstores import VectorStore +from langchain_text_splitters.base import TextSplitter +from langchain_openai import OpenAIEmbeddings + +from agentscope.rag import RAGBase +from agentscope.models import ModelWrapperBase + + +class LangChainRAG(RAGBase): + """ + This class is a wrapper around the LangChain RAG. + """ + + def __init__( + self, + model: ModelWrapperBase, + loader_type: Type[BaseLoader], + splitter_type: Type[TextSplitter], + vector_store_type: Type[VectorStore], + embedding_model: Type[OpenAIEmbeddings], + **kwargs: Any, + ) -> None: + super().__init__(model, **kwargs) + self.loader_type = loader_type + self.splitter_type = splitter_type + self.config = kwargs + self.vector_store_type = vector_store_type + self.loader = None + self.splitter = splitter_type + self.vector_store = None + self.embedding_model = embedding_model + self.retriever = None + + def load_data(self, path: Union[Path, str], **kwargs: Any) -> Any: + """loading data from a directory""" + docs = self.loader.load() + all_splits = self.splitter.split_documents(docs) + return all_splits + + def store_and_index(self, docs: Any) -> None: + """indexing the documents and store them into the vector store""" + self.vector_store = self.vector_store_type.from_documents( + documents=docs, + embedding=self.embedding_model(), + ) + # build retriever + k = self.config.get("k", 6) + search_type = self.config.get("search_type", "similarity") + self.retriever = self.vector_store.as_retriever( + search_type=search_type, + search_kwargs={"k": k}, + ) + + def retrieve(self, query: Any) -> list[Any]: + """retrieve the documents based on the query""" + retrieved_docs = self.retriever.invoke(query) + return retrieved_docs diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py new file mode 100644 index 000000000..d9e6fbf05 --- /dev/null +++ b/src/agentscope/rag/llama_index_rag.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +""" +This module is integrate the Llama index RAG model into our AgentScope package +""" + +from typing import Any, Union, Type +from pathlib import Path +from llama_index.core import VectorStoreIndex +from llama_index.core.readers.base import BaseReader +from llama_index.core.base.base_retriever import BaseRetriever + +from agentscope.rag import RAGBase +from agentscope.models import ModelWrapperBase + + +class LlamaIndexRAG(RAGBase): + """ + This class is a wrapper around the Llama index RAG. + """ + + def __init__( + self, + model: ModelWrapperBase, + loader_type: Type[BaseReader], + vector_store_type: Type[VectorStoreIndex], + retriever_type: Type[BaseRetriever], + **kwargs: Any, + ) -> None: + super().__init__(model, **kwargs) + self.loader_type = loader_type + self.vector_store_type = vector_store_type + self.retriever_type = retriever_type + self.retriever = None + self.query_engine = None + self.persist_dir = kwargs.get("persist_dir", "./") + + def load_data( + self, + path: Union[Path, str], + **kwargs: Any, + ) -> Any: + documents = self.loader_type(path).load_data() + return documents + + def store_and_index(self, docs: Any) -> None: + index = self.vector_store_type.from_documents(docs) + index.storage_context.persist(persist_dir=self.persist_dir) + self.query_engine = index.as_query_engine() + + def retrieve(self, query: Any) -> list[Any]: + response = self.query_engine.query(query) + return response diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py new file mode 100644 index 000000000..713f53a6e --- /dev/null +++ b/src/agentscope/rag/rag.py @@ -0,0 +1,66 @@ +# -*- coding: utf-8 -*- +""" +Base class module for retrieval augmented generation (RAG). +To accommodate the RAG process of different packages, +we abstract the RAG process into four stages: +- data loading: loading data into memory for following processing; +- data indexing and storage: document chunking, embedding generation, +and off-load the data into VDB; +- data retrieval: taking a query and return a batch of documents or +document chunks; +- post-processing of the retrieved data: use the retrieved data to +generate an answer. +""" + +from abc import ABC, abstractmethod +from typing import Any, Optional, Union +from pathlib import Path + +from agentscope.models import ModelWrapperBase + + +class RAGBase(ABC): + """Base class for RAG""" + + def __init__( + self, + model: Optional[ModelWrapperBase], + **kwargs: Any, + ) -> None: + # pylint: disable=unused-argument + self.postprocessing_model = model + + @abstractmethod + def load_data(self, path: Union[Path, str], **kwargs: Any) -> Any: + """ + load data from disk to memory for following preprocessing + """ + + @abstractmethod + def store_and_index(self, docs: Any) -> None: + """ + preprocessing the loaded documents, for example: + 1) chunking, + 2) generate embedding, + 3) store the embedding-content to vdb + """ + + @abstractmethod + def retrieve(self, query: Any) -> list[Any]: + """ + retrieve list of content from vdb to memory + """ + + def generate_answer( + self, + retrieved_docs: list[str], + prompt: str, + **kwargs: Any, + ) -> Any: + """ + post-processing function, generates answer based on the + retrieved documents. + Example: + self.postprocessing_model(prompt.format(retrieved_docs)) + """ + raise NotImplementedError From 3330cf5f6a5c9d78441b47fc6cb6560813f7fbf8 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 13 Mar 2024 14:40:46 +0800 Subject: [PATCH 02/21] unify types --- src/agentscope/rag/langchain_rag.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index 3fbea2a15..5651d97cb 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -36,14 +36,16 @@ def __init__( self.config = kwargs self.vector_store_type = vector_store_type self.loader = None - self.splitter = splitter_type + self.splitter = None self.vector_store = None self.embedding_model = embedding_model self.retriever = None def load_data(self, path: Union[Path, str], **kwargs: Any) -> Any: """loading data from a directory""" + self.loader = self.loader_type(path) docs = self.loader.load() + self.splitter = self.splitter_type(**kwargs) all_splits = self.splitter.split_documents(docs) return all_splits From 3df40f233f327338c0876cfd0c891da09114989c Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Thu, 14 Mar 2024 20:19:19 +0800 Subject: [PATCH 03/21] added preliminary examples --- examples/rag/__init__.py | 0 examples/rag/rag_agents.py | 107 +++++++++++++++++++++++ examples/rag/rag_example.py | 41 +++++++++ src/agentscope/models/dashscope_model.py | 2 +- src/agentscope/rag/llama_index_rag.py | 33 ++++--- 5 files changed, 169 insertions(+), 14 deletions(-) create mode 100644 examples/rag/__init__.py create mode 100644 examples/rag/rag_agents.py create mode 100644 examples/rag/rag_example.py diff --git a/examples/rag/__init__.py b/examples/rag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py new file mode 100644 index 000000000..cd34c0447 --- /dev/null +++ b/examples/rag/rag_agents.py @@ -0,0 +1,107 @@ +# -*- coding: utf-8 -*- +""" +This example shows how to build a agent with RAG (backup by LlamaIndex) +""" + +from typing import Optional +from llama_index.embeddings.dashscope import DashScopeEmbedding +from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader + +from agentscope.prompt import PromptType +from agentscope.agents.agent import AgentBase +from agentscope.prompt import PromptEngine +from agentscope.message import Msg +from agentscope.rag.llama_index_rag import LlamaIndexRAG + + +Settings.embed_model = DashScopeEmbedding() + + +class LlamaIndexAgent(AgentBase): + """ + A LlamaIndex agent build on LlamaIndex. + """ + + def __init__( + self, + name: str, + sys_prompt: Optional[str] = None, + model_config_name: str = None, + use_memory: bool = True, + memory_config: Optional[dict] = None, + prompt_type: Optional[PromptType] = PromptType.LIST, + ) -> None: + super().__init__( + name=name, + sys_prompt=sys_prompt, + model_config_name=model_config_name, + use_memory=use_memory, + memory_config=memory_config, + ) + # init prompt engine + self.engine = PromptEngine(self.model, prompt_type=prompt_type) + self.rag = LlamaIndexRAG( + model=self.model, + loader_type=SimpleDirectoryReader, + vector_store_type=VectorStoreIndex, + ) + docs = self.rag.load_data(path="./data") + self.rag.store_and_index(docs) + + def reply(self, x: dict = None) -> dict: + """ + Reply function of the LlamaIndex agent. + Processes the input data, + 1) use the input data to retrieve with RAG function; + 2) generates a prompt using the current memory and system + prompt; + 3) invokes the language model to produce a response. The + response is then formatted and added to the dialogue memory. + + Args: + x (`dict`, defaults to `None`): + A dictionary representing the user's input to the agent. This + input is added to the memory if provided. Defaults to + None. + Returns: + A dictionary representing the message generated by the agent in + response to the user's input. + """ + # record the input if needed + if x is not None: + self.memory.add(x) + + content = x.content + retrieved_docs = self.rag.retrieve(content) + + retrieved_docs_to_string = "" + for node in retrieved_docs: + print(node) + retrieved_docs_to_string += node.get_text() + + print(retrieved_docs_to_string) + + # prepare prompt + prompt = self.engine.join( + { + "role": "system", + "content": self.sys_prompt.format_map( + {"retrieved_context": retrieved_docs_to_string}, + ), + }, + # {"role": "system", "content": retrieved_docs_to_string}, + self.memory.get_memory(), + ) + + print(prompt) + # call llm and generate response + response = self.model(prompt).text + msg = Msg(self.name, response) + + # Print/speak the message in this agent's voice + self.speak(msg) + + # Record the message in memory + self.memory.add(msg) + + return msg diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py new file mode 100644 index 000000000..e9414f726 --- /dev/null +++ b/examples/rag/rag_example.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +""" +A simple example for conversation between user and +a agent with RAG capability. +""" +import os +from rag_agents import LlamaIndexAgent +import agentscope +from agentscope.agents import UserAgent + + +def main() -> None: + """A conversation demo""" + + agentscope.init( + model_configs=[ + { + "model_type": "dashscope_chat", + "config_name": "qwen_config", + "model_name": "qwen-max", + "api_key": f"{os.environ.get('DASHSCOPE_API_KEY')}", + }, + ], + ) + + # Init two agents + rag_agent = LlamaIndexAgent( + name="Assistant", + sys_prompt="You're a helpful assistant. You need to generate answers " + "based on the provided context:\n " + "Context: \n {retrieved_context}\n ", + model_config_name="qwen_config", # replace by your model config name + ) + user_agent = UserAgent() + # start the conversation between user and assistant + x = user_agent() + rag_agent(x) + + +if __name__ == "__main__": + main() diff --git a/src/agentscope/models/dashscope_model.py b/src/agentscope/models/dashscope_model.py index b246909c5..da0be6077 100644 --- a/src/agentscope/models/dashscope_model.py +++ b/src/agentscope/models/dashscope_model.py @@ -182,7 +182,7 @@ def __call__( "Each message in the 'messages' list must contain a 'role' " "and 'content' key for DashScope API.", ) - + print(messages) # step3: forward to generate response response = dashscope.Generation.call( model=self.model, diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index d9e6fbf05..15bd3d815 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -3,11 +3,11 @@ This module is integrate the Llama index RAG model into our AgentScope package """ -from typing import Any, Union, Type +from typing import Any, Union, Type, Optional from pathlib import Path -from llama_index.core import VectorStoreIndex from llama_index.core.readers.base import BaseReader from llama_index.core.base.base_retriever import BaseRetriever +from llama_index.core import VectorStoreIndex, SimpleDirectoryReader from agentscope.rag import RAGBase from agentscope.models import ModelWrapperBase @@ -21,17 +21,17 @@ class LlamaIndexRAG(RAGBase): def __init__( self, model: ModelWrapperBase, - loader_type: Type[BaseReader], - vector_store_type: Type[VectorStoreIndex], - retriever_type: Type[BaseRetriever], + loader_type: Optional[Type[BaseReader]] = None, + vector_store_type: Optional[Type[VectorStoreIndex]] = None, + retriever_type: Optional[Type[BaseRetriever]] = None, **kwargs: Any, ) -> None: super().__init__(model, **kwargs) - self.loader_type = loader_type - self.vector_store_type = vector_store_type + self.loader_type = loader_type or SimpleDirectoryReader + self.vector_store_type = vector_store_type or VectorStoreIndex self.retriever_type = retriever_type self.retriever = None - self.query_engine = None + self.index = None self.persist_dir = kwargs.get("persist_dir", "./") def load_data( @@ -43,10 +43,17 @@ def load_data( return documents def store_and_index(self, docs: Any) -> None: - index = self.vector_store_type.from_documents(docs) - index.storage_context.persist(persist_dir=self.persist_dir) - self.query_engine = index.as_query_engine() + """ + In LlamaIndex terms, an Index is a data structure composed + of Document objects, designed to enable querying by an LLM. + A vector store index takes Documents and splits them up into + Nodes (chunks). It then creates vector embeddings of the + text of every node, ready to be queried by an LLM. + """ + self.index = self.vector_store_type.from_documents(docs) + self.index.storage_context.persist(persist_dir=self.persist_dir) + self.retriever = self.index.as_retriever() def retrieve(self, query: Any) -> list[Any]: - response = self.query_engine.query(query) - return response + retrieved_docs = self.retriever.retrieve(str(query)) + return retrieved_docs From d9d42d569a0f8b95d90e7a3aea245e6f917e1a75 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Mon, 18 Mar 2024 13:58:24 +0800 Subject: [PATCH 04/21] update to allow embedding model wrapper --- examples/rag/rag_agents.py | 48 ++++--- examples/rag/rag_example.py | 8 ++ src/agentscope/rag/langchain_rag.py | 26 +++- src/agentscope/rag/llama_index_rag.py | 198 ++++++++++++++++++++++---- src/agentscope/rag/rag.py | 25 +++- 5 files changed, 250 insertions(+), 55 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index cd34c0447..64a51839a 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -1,16 +1,18 @@ # -*- coding: utf-8 -*- """ -This example shows how to build a agent with RAG (backup by LlamaIndex) +This example shows how to build an agent with RAG (backup by LlamaIndex) """ from typing import Optional from llama_index.embeddings.dashscope import DashScopeEmbedding -from llama_index.core import Settings, VectorStoreIndex, SimpleDirectoryReader +from llama_index.core import Settings, SimpleDirectoryReader +from llama_index.core.node_parser import TokenTextSplitter from agentscope.prompt import PromptType from agentscope.agents.agent import AgentBase from agentscope.prompt import PromptEngine from agentscope.message import Msg +from agentscope.models import load_model_by_config_name from agentscope.rag.llama_index_rag import LlamaIndexRAG @@ -27,6 +29,7 @@ def __init__( name: str, sys_prompt: Optional[str] = None, model_config_name: str = None, + emb_model_config_name: str = None, use_memory: bool = True, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, @@ -40,15 +43,27 @@ def __init__( ) # init prompt engine self.engine = PromptEngine(self.model, prompt_type=prompt_type) + self.emb_model = load_model_by_config_name(emb_model_config_name) + # init rag related attributes self.rag = LlamaIndexRAG( model=self.model, - loader_type=SimpleDirectoryReader, - vector_store_type=VectorStoreIndex, + emb_model=self.emb_model, + ) + docs = self.rag.load_data( + loader=SimpleDirectoryReader("./data"), + ) + self.rag.store_and_index( + docs, + transformations=[ + TokenTextSplitter(), + ], ) - docs = self.rag.load_data(path="./data") - self.rag.store_and_index(docs) - def reply(self, x: dict = None) -> dict: + def reply( + self, + x: dict = None, + min_score: float = 0, + ) -> dict: """ Reply function of the LlamaIndex agent. Processes the input data, @@ -63,23 +78,23 @@ def reply(self, x: dict = None) -> dict: A dictionary representing the user's input to the agent. This input is added to the memory if provided. Defaults to None. + min_score (`float`, defaults to `0`): + A minimum score of the retrieved chunk to be used in the prompt Returns: A dictionary representing the message generated by the agent in response to the user's input. """ + retrieved_docs_to_string = "" # record the input if needed if x is not None: self.memory.add(x) + # retrieve when the input is not None + content = x.get("content", "") + retrieved_docs = self.rag.retrieve(content) - content = x.content - retrieved_docs = self.rag.retrieve(content) - - retrieved_docs_to_string = "" - for node in retrieved_docs: - print(node) - retrieved_docs_to_string += node.get_text() - - print(retrieved_docs_to_string) + for node in retrieved_docs: + if node.get_score() > min_score: + retrieved_docs_to_string += node.get_text() # prepare prompt prompt = self.engine.join( @@ -93,7 +108,6 @@ def reply(self, x: dict = None) -> dict: self.memory.get_memory(), ) - print(prompt) # call llm and generate response response = self.model(prompt).text msg = Msg(self.name, response) diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index e9414f726..94d09ce36 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -20,6 +20,12 @@ def main() -> None: "model_name": "qwen-max", "api_key": f"{os.environ.get('DASHSCOPE_API_KEY')}", }, + { + "model_type": "dashscope_text_embedding", + "config_name": "qwen_emb_config", + "model_name": "text-embedding-v2", + "api_key": f"{os.environ.get('DASHSCOPE_API_KEY')}", + }, ], ) @@ -30,10 +36,12 @@ def main() -> None: "based on the provided context:\n " "Context: \n {retrieved_context}\n ", model_config_name="qwen_config", # replace by your model config name + emb_model_config_name="qwen_emb_config", ) user_agent = UserAgent() # start the conversation between user and assistant x = user_agent() + x.role = "user" # to enforce dashscope requirement on roles rag_agent(x) diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index 5651d97cb..6dcd89eee 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -4,8 +4,9 @@ """ -from typing import Type, Any, Union -from pathlib import Path +from typing import Type, Any, Optional + +# from pathlib import Path from langchain_community.document_loaders.base import BaseLoader from langchain_core.vectorstores import VectorStore @@ -23,14 +24,15 @@ class LangChainRAG(RAGBase): def __init__( self, - model: ModelWrapperBase, + model: Optional[ModelWrapperBase], + emb_model: Optional[ModelWrapperBase], loader_type: Type[BaseLoader], splitter_type: Type[TextSplitter], vector_store_type: Type[VectorStore], embedding_model: Type[OpenAIEmbeddings], **kwargs: Any, ) -> None: - super().__init__(model, **kwargs) + super().__init__(model, emb_model, **kwargs) self.loader_type = loader_type self.splitter_type = splitter_type self.config = kwargs @@ -41,15 +43,25 @@ def __init__( self.embedding_model = embedding_model self.retriever = None - def load_data(self, path: Union[Path, str], **kwargs: Any) -> Any: + def load_data( + self, + loader: Any, + query: Any, + **kwargs: Any, + ) -> Any: """loading data from a directory""" - self.loader = self.loader_type(path) + self.loader = loader docs = self.loader.load() self.splitter = self.splitter_type(**kwargs) all_splits = self.splitter.split_documents(docs) return all_splits - def store_and_index(self, docs: Any) -> None: + def store_and_index( + self, + docs: Any, + vector_store: Any, + **kwargs: Any, + ) -> None: """indexing the documents and store them into the vector store""" self.vector_store = self.vector_store_type.from_documents( documents=docs, diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index 15bd3d815..830fe4bac 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -1,18 +1,80 @@ # -*- coding: utf-8 -*- """ -This module is integrate the Llama index RAG model into our AgentScope package +This module is an integration of the Llama index RAG +into AgentScope package """ -from typing import Any, Union, Type, Optional -from pathlib import Path +from typing import Any, Optional, List, Union + from llama_index.core.readers.base import BaseReader from llama_index.core.base.base_retriever import BaseRetriever -from llama_index.core import VectorStoreIndex, SimpleDirectoryReader +from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding +from llama_index.core.node_parser.interface import NodeParser +from llama_index.core.ingestion import IngestionPipeline +from llama_index.core.vector_stores.types import ( + BasePydanticVectorStore, + VectorStore, +) +from llama_index.core.bridge.pydantic import PrivateAttr + + +from llama_index.core import ( + VectorStoreIndex, + SimpleDirectoryReader, +) from agentscope.rag import RAGBase from agentscope.models import ModelWrapperBase +class _EmbeddingModel(BaseEmbedding): + """ + wrapp a ModelWrapperBase to an embedding mode in Llama Index. + """ + + _emb_model_wrapper: ModelWrapperBase = PrivateAttr() + + def __init__( + self, + emb_model: ModelWrapperBase, + embed_batch_size: int = 1, + ) -> None: + super().__init__( + model_name="Temporary_embedding_wrapper", + embed_batch_size=embed_batch_size, + ) + self._emb_model_wrapper = emb_model + + def _get_query_embedding(self, query: str) -> List[float]: + # Note: AgentScope embedding model wrapper returns list of embedding + return list(self._emb_model_wrapper(query).embedding[0]) + + def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]: + results = [ + list(self._emb_model_wrapper(t).embedding[0]) for t in texts + ] + return results + + def _get_text_embedding(self, text: str) -> Embedding: + return list(self._emb_model_wrapper(text).embedding[0]) + + # TODO: use proper async methods, but depends on model wrapper + async def _aget_query_embedding(self, query: str) -> List[float]: + """The asynchronous version of _get_query_embedding.""" + return self._get_query_embedding(query) + + async def _aget_text_embedding(self, text: str) -> List[float]: + """Asynchronously get text embedding.""" + return self._get_text_embedding(text) + + async def _aget_text_embeddings( + self, + texts: List[str], + ) -> List[List[float]]: + """Asynchronously get text embeddings.""" + return self._get_text_embeddings(texts) + + class LlamaIndexRAG(RAGBase): """ This class is a wrapper around the Llama index RAG. @@ -20,40 +82,128 @@ class LlamaIndexRAG(RAGBase): def __init__( self, - model: ModelWrapperBase, - loader_type: Optional[Type[BaseReader]] = None, - vector_store_type: Optional[Type[VectorStoreIndex]] = None, - retriever_type: Optional[Type[BaseRetriever]] = None, + model: Optional[ModelWrapperBase], + emb_model: Union[ModelWrapperBase, BaseEmbedding], **kwargs: Any, ) -> None: - super().__init__(model, **kwargs) - self.loader_type = loader_type or SimpleDirectoryReader - self.vector_store_type = vector_store_type or VectorStoreIndex - self.retriever_type = retriever_type + super().__init__(model, emb_model, **kwargs) self.retriever = None self.index = None self.persist_dir = kwargs.get("persist_dir", "./") def load_data( self, - path: Union[Path, str], + loader: BaseReader = SimpleDirectoryReader, + query: Optional[str] = None, **kwargs: Any, ) -> Any: - documents = self.loader_type(path).load_data() + """ + Accept a loader, loading the desired data (no chunking) + :param loader: object to load data, expected be an instance of class + inheriting from BaseReader. + :param query: optional, used when the data is in a database. + :return: the loaded documents (un-chunked) + + Example 1: use simple directory loader to load general documents, + including Markdown, PDFs, Word documents, PowerPoint decks, images, + audio and video. + ``` + load_data_to_chunks( + loader=SimpleDirectoryReader("./data") + ) + ``` + + Example 2: use SQL loader + ``` + load_data_to_chunks( + DatabaseReader( + scheme=os.getenv("DB_SCHEME"), + host=os.getenv("DB_HOST"), + port=os.getenv("DB_PORT"), + user=os.getenv("DB_USER"), + password=os.getenv("DB_PASS"), + dbname=os.getenv("DB_NAME"), + ), + query = "SELECT * FROM users" + ) + ``` + """ + if query is None: + documents = loader.load_data() + else: + documents = loader.load_data(query) return documents - def store_and_index(self, docs: Any) -> None: + def store_and_index( + self, + docs: Any, + vector_store: Union[BasePydanticVectorStore, VectorStore, None] = None, + retriever: Optional[BaseRetriever] = None, + **kwargs: Any, + ) -> Any: """ + Preprocessing the loaded documents. + :param docs: documents to be processed + :param vector_store: vector store + :param retriever: optional, specifies the retriever to use + :param args: additional + :param kwargs: + In LlamaIndex terms, an Index is a data structure composed of Document objects, designed to enable querying by an LLM. - A vector store index takes Documents and splits them up into - Nodes (chunks). It then creates vector embeddings of the - text of every node, ready to be queried by an LLM. + For example: + 1) preprocessing documents with + 2) generate embedding, + 3) store the embedding-content to vdb """ - self.index = self.vector_store_type.from_documents(docs) - self.index.storage_context.persist(persist_dir=self.persist_dir) - self.retriever = self.index.as_retriever() + # build and run preprocessing pipeline + transformations = [] + if "transformations" in kwargs: + for item in kwargs["transformations"]: + if isinstance(item, NodeParser): + transformations.append(item) - def retrieve(self, query: Any) -> list[Any]: - retrieved_docs = self.retriever.retrieve(str(query)) - return retrieved_docs + # adding embedding model as the last step of transformation + # https://docs.llamaindex.ai/en/stable/module_guides/loading/ingestion_pipeline/root.html + if isinstance(self.emb_model, ModelWrapperBase): + transformations.append(_EmbeddingModel(self.emb_model)) + elif isinstance(self.emb_model, BaseEmbedding): + transformations.append(self.emb_model) + + if vector_store is not None: + pipeline = IngestionPipeline( + transformations=transformations, + vector_store=vector_store, + ) + _ = pipeline.run(docs) + self.index = VectorStoreIndex.from_vector_store(vector_store) + else: + # No vector store is provide, use simple in memory + pipeline = IngestionPipeline( + transformations=transformations, + ) + nodes = pipeline.run(documents=docs) + self.index = VectorStoreIndex(nodes=nodes) + + if retriever is None: + self.retriever = self.index.as_retriever(**kwargs) + else: + self.retriever = retriever + return self.index + + def set_retriever(self, retriever: BaseRetriever) -> None: + """ + Reset the retriever if necessary. + """ + self.retriever = retriever + + def retrieve(self, query: str) -> list[Any]: + """ + This is a basic retrieve function + :param query: query is expected to be a question in string + + More advanced query processing can refer to + https://docs.llamaindex.ai/en/stable/examples/query_transformations/query_transform_cookbook.html + """ + retrieved = self.retriever.retrieve(str(query)) + return retrieved diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index 713f53a6e..e425fd91d 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -13,8 +13,7 @@ """ from abc import ABC, abstractmethod -from typing import Any, Optional, Union -from pathlib import Path +from typing import Any, Optional from agentscope.models import ModelWrapperBase @@ -25,19 +24,31 @@ class RAGBase(ABC): def __init__( self, model: Optional[ModelWrapperBase], + emb_model: Optional[ModelWrapperBase], **kwargs: Any, ) -> None: # pylint: disable=unused-argument self.postprocessing_model = model + self.emb_model = emb_model @abstractmethod - def load_data(self, path: Union[Path, str], **kwargs: Any) -> Any: + def load_data( + self, + loader: Any, + query: Any, + **kwargs: Any, + ) -> Any: """ - load data from disk to memory for following preprocessing + load data (documents) from disk to memory and chunking them """ @abstractmethod - def store_and_index(self, docs: Any) -> None: + def store_and_index( + self, + docs: Any, + vector_store: Any, + **kwargs: Any, + ) -> Any: """ preprocessing the loaded documents, for example: 1) chunking, @@ -51,9 +62,9 @@ def retrieve(self, query: Any) -> list[Any]: retrieve list of content from vdb to memory """ - def generate_answer( + def post_processing( self, - retrieved_docs: list[str], + retrieved_docs: list[Any], prompt: str, **kwargs: Any, ) -> Any: From 518321f8bc6e83a135eeab5bba4a52a157914472 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Mon, 18 Mar 2024 14:15:47 +0800 Subject: [PATCH 05/21] fix typos --- examples/rag/rag_agents.py | 2 ++ src/agentscope/rag/langchain_rag.py | 1 + src/agentscope/rag/llama_index_rag.py | 4 +++- src/agentscope/rag/rag.py | 15 +++++++++++---- 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 64a51839a..40ba3ea15 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -49,6 +49,7 @@ def __init__( model=self.model, emb_model=self.emb_model, ) + # load the document to memory docs = self.rag.load_data( loader=SimpleDirectoryReader("./data"), ) @@ -56,6 +57,7 @@ def __init__( docs, transformations=[ TokenTextSplitter(), + # just an example, more operators can be added ], ) diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index 6dcd89eee..e033ca54a 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -20,6 +20,7 @@ class LangChainRAG(RAGBase): """ This class is a wrapper around the LangChain RAG. + TODO: still under construction """ def __init__( diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index 830fe4bac..ca42313b2 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -29,7 +29,8 @@ class _EmbeddingModel(BaseEmbedding): """ - wrapp a ModelWrapperBase to an embedding mode in Llama Index. + wrapper for ModelWrapperBase to an embedding model can be used + in Llama Index pipeline. """ _emb_model_wrapper: ModelWrapperBase = PrivateAttr() @@ -185,6 +186,7 @@ def store_and_index( nodes = pipeline.run(documents=docs) self.index = VectorStoreIndex(nodes=nodes) + # set the retriever if retriever is None: self.retriever = self.index.as_retriever(**kwargs) else: diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index e425fd91d..9c82b204a 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -64,14 +64,21 @@ def retrieve(self, query: Any) -> list[Any]: def post_processing( self, - retrieved_docs: list[Any], + retrieved_docs: list[str], prompt: str, **kwargs: Any, ) -> Any: """ - post-processing function, generates answer based on the - retrieved documents. + A default solution for post-processing function, generates answer + based on the retrieved documents. + :param retrieved_docs: list of retrieved documents + :param prompt: prompt for LLM generating answer with the + retrieved documents + + Example: self.postprocessing_model(prompt.format(retrieved_docs)) """ - raise NotImplementedError + assert self.postprocessing_model + prompt = prompt.format(retrieved_docs) + return self.postprocessing_model(prompt, **kwargs).text From b7b60f43c82f4fb79b759670f68c679e13a699c4 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 11:23:54 +0800 Subject: [PATCH 06/21] update example --- examples/rag/rag_agents.py | 23 ++++-------- examples/rag/rag_example.py | 11 ++++-- src/agentscope/rag/llama_index_rag.py | 54 +++++++++++++++++++-------- src/agentscope/rag/rag.py | 7 +++- 4 files changed, 59 insertions(+), 36 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 40ba3ea15..55286ba7d 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -4,9 +4,7 @@ """ from typing import Optional -from llama_index.embeddings.dashscope import DashScopeEmbedding -from llama_index.core import Settings, SimpleDirectoryReader -from llama_index.core.node_parser import TokenTextSplitter +from llama_index.core import SimpleDirectoryReader from agentscope.prompt import PromptType from agentscope.agents.agent import AgentBase @@ -16,9 +14,6 @@ from agentscope.rag.llama_index_rag import LlamaIndexRAG -Settings.embed_model = DashScopeEmbedding() - - class LlamaIndexAgent(AgentBase): """ A LlamaIndex agent build on LlamaIndex. @@ -50,16 +45,14 @@ def __init__( emb_model=self.emb_model, ) # load the document to memory + # Feed the AgentScope tutorial documents, so that + # the agent can answer questions related to AgentScope! docs = self.rag.load_data( - loader=SimpleDirectoryReader("./data"), - ) - self.rag.store_and_index( - docs, - transformations=[ - TokenTextSplitter(), - # just an example, more operators can be added - ], + loader=SimpleDirectoryReader( + "../../docs/sphinx_doc/en/source/tutorial", + ), ) + self.rag.store_and_index(docs) def reply( self, @@ -93,11 +86,9 @@ def reply( # retrieve when the input is not None content = x.get("content", "") retrieved_docs = self.rag.retrieve(content) - for node in retrieved_docs: if node.get_score() > min_score: retrieved_docs_to_string += node.get_text() - # prepare prompt prompt = self.engine.join( { diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index 94d09ce36..8f839b66b 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- """ A simple example for conversation between user and -a agent with RAG capability. +an agent with RAG capability. """ import os from rag_agents import LlamaIndexAgent @@ -40,9 +40,12 @@ def main() -> None: ) user_agent = UserAgent() # start the conversation between user and assistant - x = user_agent() - x.role = "user" # to enforce dashscope requirement on roles - rag_agent(x) + while True: + x = user_agent() + x.role = "user" # to enforce dashscope requirement on roles + if len(x["content"]) == 0 or str(x["content"]).startswith("exit"): + break + rag_agent(x) if __name__ == "__main__": diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index ca42313b2..396114b92 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -9,13 +9,14 @@ from llama_index.core.readers.base import BaseReader from llama_index.core.base.base_retriever import BaseRetriever from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding -from llama_index.core.node_parser.interface import NodeParser from llama_index.core.ingestion import IngestionPipeline from llama_index.core.vector_stores.types import ( BasePydanticVectorStore, VectorStore, ) from llama_index.core.bridge.pydantic import PrivateAttr +from llama_index.core.node_parser.interface import NodeParser +from llama_index.core.node_parser import SentenceSplitter from llama_index.core import ( @@ -24,6 +25,7 @@ ) from agentscope.rag import RAGBase +from agentscope.rag.rag import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP from agentscope.models import ModelWrapperBase @@ -84,17 +86,29 @@ class LlamaIndexRAG(RAGBase): def __init__( self, model: Optional[ModelWrapperBase], - emb_model: Union[ModelWrapperBase, BaseEmbedding], + emb_model: Any = None, + config: Optional[dict] = None, **kwargs: Any, ) -> None: - super().__init__(model, emb_model, **kwargs) + super().__init__(model, emb_model, config, **kwargs) self.retriever = None self.index = None self.persist_dir = kwargs.get("persist_dir", "./") + self.emb_model = emb_model + + # ensure the emb_model is compatible with LlamaIndex + if issubclass(type(self.emb_model), ModelWrapperBase): + self.emb_model = _EmbeddingModel(self.emb_model) + elif isinstance(self.emb_model, BaseEmbedding): + pass + else: + raise TypeError( + f"Embedding model does not support {type(self.emb_model)}.", + ) def load_data( self, - loader: BaseReader = SimpleDirectoryReader, + loader: BaseReader = SimpleDirectoryReader("./data/"), query: Optional[str] = None, **kwargs: Any, ) -> Any: @@ -140,6 +154,7 @@ def store_and_index( docs: Any, vector_store: Union[BasePydanticVectorStore, VectorStore, None] = None, retriever: Optional[BaseRetriever] = None, + transformations: Optional[list[NodeParser]] = None, **kwargs: Any, ) -> Any: """ @@ -147,7 +162,8 @@ def store_and_index( :param docs: documents to be processed :param vector_store: vector store :param retriever: optional, specifies the retriever to use - :param args: additional + :param transformations: optional, specifies the transformations + to preprocess the documents :param kwargs: In LlamaIndex terms, an Index is a data structure composed @@ -158,18 +174,23 @@ def store_and_index( 3) store the embedding-content to vdb """ # build and run preprocessing pipeline - transformations = [] - if "transformations" in kwargs: - for item in kwargs["transformations"]: - if isinstance(item, NodeParser): - transformations.append(item) + if transformations is None: + transformations = [ + SentenceSplitter( + chunk_size=self.config.get( + "chunk_size", + DEFAULT_CHUNK_SIZE, + ), + chunk_overlap=self.config.get( + "chunk_overlap", + DEFAULT_CHUNK_OVERLAP, + ), + ), + ] # adding embedding model as the last step of transformation # https://docs.llamaindex.ai/en/stable/module_guides/loading/ingestion_pipeline/root.html - if isinstance(self.emb_model, ModelWrapperBase): - transformations.append(_EmbeddingModel(self.emb_model)) - elif isinstance(self.emb_model, BaseEmbedding): - transformations.append(self.emb_model) + transformations.append(self.emb_model) if vector_store is not None: pipeline = IngestionPipeline( @@ -188,7 +209,10 @@ def store_and_index( # set the retriever if retriever is None: - self.retriever = self.index.as_retriever(**kwargs) + self.retriever = self.index.as_retriever( + embed_model=self.emb_model, + **kwargs, + ) else: self.retriever = retriever return self.index diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index 9c82b204a..bc0e76b06 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -17,6 +17,9 @@ from agentscope.models import ModelWrapperBase +DEFAULT_CHUNK_SIZE = 1024 +DEFAULT_CHUNK_OVERLAP = 20 + class RAGBase(ABC): """Base class for RAG""" @@ -24,12 +27,14 @@ class RAGBase(ABC): def __init__( self, model: Optional[ModelWrapperBase], - emb_model: Optional[ModelWrapperBase], + emb_model: Any = None, + config: Optional[dict] = None, **kwargs: Any, ) -> None: # pylint: disable=unused-argument self.postprocessing_model = model self.emb_model = emb_model + self.config = config or {} @abstractmethod def load_data( From a88da63aa893548c459a125a146db453feab70f8 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 16:29:38 +0800 Subject: [PATCH 07/21] improve langchain rag --- examples/rag/rag_agents.py | 117 +++++++++++++++++----- examples/rag/rag_example.py | 45 +++++++-- src/agentscope/rag/langchain_rag.py | 133 ++++++++++++++++++++------ src/agentscope/rag/llama_index_rag.py | 11 ++- src/agentscope/rag/rag.py | 2 +- 5 files changed, 244 insertions(+), 64 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 55286ba7d..9d8efdf8a 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -5,6 +5,7 @@ from typing import Optional from llama_index.core import SimpleDirectoryReader +from langchain_community.document_loaders import DirectoryLoader from agentscope.prompt import PromptType from agentscope.agents.agent import AgentBase @@ -12,11 +13,13 @@ from agentscope.message import Msg from agentscope.models import load_model_by_config_name from agentscope.rag.llama_index_rag import LlamaIndexRAG +from agentscope.rag.langchain_rag import LangChainRAG -class LlamaIndexAgent(AgentBase): +class RAGAgent(AgentBase): """ - A LlamaIndex agent build on LlamaIndex. + Base class for RAG agents, child classes include the + RAG agents built with LlamaIndex and LangChain in this file """ def __init__( @@ -28,6 +31,7 @@ def __init__( use_memory: bool = True, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, + config: Optional[dict] = None, ) -> None: super().__init__( name=name, @@ -39,25 +43,15 @@ def __init__( # init prompt engine self.engine = PromptEngine(self.model, prompt_type=prompt_type) self.emb_model = load_model_by_config_name(emb_model_config_name) - # init rag related attributes - self.rag = LlamaIndexRAG( - model=self.model, - emb_model=self.emb_model, - ) - # load the document to memory - # Feed the AgentScope tutorial documents, so that - # the agent can answer questions related to AgentScope! - docs = self.rag.load_data( - loader=SimpleDirectoryReader( - "../../docs/sphinx_doc/en/source/tutorial", - ), - ) - self.rag.store_and_index(docs) + + # init rag as None + # MUST USE LlamaIndexAgent OR LangChainAgent + self.rag = None + self.config = config or {} def reply( self, x: dict = None, - min_score: float = 0, ) -> dict: """ Reply function of the LlamaIndex agent. @@ -73,8 +67,6 @@ def reply( A dictionary representing the user's input to the agent. This input is added to the memory if provided. Defaults to None. - min_score (`float`, defaults to `0`): - A minimum score of the retrieved chunk to be used in the prompt Returns: A dictionary representing the message generated by the agent in response to the user's input. @@ -85,10 +77,9 @@ def reply( self.memory.add(x) # retrieve when the input is not None content = x.get("content", "") - retrieved_docs = self.rag.retrieve(content) - for node in retrieved_docs: - if node.get_score() > min_score: - retrieved_docs_to_string += node.get_text() + retrieved_docs = self.rag.retrieve(content, to_list_strs=True) + for content in retrieved_docs: + retrieved_docs_to_string += content # prepare prompt prompt = self.engine.join( { @@ -112,3 +103,83 @@ def reply( self.memory.add(msg) return msg + + +class LlamaIndexAgent(RAGAgent): + """ + A LlamaIndex agent build on LlamaIndex. + """ + + def __init__( + self, + name: str, + sys_prompt: Optional[str] = None, + model_config_name: str = None, + emb_model_config_name: str = None, + use_memory: bool = True, + memory_config: Optional[dict] = None, + prompt_type: Optional[PromptType] = PromptType.LIST, + config: Optional[dict] = None, + ) -> None: + super().__init__( + name=name, + sys_prompt=sys_prompt, + model_config_name=model_config_name, + emb_model_config_name=emb_model_config_name, + use_memory=use_memory, + memory_config=memory_config, + prompt_type=prompt_type, + config=config, + ) + # init rag related attributes + self.rag = LlamaIndexRAG( + model=self.model, + emb_model=self.emb_model, + ) + # load the document to memory + # Feed the AgentScope tutorial documents, so that + # the agent can answer questions related to AgentScope! + docs = self.rag.load_data( + loader=SimpleDirectoryReader(self.config["data_path"]), + ) + self.rag.store_and_index(docs) + + +class LangChainRAGAgent(RAGAgent): + """ + A LlamaIndex agent build on LlamaIndex. + """ + + def __init__( + self, + name: str, + sys_prompt: Optional[str] = None, + model_config_name: str = None, + emb_model_config_name: str = None, + use_memory: bool = True, + memory_config: Optional[dict] = None, + prompt_type: Optional[PromptType] = PromptType.LIST, + config: Optional[dict] = None, + ) -> None: + super().__init__( + name=name, + sys_prompt=sys_prompt, + model_config_name=model_config_name, + emb_model_config_name=emb_model_config_name, + use_memory=use_memory, + memory_config=memory_config, + prompt_type=prompt_type, + config=config, + ) + # init rag related attributes + self.rag = LangChainRAG( + model=self.model, + emb_model=self.emb_model, + ) + # load the document to memory + # Feed the AgentScope tutorial documents, so that + # the agent can answer questions related to AgentScope! + docs = self.rag.load_data( + loader=DirectoryLoader(self.config["data_path"]), + ) + self.rag.store_and_index(docs) diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index 8f839b66b..87f1b6cc1 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -4,7 +4,8 @@ an agent with RAG capability. """ import os -from rag_agents import LlamaIndexAgent +import argparse +from rag_agents import LlamaIndexAgent, LangChainRAGAgent import agentscope from agentscope.agents import UserAgent @@ -29,15 +30,27 @@ def main() -> None: ], ) - # Init two agents - rag_agent = LlamaIndexAgent( - name="Assistant", - sys_prompt="You're a helpful assistant. You need to generate answers " - "based on the provided context:\n " - "Context: \n {retrieved_context}\n ", - model_config_name="qwen_config", # replace by your model config name - emb_model_config_name="qwen_emb_config", - ) + # Init RAG agent and user + if args.module == "llamaindex": + rag_agent = LlamaIndexAgent( + name="Assistant", + sys_prompt="You're a helpful assistant. You need to generate " + "answers based on the provided context:\n " + "Context: \n {retrieved_context}\n ", + model_config_name="qwen_config", # model config name + emb_model_config_name="qwen_emb_config", + config={"data_path": args.data_path}, + ) + else: + rag_agent = LangChainRAGAgent( + name="Assistant", + sys_prompt="You're a helpful assistant. You need to generate" + " answers based on the provided context:\n " + "Context: \n {retrieved_context}\n ", + model_config_name="qwen_config", # your model config name + emb_model_config_name="qwen_emb_config", + config={"data_path": args.data_path}, + ) user_agent = UserAgent() # start the conversation between user and assistant while True: @@ -49,4 +62,16 @@ def main() -> None: if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--module", + choices=["llamaindex", "langchain"], + default="llamaindex", + ) + parser.add_argument( + "--data_path", + type=str, + default="./data/", + ) + args = parser.parse_args() main() diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index e033ca54a..3211f2e72 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -4,19 +4,45 @@ """ -from typing import Type, Any, Optional +from typing import Any, Optional # from pathlib import Path -from langchain_community.document_loaders.base import BaseLoader from langchain_core.vectorstores import VectorStore +from langchain_core.documents import Document +from langchain_core.embeddings import Embeddings +from langchain_community.document_loaders.base import BaseLoader +from langchain_community.vectorstores import Chroma from langchain_text_splitters.base import TextSplitter -from langchain_openai import OpenAIEmbeddings + +# from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_text_splitters import CharacterTextSplitter from agentscope.rag import RAGBase +from agentscope.rag.rag import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE from agentscope.models import ModelWrapperBase +class _LangChainEmbModel(Embeddings): + def __init__(self, emb_model: ModelWrapperBase): + self._emb_model_wrapper = emb_model + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """ + Wrapper function for embedding list of documents + """ + results = [ + list(self._emb_model_wrapper(t).embedding[0]) for t in texts + ] + return results + + def embed_query(self, text: str) -> list[float]: + """ + Wrapper function for embedding a single query + """ + return list(self._emb_model_wrapper(text).embedding[0]) + + class LangChainRAG(RAGBase): """ This class is a wrapper around the LangChain RAG. @@ -27,47 +53,88 @@ def __init__( self, model: Optional[ModelWrapperBase], emb_model: Optional[ModelWrapperBase], - loader_type: Type[BaseLoader], - splitter_type: Type[TextSplitter], - vector_store_type: Type[VectorStore], - embedding_model: Type[OpenAIEmbeddings], + config: Optional[dict] = None, **kwargs: Any, ) -> None: super().__init__(model, emb_model, **kwargs) - self.loader_type = loader_type - self.splitter_type = splitter_type - self.config = kwargs - self.vector_store_type = vector_store_type + self.loader = None self.splitter = None - self.vector_store = None - self.embedding_model = embedding_model self.retriever = None + self.vector_store = None + + self.config = config or {} + if isinstance(emb_model, ModelWrapperBase): + self.emb_model = _LangChainEmbModel(emb_model) + elif isinstance(emb_model, Embeddings): + self.emb_model = emb_model + else: + raise TypeError( + f"Embedding model does not support {type(self.emb_model)}.", + ) def load_data( self, - loader: Any, - query: Any, + loader: BaseLoader, + query: Optional[Any] = None, **kwargs: Any, - ) -> Any: - """loading data from a directory""" + ) -> list[Document]: + # pylint: disable=unused-argument + """ + loading data from a directory + :param loader: accepting a LangChain loader instance, default is a + :param _: accepting a query, LangChain does not rely on this + :param kwargs: other parameters for loader and splitter + :return: a list of documents + + Notice: currently LangChain supports + """ self.loader = loader docs = self.loader.load() - self.splitter = self.splitter_type(**kwargs) - all_splits = self.splitter.split_documents(docs) - return all_splits + # self.splitter = self.splitter_type(**kwargs) + # all_splits = self.splitter.split_documents(docs) + return docs def store_and_index( self, docs: Any, - vector_store: Any, + vector_store: Optional[VectorStore] = None, + splitter: Optional[TextSplitter] = None, **kwargs: Any, ) -> None: - """indexing the documents and store them into the vector store""" - self.vector_store = self.vector_store_type.from_documents( - documents=docs, - embedding=self.embedding_model(), + """ + Preprocessing the loaded documents. + :param docs: documents to be processed + :param vector_store: vector store + :param retriever: optional, specifies the retriever to use + :param splitter: optional, specifies the splitter to preprocess + the documents + :param kwargs: + + In LlamaIndex terms, an Index is a data structure composed + of Document objects, designed to enable querying by an LLM. + For example: + 1) preprocessing documents with + 2) generate embedding, + 3) store the embedding-content to vdb + """ + self.splitter = splitter or CharacterTextSplitter( + chunk_size=self.config.get("chunk_size", DEFAULT_CHUNK_SIZE), + chunk_overlap=self.config.get( + "chunk_overlap", + DEFAULT_CHUNK_OVERLAP, + ), + ) + all_splits = self.splitter.split_documents(docs) + + # indexing the chunks and store them into the vector store + if vector_store is None: + vector_store = Chroma() + self.vector_store = vector_store.from_documents( + documents=all_splits, + embedding=self.emb_model, ) + # build retriever k = self.config.get("k", 6) search_type = self.config.get("search_type", "similarity") @@ -76,7 +143,19 @@ def store_and_index( search_kwargs={"k": k}, ) - def retrieve(self, query: Any) -> list[Any]: - """retrieve the documents based on the query""" + def retrieve(self, query: Any, to_list_strs: bool = False) -> list[Any]: + """ + This is a basic retrieve function with LangChain APIs + :param query: query is expected to be a question in string + + More advanced retriever can refer to + https://python.langchain.com/docs/modules/data_connection/retrievers/ + """ + retrieved_docs = self.retriever.invoke(query) + if to_list_strs: + results = [] + for doc in retrieved_docs: + results.append(doc.page_content) + return results return retrieved_docs diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index 396114b92..c7a08c8ae 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -97,8 +97,8 @@ def __init__( self.emb_model = emb_model # ensure the emb_model is compatible with LlamaIndex - if issubclass(type(self.emb_model), ModelWrapperBase): - self.emb_model = _EmbeddingModel(self.emb_model) + if isinstance(emb_model, ModelWrapperBase): + self.emb_model = _EmbeddingModel(emb_model) elif isinstance(self.emb_model, BaseEmbedding): pass else: @@ -223,7 +223,7 @@ def set_retriever(self, retriever: BaseRetriever) -> None: """ self.retriever = retriever - def retrieve(self, query: str) -> list[Any]: + def retrieve(self, query: str, to_list_strs: bool = False) -> list[Any]: """ This is a basic retrieve function :param query: query is expected to be a question in string @@ -232,4 +232,9 @@ def retrieve(self, query: str) -> list[Any]: https://docs.llamaindex.ai/en/stable/examples/query_transformations/query_transform_cookbook.html """ retrieved = self.retriever.retrieve(str(query)) + if to_list_strs: + results = [] + for node in retrieved: + results.append(node.get_text()) + return results return retrieved diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index bc0e76b06..a6086b3ef 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -62,7 +62,7 @@ def store_and_index( """ @abstractmethod - def retrieve(self, query: Any) -> list[Any]: + def retrieve(self, query: Any, to_list_strs: bool = False) -> list[Any]: """ retrieve list of content from vdb to memory """ From 7c3a2a4d23b816cec549efcea9e0cce14dcc0c03 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 17:16:20 +0800 Subject: [PATCH 08/21] update example --- examples/rag/rag_agents.py | 2 ++ examples/rag/rag_example.py | 30 +++++++++++------------------- 2 files changed, 13 insertions(+), 19 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 9d8efdf8a..fad261a3f 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -80,6 +80,8 @@ def reply( retrieved_docs = self.rag.retrieve(content, to_list_strs=True) for content in retrieved_docs: retrieved_docs_to_string += content + + self.speak("[retrieved]:" + retrieved_docs_to_string) # prepare prompt prompt = self.engine.join( { diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index 87f1b6cc1..16200e945 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -30,27 +30,19 @@ def main() -> None: ], ) - # Init RAG agent and user if args.module == "llamaindex": - rag_agent = LlamaIndexAgent( - name="Assistant", - sys_prompt="You're a helpful assistant. You need to generate " - "answers based on the provided context:\n " - "Context: \n {retrieved_context}\n ", - model_config_name="qwen_config", # model config name - emb_model_config_name="qwen_emb_config", - config={"data_path": args.data_path}, - ) + AgentClass = LlamaIndexAgent else: - rag_agent = LangChainRAGAgent( - name="Assistant", - sys_prompt="You're a helpful assistant. You need to generate" - " answers based on the provided context:\n " - "Context: \n {retrieved_context}\n ", - model_config_name="qwen_config", # your model config name - emb_model_config_name="qwen_emb_config", - config={"data_path": args.data_path}, - ) + AgentClass = LangChainRAGAgent + rag_agent = AgentClass( + name="Assistant", + sys_prompt="You're a helpful assistant. You need to generate" + " answers based on the provided context:\n " + "Context: \n {retrieved_context}\n ", + model_config_name="qwen_config", # your model config name + emb_model_config_name="qwen_emb_config", + config={"data_path": args.data_path}, + ) user_agent = UserAgent() # start the conversation between user and assistant while True: From d572fd0723cdbf2df28863d8dbaa48b1fc9df482 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 18:17:33 +0800 Subject: [PATCH 09/21] update parameter for retrieval --- examples/rag/__init__.py | 0 examples/rag/rag_agents.py | 6 ++++-- examples/rag/rag_example.py | 18 ++++++++++++++++-- src/agentscope/rag/llama_index_rag.py | 12 +++++++++++- src/agentscope/rag/rag.py | 1 + 5 files changed, 32 insertions(+), 5 deletions(-) delete mode 100644 examples/rag/__init__.py diff --git a/examples/rag/__init__.py b/examples/rag/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index fad261a3f..5a67dcfc0 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -54,7 +54,7 @@ def reply( x: dict = None, ) -> dict: """ - Reply function of the LlamaIndex agent. + Reply function of the RAG agent. Processes the input data, 1) use the input data to retrieve with RAG function; 2) generates a prompt using the current memory and system @@ -79,7 +79,7 @@ def reply( content = x.get("content", "") retrieved_docs = self.rag.retrieve(content, to_list_strs=True) for content in retrieved_docs: - retrieved_docs_to_string += content + retrieved_docs_to_string += "\n>>>> " + content self.speak("[retrieved]:" + retrieved_docs_to_string) # prepare prompt @@ -137,6 +137,7 @@ def __init__( self.rag = LlamaIndexRAG( model=self.model, emb_model=self.emb_model, + config=config, ) # load the document to memory # Feed the AgentScope tutorial documents, so that @@ -177,6 +178,7 @@ def __init__( self.rag = LangChainRAG( model=self.model, emb_model=self.emb_model, + config=config, ) # load the document to memory # Feed the AgentScope tutorial documents, so that diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index 16200e945..c63691239 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -5,6 +5,8 @@ """ import os import argparse +from loguru import logger + from rag_agents import LlamaIndexAgent, LangChainRAGAgent import agentscope from agentscope.agents import UserAgent @@ -41,7 +43,12 @@ def main() -> None: "Context: \n {retrieved_context}\n ", model_config_name="qwen_config", # your model config name emb_model_config_name="qwen_emb_config", - config={"data_path": args.data_path}, + config={ + "data_path": args.data_path, + "chunk_size": 2048, + "chunk_overlap": 40, + "similarity_top_k": 10, + }, ) user_agent = UserAgent() # start the conversation between user and assistant @@ -54,6 +61,8 @@ def main() -> None: if __name__ == "__main__": + # The default parameters can set a AgentScope consultant to + # answer question about agentscope based on tutorial. parser = argparse.ArgumentParser() parser.add_argument( "--module", @@ -63,7 +72,12 @@ def main() -> None: parser.add_argument( "--data_path", type=str, - default="./data/", + default="../../docs/sphinx_doc/en/source/tutorial", ) args = parser.parse_args() + if args.module == "langchain": + logger.warning( + "LangChain RAG Chosen. May require install pandoc in advanced.", + "For example, run ` brew install pandoc` on MacOS.", + ) main() diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index c7a08c8ae..7b111bd77 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -25,7 +25,11 @@ ) from agentscope.rag import RAGBase -from agentscope.rag.rag import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP +from agentscope.rag.rag import ( + DEFAULT_CHUNK_SIZE, + DEFAULT_CHUNK_OVERLAP, + DEFAULT_TOP_K, +) from agentscope.models import ModelWrapperBase @@ -95,6 +99,7 @@ def __init__( self.index = None self.persist_dir = kwargs.get("persist_dir", "./") self.emb_model = emb_model + print(self.config) # ensure the emb_model is compatible with LlamaIndex if isinstance(emb_model, ModelWrapperBase): @@ -209,8 +214,13 @@ def store_and_index( # set the retriever if retriever is None: + print(self.config.get("similarity_top_k", DEFAULT_TOP_K)) self.retriever = self.index.as_retriever( embed_model=self.emb_model, + similarity_top_k=self.config.get( + "similarity_top_k", + DEFAULT_TOP_K, + ), **kwargs, ) else: diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index a6086b3ef..c3a2c12ea 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -19,6 +19,7 @@ DEFAULT_CHUNK_SIZE = 1024 DEFAULT_CHUNK_OVERLAP = 20 +DEFAULT_TOP_K = 5 class RAGBase(ABC): From 3c894ef2d7d9a6c2a7d35b77639d38eae1a84c8b Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 19:33:09 +0800 Subject: [PATCH 10/21] update langchain parameter --- src/agentscope/rag/langchain_rag.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index 3211f2e72..aac070229 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -136,11 +136,13 @@ def store_and_index( ) # build retriever - k = self.config.get("k", 6) search_type = self.config.get("search_type", "similarity") self.retriever = self.vector_store.as_retriever( search_type=search_type, - search_kwargs={"k": k}, + search_kwargs={ + "k": self.config.get("similarity_top_k", 6), + "score_threshold": self.config.get("score_threshold", 0), + }, ) def retrieve(self, query: Any, to_list_strs: bool = False) -> list[Any]: From be8d2d19e6799b15485bda0324d20a01848fd5d7 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Tue, 19 Mar 2024 19:40:51 +0800 Subject: [PATCH 11/21] fix --- src/agentscope/rag/langchain_rag.py | 1 - src/agentscope/rag/rag.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index aac070229..f3c6692e3 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -141,7 +141,6 @@ def store_and_index( search_type=search_type, search_kwargs={ "k": self.config.get("similarity_top_k", 6), - "score_threshold": self.config.get("score_threshold", 0), }, ) diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index c3a2c12ea..ca84810a1 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -86,5 +86,5 @@ def post_processing( self.postprocessing_model(prompt.format(retrieved_docs)) """ assert self.postprocessing_model - prompt = prompt.format(retrieved_docs) + prompt = prompt.format("\n".join(retrieved_docs)) return self.postprocessing_model(prompt, **kwargs).text From e2eb025001c8ebc545d0f323b5fb74e6e77844c7 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 20 Mar 2024 20:10:32 +0800 Subject: [PATCH 12/21] update doc strings --- examples/rag/rag_agents.py | 40 +++++--- examples/rag/rag_example.py | 4 +- src/agentscope/rag/langchain_rag.py | 108 ++++++++++++++------ src/agentscope/rag/llama_index_rag.py | 138 ++++++++++++++++++++------ src/agentscope/rag/rag.py | 34 ++++++- 5 files changed, 240 insertions(+), 84 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 5a67dcfc0..ae3ba4205 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -4,8 +4,6 @@ """ from typing import Optional -from llama_index.core import SimpleDirectoryReader -from langchain_community.document_loaders import DirectoryLoader from agentscope.prompt import PromptType from agentscope.agents.agent import AgentBase @@ -28,7 +26,6 @@ def __init__( sys_prompt: Optional[str] = None, model_config_name: str = None, emb_model_config_name: str = None, - use_memory: bool = True, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, @@ -37,7 +34,7 @@ def __init__( name=name, sys_prompt=sys_prompt, model_config_name=model_config_name, - use_memory=use_memory, + use_memory=True, memory_config=memory_config, ) # init prompt engine @@ -48,6 +45,8 @@ def __init__( # MUST USE LlamaIndexAgent OR LangChainAgent self.rag = None self.config = config or {} + if "log_retrieval" not in self.config: + self.config["log_retrieval"] = True def reply( self, @@ -73,25 +72,28 @@ def reply( """ retrieved_docs_to_string = "" # record the input if needed - if x is not None: + if self.memory: self.memory.add(x) + + if x is not None: # retrieve when the input is not None content = x.get("content", "") retrieved_docs = self.rag.retrieve(content, to_list_strs=True) for content in retrieved_docs: retrieved_docs_to_string += "\n>>>> " + content - self.speak("[retrieved]:" + retrieved_docs_to_string) + if self.config["log_retrieval"]: + self.speak("[retrieved]:" + retrieved_docs_to_string) + # prepare prompt prompt = self.engine.join( { "role": "system", - "content": self.sys_prompt.format_map( - {"retrieved_context": retrieved_docs_to_string}, - ), + "content": self.sys_prompt, }, # {"role": "system", "content": retrieved_docs_to_string}, self.memory.get_memory(), + "Context: " + retrieved_docs_to_string, ) # call llm and generate response @@ -118,7 +120,6 @@ def __init__( sys_prompt: Optional[str] = None, model_config_name: str = None, emb_model_config_name: str = None, - use_memory: bool = True, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, @@ -128,11 +129,17 @@ def __init__( sys_prompt=sys_prompt, model_config_name=model_config_name, emb_model_config_name=emb_model_config_name, - use_memory=use_memory, memory_config=memory_config, prompt_type=prompt_type, config=config, ) + try: + from llama_index.core import SimpleDirectoryReader + except ImportError as exc: + raise ImportError( + " LlamaIndexAgent requires llama-index to be install." + "Please run `pip install llama-index`", + ) from exc # init rag related attributes self.rag = LlamaIndexRAG( model=self.model, @@ -159,7 +166,6 @@ def __init__( sys_prompt: Optional[str] = None, model_config_name: str = None, emb_model_config_name: str = None, - use_memory: bool = True, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, @@ -169,11 +175,19 @@ def __init__( sys_prompt=sys_prompt, model_config_name=model_config_name, emb_model_config_name=emb_model_config_name, - use_memory=use_memory, memory_config=memory_config, prompt_type=prompt_type, config=config, ) + try: + from langchain_community.document_loaders import DirectoryLoader + except ImportError as exc: + raise ImportError( + "LangChainRAGAgent requires LangChain related packages " + "installed. Please run `pip install langchain " + "unstructured[all-docs] langchain-text-splitters`", + ) from exc + # init rag related attributes self.rag = LangChainRAG( model=self.model, diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index c63691239..c997ea292 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -39,8 +39,7 @@ def main() -> None: rag_agent = AgentClass( name="Assistant", sys_prompt="You're a helpful assistant. You need to generate" - " answers based on the provided context:\n " - "Context: \n {retrieved_context}\n ", + " answers based on the provided context:\n ", model_config_name="qwen_config", # your model config name emb_model_config_name="qwen_emb_config", config={ @@ -48,6 +47,7 @@ def main() -> None: "chunk_size": 2048, "chunk_overlap": 40, "similarity_top_k": 10, + "log_retrieval": True, }, ) user_agent = UserAgent() diff --git a/src/agentscope/rag/langchain_rag.py b/src/agentscope/rag/langchain_rag.py index f3c6692e3..d9b176691 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/src/agentscope/rag/langchain_rag.py @@ -4,19 +4,24 @@ """ -from typing import Any, Optional - -# from pathlib import Path - -from langchain_core.vectorstores import VectorStore -from langchain_core.documents import Document -from langchain_core.embeddings import Embeddings -from langchain_community.document_loaders.base import BaseLoader -from langchain_community.vectorstores import Chroma -from langchain_text_splitters.base import TextSplitter - -# from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_text_splitters import CharacterTextSplitter +from typing import Any, Optional, Union + +try: + from langchain_core.vectorstores import VectorStore + from langchain_core.documents import Document + from langchain_core.embeddings import Embeddings + from langchain_community.document_loaders.base import BaseLoader + from langchain_community.vectorstores import Chroma + from langchain_text_splitters.base import TextSplitter + from langchain_text_splitters import CharacterTextSplitter +except ImportError: + VectorStore = None + Document = None + Embeddings = None + BaseLoader = None + Chroma = None + TextSplitter = None + CharacterTextSplitter = None from agentscope.rag import RAGBase from agentscope.rag.rag import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE @@ -24,12 +29,25 @@ class _LangChainEmbModel(Embeddings): - def __init__(self, emb_model: ModelWrapperBase): + """ + Dummy wrapper to convert the ModelWrapperBase embedding model + to a LanguageChain RAG model + """ + + def __init__(self, emb_model: ModelWrapperBase) -> None: + """ + Dummy wrapper + Args: + emb_model (ModelWrapperBase): embedding model of + ModelWrapperBase type + """ self._emb_model_wrapper = emb_model def embed_documents(self, texts: list[str]) -> list[list[float]]: """ Wrapper function for embedding list of documents + Args: + texts (list[str]): list of texts to be embedded """ results = [ list(self._emb_model_wrapper(t).embedding[0]) for t in texts @@ -39,6 +57,8 @@ def embed_documents(self, texts: list[str]) -> list[list[float]]: def embed_query(self, text: str) -> list[float]: """ Wrapper function for embedding a single query + Args: + text (str): query to be embedded """ return list(self._emb_model_wrapper(text).embedding[0]) @@ -46,16 +66,25 @@ def embed_query(self, text: str) -> list[float]: class LangChainRAG(RAGBase): """ This class is a wrapper around the LangChain RAG. - TODO: still under construction """ def __init__( self, model: Optional[ModelWrapperBase], - emb_model: Optional[ModelWrapperBase], + emb_model: Union[ModelWrapperBase, Embeddings, None], config: Optional[dict] = None, **kwargs: Any, ) -> None: + """ + Initializes the LangChainRAG + Args: + model (ModelWrapperBase): + The language model used for final synthesis + emb_model ( Union[ModelWrapperBase, Embeddings, None]): + The embedding model used for generate embeddings + config (dict): + The additional configuration for llama index rag + """ super().__init__(model, emb_model, **kwargs) self.loader = None @@ -63,6 +92,11 @@ def __init__( self.retriever = None self.vector_store = None + if VectorStore is None: + raise ImportError( + "Please install LangChain RAG packages to use LangChain RAG.", + ) + self.config = config or {} if isinstance(emb_model, ModelWrapperBase): self.emb_model = _LangChainEmbModel(emb_model) @@ -81,18 +115,17 @@ def load_data( ) -> list[Document]: # pylint: disable=unused-argument """ - loading data from a directory - :param loader: accepting a LangChain loader instance, default is a - :param _: accepting a query, LangChain does not rely on this - :param kwargs: other parameters for loader and splitter - :return: a list of documents - - Notice: currently LangChain supports + Loading data from a directory + Args: + loader (BaseLoader): + accepting a LangChain loader instance + query (str): + accepting a query, LangChain does not rely on this + Returns: + list[Document]: a list of documents loaded """ self.loader = loader docs = self.loader.load() - # self.splitter = self.splitter_type(**kwargs) - # all_splits = self.splitter.split_documents(docs) return docs def store_and_index( @@ -101,15 +134,20 @@ def store_and_index( vector_store: Optional[VectorStore] = None, splitter: Optional[TextSplitter] = None, **kwargs: Any, - ) -> None: + ) -> Any: """ Preprocessing the loaded documents. - :param docs: documents to be processed - :param vector_store: vector store - :param retriever: optional, specifies the retriever to use - :param splitter: optional, specifies the splitter to preprocess - the documents - :param kwargs: + Args: + docs (Any): + documents to be processed + vector_store (Optional[VectorStore]): + vector store in LangChain RAG + splitter (Optional[TextSplitter]): + optional, specifies the splitter to preprocess + the documents + + Returns: + None In LlamaIndex terms, an Index is a data structure composed of Document objects, designed to enable querying by an LLM. @@ -147,7 +185,11 @@ def store_and_index( def retrieve(self, query: Any, to_list_strs: bool = False) -> list[Any]: """ This is a basic retrieve function with LangChain APIs - :param query: query is expected to be a question in string + Args: + query: query is expected to be a question in string + + Returns: + list of answers More advanced retriever can refer to https://python.langchain.com/docs/modules/data_connection/retrievers/ diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index 7b111bd77..60b0e1838 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -5,24 +5,34 @@ """ from typing import Any, Optional, List, Union +from loguru import logger -from llama_index.core.readers.base import BaseReader -from llama_index.core.base.base_retriever import BaseRetriever -from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding -from llama_index.core.ingestion import IngestionPipeline -from llama_index.core.vector_stores.types import ( - BasePydanticVectorStore, - VectorStore, -) -from llama_index.core.bridge.pydantic import PrivateAttr -from llama_index.core.node_parser.interface import NodeParser -from llama_index.core.node_parser import SentenceSplitter - - -from llama_index.core import ( - VectorStoreIndex, - SimpleDirectoryReader, -) +try: + from llama_index.core.readers.base import BaseReader + from llama_index.core.base.base_retriever import BaseRetriever + from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding + from llama_index.core.ingestion import IngestionPipeline + from llama_index.core.vector_stores.types import ( + BasePydanticVectorStore, + VectorStore, + ) + from llama_index.core.bridge.pydantic import PrivateAttr + from llama_index.core.node_parser.interface import NodeParser + from llama_index.core.node_parser import SentenceSplitter + from llama_index.core import ( + VectorStoreIndex, + SimpleDirectoryReader, + ) +except ImportError: + BaseReader, BaseRetriever, BaseEmbedding, Embedding = ( + None, + None, + None, + None, + ) + IngestionPipeline, BasePydanticVectorStore, VectorStore = None, None, None + PrivateAttr, NodeParser, SentenceSplitter = None, None, None + VectorStoreIndex, SimpleDirectoryReader = None, None from agentscope.rag import RAGBase from agentscope.rag.rag import ( @@ -46,6 +56,14 @@ def __init__( emb_model: ModelWrapperBase, embed_batch_size: int = 1, ) -> None: + """ + Dummy wrapper to convert a ModelWrapperBase to llama Index + embedding model + + Args: + emb_model (ModelWrapperBase): embedding model in ModelWrapperBase + embed_batch_size (int): batch size, defaults to 1 + """ super().__init__( model_name="Temporary_embedding_wrapper", embed_batch_size=embed_batch_size, @@ -53,16 +71,31 @@ def __init__( self._emb_model_wrapper = emb_model def _get_query_embedding(self, query: str) -> List[float]: + """ + get embedding for query + Args: + query (str): query to be embedded + """ # Note: AgentScope embedding model wrapper returns list of embedding return list(self._emb_model_wrapper(query).embedding[0]) def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]: + """ + get embedding for list of strings + Args: + texts ( List[str]): texts to be embedded + """ results = [ list(self._emb_model_wrapper(t).embedding[0]) for t in texts ] return results def _get_text_embedding(self, text: str) -> Embedding: + """ + get embedding for a single string + Args: + text (str): texts to be embedded + """ return list(self._emb_model_wrapper(text).embedding[0]) # TODO: use proper async methods, but depends on model wrapper @@ -84,16 +117,26 @@ async def _aget_text_embeddings( class LlamaIndexRAG(RAGBase): """ - This class is a wrapper around the Llama index RAG. + This class is a wrapper with the llama index RAG. """ def __init__( self, model: Optional[ModelWrapperBase], - emb_model: Any = None, + emb_model: Union[ModelWrapperBase, BaseEmbedding, None] = None, config: Optional[dict] = None, **kwargs: Any, ) -> None: + """ + RAG component based on llama index. + Args: + model (ModelWrapperBase): + The language model used for final synthesis + emb_model (Optional[ModelWrapperBase]): + The embedding model used for generate embeddings + config (dict): + The additional configuration for llama index rag + """ super().__init__(model, emb_model, config, **kwargs) self.retriever = None self.index = None @@ -119,10 +162,15 @@ def load_data( ) -> Any: """ Accept a loader, loading the desired data (no chunking) - :param loader: object to load data, expected be an instance of class - inheriting from BaseReader. - :param query: optional, used when the data is in a database. - :return: the loaded documents (un-chunked) + Args: + loader (BaseReader): + object to load data, expected be an instance of class + inheriting from BaseReader in llama index. + query (Optional[str]): + optional, used when the data is in a database. + + Returns: + Any: loaded documents Example 1: use simple directory loader to load general documents, including Markdown, PDFs, Word documents, PowerPoint decks, images, @@ -152,6 +200,7 @@ def load_data( documents = loader.load_data() else: documents = loader.load_data(query) + logger.info(f"loaded {len(documents)} documents") return documents def store_and_index( @@ -164,12 +213,21 @@ def store_and_index( ) -> Any: """ Preprocessing the loaded documents. - :param docs: documents to be processed - :param vector_store: vector store - :param retriever: optional, specifies the retriever to use - :param transformations: optional, specifies the transformations - to preprocess the documents - :param kwargs: + Args: + docs (Any): + documents to be processed, usually expected to be in + llama index Documents. + vector_store (Union[BasePydanticVectorStore, VectorStore, None]): + vector store in llama index + retriever (Optional[BaseRetriever]): + optional, specifies the retriever in llama index to be used + transformations (Optional[list[NodeParser]]): + optional, specifies the transformations (operators) to + process documents (e.g., split the documents into smaller + chunks) + + Return: + Any: return the index of the processed document In LlamaIndex terms, an Index is a data structure composed of Document objects, designed to enable querying by an LLM. @@ -210,11 +268,16 @@ def store_and_index( transformations=transformations, ) nodes = pipeline.run(documents=docs) - self.index = VectorStoreIndex(nodes=nodes) + self.index = VectorStoreIndex( + nodes=nodes, + embed_model=self.emb_model, + ) # set the retriever if retriever is None: - print(self.config.get("similarity_top_k", DEFAULT_TOP_K)) + logger.info( + f'{self.config.get("similarity_top_k", DEFAULT_TOP_K)}', + ) self.retriever = self.index.as_retriever( embed_model=self.emb_model, similarity_top_k=self.config.get( @@ -230,13 +293,24 @@ def store_and_index( def set_retriever(self, retriever: BaseRetriever) -> None: """ Reset the retriever if necessary. + Args: + retriever (BaseRetriever): passing a retriever in llama index. """ self.retriever = retriever def retrieve(self, query: str, to_list_strs: bool = False) -> list[Any]: """ This is a basic retrieve function - :param query: query is expected to be a question in string + Args: + query (str): + query is expected to be a question in string + to_list_strs (book): + whether returns the list of strings; + if False, return NodeWithScore + + Return: + list[Any]: list of str or NodeWithScore + More advanced query processing can refer to https://docs.llamaindex.ai/en/stable/examples/query_transformations/query_transform_cookbook.html diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index ca84810a1..746ec4d57 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -45,7 +45,13 @@ def load_data( **kwargs: Any, ) -> Any: """ - load data (documents) from disk to memory and chunking them + Load data (documents) from disk to memory and chunking them + Args: + loader (Any): data loader, depending on the package + query (str): query for getting data from DB + + Returns: + Any: loaded documents """ @abstractmethod @@ -56,6 +62,16 @@ def store_and_index( **kwargs: Any, ) -> Any: """ + Store and index the documents. + Args: + docs (Any): + documents to be processed, stored and indexed + vector_store (Any): + vector store to store the index and/or documents + + Returns: + Any: can be indices, depending on the RAG package + preprocessing the loaded documents, for example: 1) chunking, 2) generate embedding, @@ -66,6 +82,12 @@ def store_and_index( def retrieve(self, query: Any, to_list_strs: bool = False) -> list[Any]: """ retrieve list of content from vdb to memory + Args: + query (Any): query to retrieve + to_list_strs (bool): whether return a list of str + + Returns: + return a list with retrieved documents (in strings) """ def post_processing( @@ -77,10 +99,14 @@ def post_processing( """ A default solution for post-processing function, generates answer based on the retrieved documents. - :param retrieved_docs: list of retrieved documents - :param prompt: prompt for LLM generating answer with the - retrieved documents + Args: + retrieved_docs (list[str]): + list of retrieved documents + prompt (str): + prompt for LLM generating answer with the retrieved documents + Returns: + Any: a synthesized answer from LLM with retrieved documents Example: self.postprocessing_model(prompt.format(retrieved_docs)) From f038021b38e0b29db9e06bacac146169607e74f3 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 20 Mar 2024 20:51:51 +0800 Subject: [PATCH 13/21] update doc strings 2 --- examples/rag/rag_agents.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index ae3ba4205..6916ca987 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -30,6 +30,17 @@ def __init__( prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, ) -> None: + """ + Initialize the RAG base agent + Args: + name (str): the name for the agent + sys_prompt (str): system prompt for the RAG agent + model_config_name (str): language model for the agent + emb_model_config_name (str): embedding model for the agent + memory_config (dict): memory configuration + prompt_type (PromptType): prompt type, list or str + config (dict): additional config for RAG and agent + """ super().__init__( name=name, sys_prompt=sys_prompt, @@ -124,6 +135,17 @@ def __init__( prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, ) -> None: + """ + Initialize the RAG LlamaIndexAgent + Args: + name (str): the name for the agent + sys_prompt (str): system prompt for the RAG agent + model_config_name (str): language model for the agent + emb_model_config_name (str): embedding model for the agent + memory_config (dict): memory configuration + prompt_type (PromptType): prompt type, list or str + config (dict): additional config for RAG and agent + """ super().__init__( name=name, sys_prompt=sys_prompt, @@ -170,6 +192,17 @@ def __init__( prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, ) -> None: + """ + Initialize the RAG LangChainRAGAgent + Args: + name (str): the name for the agent + sys_prompt (str): system prompt for the RAG agent + model_config_name (str): language model for the agent + emb_model_config_name (str): embedding model for the agent + memory_config (dict): memory configuration + prompt_type (PromptType): prompt type, list or str + config (dict): additional config for RAG and agent + """ super().__init__( name=name, sys_prompt=sys_prompt, From c07b9a78de657c1bfaf7d3500d5803d78965980a Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Thu, 21 Mar 2024 13:01:15 +0800 Subject: [PATCH 14/21] update following comments --- examples/rag/rag_agents.py | 119 ++++++++++++++++++++++++++++++------ examples/rag/rag_example.py | 16 ++++- src/agentscope/rag/rag.py | 4 +- 3 files changed, 119 insertions(+), 20 deletions(-) diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py index 6916ca987..b1120ad81 100644 --- a/examples/rag/rag_agents.py +++ b/examples/rag/rag_agents.py @@ -4,6 +4,7 @@ """ from typing import Optional +from loguru import logger from agentscope.prompt import PromptType from agentscope.agents.agent import AgentBase @@ -14,7 +15,7 @@ from agentscope.rag.langchain_rag import LangChainRAG -class RAGAgent(AgentBase): +class RAGAgentBase(AgentBase): """ Base class for RAG agents, child classes include the RAG agents built with LlamaIndex and LangChain in this file @@ -23,9 +24,9 @@ class RAGAgent(AgentBase): def __init__( self, name: str, - sys_prompt: Optional[str] = None, - model_config_name: str = None, - emb_model_config_name: str = None, + sys_prompt: str, + model_config_name: str, + emb_model_config_name: str, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, @@ -40,6 +41,24 @@ def __init__( memory_config (dict): memory configuration prompt_type (PromptType): prompt type, list or str config (dict): additional config for RAG and agent + Current support adjustable parameters includes: + "data_path": + path to data directory where data is stored, + "chunk_size": + maximum chunk size for preprocessed documents, + default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE + "chunk_overlap": + overlap between preprocessed chunks, default is + agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP + "similarity_top_k": + number of chunks in each retrieval, default + is agentscope.rag.rag.DEFAULT_TOP_K + "log_retrieval" (bool): + whether to user agent.speak() to print the + retrieved documents, default is False. + "recent_n_mem": + how many messages in memory is used as query + for retrieval if memory is used, default is 1 """ super().__init__( name=name, @@ -85,11 +104,28 @@ def reply( # record the input if needed if self.memory: self.memory.add(x) + # in case no input is provided (e.g., in msghub), + # use the memory as query + history = self.engine.join( + self.memory.get_memory( + recent_n=self.config.get("recent_n_mem", 1), + ), + ) + query = ( + "/n".join( + [msg["content"] for msg in history], + ) + if isinstance(history, list) + else str(history) + ) + elif x is not None: + query = x["content"] + else: + query = "" - if x is not None: - # retrieve when the input is not None - content = x.get("content", "") - retrieved_docs = self.rag.retrieve(content, to_list_strs=True) + if len(query) > 0: + # when content has information, do retrieval + retrieved_docs = self.rag.retrieve(query, to_list_strs=True) for content in retrieved_docs: retrieved_docs_to_string += "\n>>>> " + content @@ -114,13 +150,14 @@ def reply( # Print/speak the message in this agent's voice self.speak(msg) - # Record the message in memory - self.memory.add(msg) + if self.memory: + # Record the message in memory + self.memory.add(msg) return msg -class LlamaIndexAgent(RAGAgent): +class LlamaIndexAgent(RAGAgentBase): """ A LlamaIndex agent build on LlamaIndex. """ @@ -128,8 +165,8 @@ class LlamaIndexAgent(RAGAgent): def __init__( self, name: str, - sys_prompt: Optional[str] = None, - model_config_name: str = None, + sys_prompt: str, + model_config_name: str, emb_model_config_name: str = None, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, @@ -145,6 +182,24 @@ def __init__( memory_config (dict): memory configuration prompt_type (PromptType): prompt type, list or str config (dict): additional config for RAG and agent + Current support adjustable parameters includes: + "data_path": + path to data directory where data is stored, + "chunk_size": + maximum chunk size for preprocessed documents, + default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE + "chunk_overlap": + overlap between preprocessed chunks, default is + agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP + "similarity_top_k": + number of chunks in each retrieval, default + is agentscope.rag.rag.DEFAULT_TOP_K + "log_retrieval" (bool): + whether to user agent.speak() to print the + retrieved documents, default is False. + "recent_n_mem": + how many messages in memory is used as query + for retrieval if memory is used, default is 1 """ super().__init__( name=name, @@ -171,13 +226,19 @@ def __init__( # load the document to memory # Feed the AgentScope tutorial documents, so that # the agent can answer questions related to AgentScope! + if "data_path" not in self.config: + self.config[" data_path"] = "./data" + logger.warning( + "No data_path provided in RAG agent config," + "use default path `./data`", + ) docs = self.rag.load_data( loader=SimpleDirectoryReader(self.config["data_path"]), ) self.rag.store_and_index(docs) -class LangChainRAGAgent(RAGAgent): +class LangChainRAGAgent(RAGAgentBase): """ A LlamaIndex agent build on LlamaIndex. """ @@ -185,9 +246,9 @@ class LangChainRAGAgent(RAGAgent): def __init__( self, name: str, - sys_prompt: Optional[str] = None, - model_config_name: str = None, - emb_model_config_name: str = None, + sys_prompt: str, + model_config_name: str, + emb_model_config_name: str, memory_config: Optional[dict] = None, prompt_type: Optional[PromptType] = PromptType.LIST, config: Optional[dict] = None, @@ -202,6 +263,24 @@ def __init__( memory_config (dict): memory configuration prompt_type (PromptType): prompt type, list or str config (dict): additional config for RAG and agent + Current support adjustable parameters includes: + "data_path": + path to data directory where data is stored, + "chunk_size": + maximum chunk size for preprocessed documents, + default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE + "chunk_overlap": + overlap between preprocessed chunks, default is + agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP + "similarity_top_k": + number of chunks in each retrieval, default + is agentscope.rag.rag.DEFAULT_TOP_K + "log_retrieval" (bool): + whether to user agent.speak() to print the + retrieved documents, default is False. + "recent_n_mem": + how many messages in memory is used as query + for retrieval if memory is used, default is 1 """ super().__init__( name=name, @@ -230,6 +309,12 @@ def __init__( # load the document to memory # Feed the AgentScope tutorial documents, so that # the agent can answer questions related to AgentScope! + if "data_path" not in self.config: + self.config[" data_path"] = "./data" + logger.warning( + "No data_path provided in RAG agent config," + "use default path `./data`", + ) docs = self.rag.load_data( loader=DirectoryLoader(self.config["data_path"]), ) diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index c997ea292..600bfcfec 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -7,7 +7,6 @@ import argparse from loguru import logger -from rag_agents import LlamaIndexAgent, LangChainRAGAgent import agentscope from agentscope.agents import UserAgent @@ -33,13 +32,25 @@ def main() -> None: ) if args.module == "llamaindex": + try: + from rag_agents import LlamaIndexAgent + except ImportError as exc: + raise ImportError( + "Please install llamaindex packages.", + ) from exc AgentClass = LlamaIndexAgent else: + try: + from rag_agents import LangChainRAGAgent + except ImportError as exc: + raise ImportError( + "Please install LangChain RAG packages.", + ) from exc AgentClass = LangChainRAGAgent rag_agent = AgentClass( name="Assistant", sys_prompt="You're a helpful assistant. You need to generate" - " answers based on the provided context:\n ", + " answers based on the provided context.", model_config_name="qwen_config", # your model config name emb_model_config_name="qwen_emb_config", config={ @@ -48,6 +59,7 @@ def main() -> None: "chunk_overlap": 40, "similarity_top_k": 10, "log_retrieval": True, + "recent_n_mem": 1, }, ) user_agent = UserAgent() diff --git a/src/agentscope/rag/rag.py b/src/agentscope/rag/rag.py index 746ec4d57..0de27ca37 100644 --- a/src/agentscope/rag/rag.py +++ b/src/agentscope/rag/rag.py @@ -23,7 +23,9 @@ class RAGBase(ABC): - """Base class for RAG""" + """ + Base class for RAG, CANNOT be instantiated directly + """ def __init__( self, From 82c79871dbe42510c90c2965215b55380d798b7a Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Thu, 28 Mar 2024 19:38:38 +0800 Subject: [PATCH 15/21] update to be more flexible --- examples/rag/agent_config.json | 79 +++++++ examples/rag/rag_agents.py | 321 --------------------------- examples/rag/rag_example.py | 74 ++----- src/agentscope/agents/__init__.py | 3 + src/agentscope/agents/rag_agents.py | 330 ++++++++++++++++++++++++++++ 5 files changed, 430 insertions(+), 377 deletions(-) create mode 100644 examples/rag/agent_config.json delete mode 100644 examples/rag/rag_agents.py create mode 100644 src/agentscope/agents/rag_agents.py diff --git a/examples/rag/agent_config.json b/examples/rag/agent_config.json new file mode 100644 index 000000000..742cf5a67 --- /dev/null +++ b/examples/rag/agent_config.json @@ -0,0 +1,79 @@ +[ + { + "class": "LlamaIndexAgent", + "args": { + "name": "AgentScope Tutorial Assistant", + "sys_prompt": "You're a helpful assistant. You need to generate answers based on the provided context.", + "model_config_name": "qwen_config", + "emb_model_config_name": "qwen_emb_config", + "rag_config": { + "load_data": { + "loader": { + "create_object": true, + "module": "llama_index.core", + "class": "SimpleDirectoryReader", + "init_args": { + "input_dir": "../../docs/sphinx_doc/en/source/tutorial/", + "required_exts": [".md"] + } + } + }, + "chunk_size": 2048, + "chunk_overlap": 40, + "similarity_top_k": 10, + "log_retrieval": true, + "recent_n_mem": 1 + } + } + }, + { + "class": "LlamaIndexAgent", + "args": { + "name": "AgentScope Framework Code Assistant", + "sys_prompt": "You're a helpful assistant about coding. You can very familiar with the framework code of AgentScope.", + "model_config_name": "qwen_config", + "emb_model_config_name": "qwen_emb_config", + "rag_config": { + "load_data": { + "loader": { + "create_object": true, + "module": "llama_index.core", + "class": "SimpleDirectoryReader", + "init_args": { + "input_dir": "../../src/agentscope", + "recursive": true, + "required_exts": [".py"] + } + } + }, + "store_and_index": { + "transformations": [ + { + "create_object": true, + "module": "llama_index.core.node_parser", + "class": "CodeSplitter", + "init_args": { + "language": "python", + "chunk_lines": 100 + } + } + ] + }, + "chunk_size": 2048, + "chunk_overlap": 40, + "similarity_top_k": 10, + "log_retrieval": true, + "recent_n_mem": 1 + } + } + }, + { + "class": "DialogAgent", + "args": { + "name": "Summarize Assistant", + "sys_prompt": "You are a helpful assistant that can summarize the answers of the previous two messages.", + "model_config_name": "qwen_config", + "use_memory": true + } + } +] \ No newline at end of file diff --git a/examples/rag/rag_agents.py b/examples/rag/rag_agents.py deleted file mode 100644 index b1120ad81..000000000 --- a/examples/rag/rag_agents.py +++ /dev/null @@ -1,321 +0,0 @@ -# -*- coding: utf-8 -*- -""" -This example shows how to build an agent with RAG (backup by LlamaIndex) -""" - -from typing import Optional -from loguru import logger - -from agentscope.prompt import PromptType -from agentscope.agents.agent import AgentBase -from agentscope.prompt import PromptEngine -from agentscope.message import Msg -from agentscope.models import load_model_by_config_name -from agentscope.rag.llama_index_rag import LlamaIndexRAG -from agentscope.rag.langchain_rag import LangChainRAG - - -class RAGAgentBase(AgentBase): - """ - Base class for RAG agents, child classes include the - RAG agents built with LlamaIndex and LangChain in this file - """ - - def __init__( - self, - name: str, - sys_prompt: str, - model_config_name: str, - emb_model_config_name: str, - memory_config: Optional[dict] = None, - prompt_type: Optional[PromptType] = PromptType.LIST, - config: Optional[dict] = None, - ) -> None: - """ - Initialize the RAG base agent - Args: - name (str): the name for the agent - sys_prompt (str): system prompt for the RAG agent - model_config_name (str): language model for the agent - emb_model_config_name (str): embedding model for the agent - memory_config (dict): memory configuration - prompt_type (PromptType): prompt type, list or str - config (dict): additional config for RAG and agent - Current support adjustable parameters includes: - "data_path": - path to data directory where data is stored, - "chunk_size": - maximum chunk size for preprocessed documents, - default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE - "chunk_overlap": - overlap between preprocessed chunks, default is - agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP - "similarity_top_k": - number of chunks in each retrieval, default - is agentscope.rag.rag.DEFAULT_TOP_K - "log_retrieval" (bool): - whether to user agent.speak() to print the - retrieved documents, default is False. - "recent_n_mem": - how many messages in memory is used as query - for retrieval if memory is used, default is 1 - """ - super().__init__( - name=name, - sys_prompt=sys_prompt, - model_config_name=model_config_name, - use_memory=True, - memory_config=memory_config, - ) - # init prompt engine - self.engine = PromptEngine(self.model, prompt_type=prompt_type) - self.emb_model = load_model_by_config_name(emb_model_config_name) - - # init rag as None - # MUST USE LlamaIndexAgent OR LangChainAgent - self.rag = None - self.config = config or {} - if "log_retrieval" not in self.config: - self.config["log_retrieval"] = True - - def reply( - self, - x: dict = None, - ) -> dict: - """ - Reply function of the RAG agent. - Processes the input data, - 1) use the input data to retrieve with RAG function; - 2) generates a prompt using the current memory and system - prompt; - 3) invokes the language model to produce a response. The - response is then formatted and added to the dialogue memory. - - Args: - x (`dict`, defaults to `None`): - A dictionary representing the user's input to the agent. This - input is added to the memory if provided. Defaults to - None. - Returns: - A dictionary representing the message generated by the agent in - response to the user's input. - """ - retrieved_docs_to_string = "" - # record the input if needed - if self.memory: - self.memory.add(x) - # in case no input is provided (e.g., in msghub), - # use the memory as query - history = self.engine.join( - self.memory.get_memory( - recent_n=self.config.get("recent_n_mem", 1), - ), - ) - query = ( - "/n".join( - [msg["content"] for msg in history], - ) - if isinstance(history, list) - else str(history) - ) - elif x is not None: - query = x["content"] - else: - query = "" - - if len(query) > 0: - # when content has information, do retrieval - retrieved_docs = self.rag.retrieve(query, to_list_strs=True) - for content in retrieved_docs: - retrieved_docs_to_string += "\n>>>> " + content - - if self.config["log_retrieval"]: - self.speak("[retrieved]:" + retrieved_docs_to_string) - - # prepare prompt - prompt = self.engine.join( - { - "role": "system", - "content": self.sys_prompt, - }, - # {"role": "system", "content": retrieved_docs_to_string}, - self.memory.get_memory(), - "Context: " + retrieved_docs_to_string, - ) - - # call llm and generate response - response = self.model(prompt).text - msg = Msg(self.name, response) - - # Print/speak the message in this agent's voice - self.speak(msg) - - if self.memory: - # Record the message in memory - self.memory.add(msg) - - return msg - - -class LlamaIndexAgent(RAGAgentBase): - """ - A LlamaIndex agent build on LlamaIndex. - """ - - def __init__( - self, - name: str, - sys_prompt: str, - model_config_name: str, - emb_model_config_name: str = None, - memory_config: Optional[dict] = None, - prompt_type: Optional[PromptType] = PromptType.LIST, - config: Optional[dict] = None, - ) -> None: - """ - Initialize the RAG LlamaIndexAgent - Args: - name (str): the name for the agent - sys_prompt (str): system prompt for the RAG agent - model_config_name (str): language model for the agent - emb_model_config_name (str): embedding model for the agent - memory_config (dict): memory configuration - prompt_type (PromptType): prompt type, list or str - config (dict): additional config for RAG and agent - Current support adjustable parameters includes: - "data_path": - path to data directory where data is stored, - "chunk_size": - maximum chunk size for preprocessed documents, - default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE - "chunk_overlap": - overlap between preprocessed chunks, default is - agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP - "similarity_top_k": - number of chunks in each retrieval, default - is agentscope.rag.rag.DEFAULT_TOP_K - "log_retrieval" (bool): - whether to user agent.speak() to print the - retrieved documents, default is False. - "recent_n_mem": - how many messages in memory is used as query - for retrieval if memory is used, default is 1 - """ - super().__init__( - name=name, - sys_prompt=sys_prompt, - model_config_name=model_config_name, - emb_model_config_name=emb_model_config_name, - memory_config=memory_config, - prompt_type=prompt_type, - config=config, - ) - try: - from llama_index.core import SimpleDirectoryReader - except ImportError as exc: - raise ImportError( - " LlamaIndexAgent requires llama-index to be install." - "Please run `pip install llama-index`", - ) from exc - # init rag related attributes - self.rag = LlamaIndexRAG( - model=self.model, - emb_model=self.emb_model, - config=config, - ) - # load the document to memory - # Feed the AgentScope tutorial documents, so that - # the agent can answer questions related to AgentScope! - if "data_path" not in self.config: - self.config[" data_path"] = "./data" - logger.warning( - "No data_path provided in RAG agent config," - "use default path `./data`", - ) - docs = self.rag.load_data( - loader=SimpleDirectoryReader(self.config["data_path"]), - ) - self.rag.store_and_index(docs) - - -class LangChainRAGAgent(RAGAgentBase): - """ - A LlamaIndex agent build on LlamaIndex. - """ - - def __init__( - self, - name: str, - sys_prompt: str, - model_config_name: str, - emb_model_config_name: str, - memory_config: Optional[dict] = None, - prompt_type: Optional[PromptType] = PromptType.LIST, - config: Optional[dict] = None, - ) -> None: - """ - Initialize the RAG LangChainRAGAgent - Args: - name (str): the name for the agent - sys_prompt (str): system prompt for the RAG agent - model_config_name (str): language model for the agent - emb_model_config_name (str): embedding model for the agent - memory_config (dict): memory configuration - prompt_type (PromptType): prompt type, list or str - config (dict): additional config for RAG and agent - Current support adjustable parameters includes: - "data_path": - path to data directory where data is stored, - "chunk_size": - maximum chunk size for preprocessed documents, - default is agentscope.rag.rag.DEFAULT_CHUNK_SIZE - "chunk_overlap": - overlap between preprocessed chunks, default is - agentscope.rag.rag.DEFAULT_CHUNK_OVERLAP - "similarity_top_k": - number of chunks in each retrieval, default - is agentscope.rag.rag.DEFAULT_TOP_K - "log_retrieval" (bool): - whether to user agent.speak() to print the - retrieved documents, default is False. - "recent_n_mem": - how many messages in memory is used as query - for retrieval if memory is used, default is 1 - """ - super().__init__( - name=name, - sys_prompt=sys_prompt, - model_config_name=model_config_name, - emb_model_config_name=emb_model_config_name, - memory_config=memory_config, - prompt_type=prompt_type, - config=config, - ) - try: - from langchain_community.document_loaders import DirectoryLoader - except ImportError as exc: - raise ImportError( - "LangChainRAGAgent requires LangChain related packages " - "installed. Please run `pip install langchain " - "unstructured[all-docs] langchain-text-splitters`", - ) from exc - - # init rag related attributes - self.rag = LangChainRAG( - model=self.model, - emb_model=self.emb_model, - config=config, - ) - # load the document to memory - # Feed the AgentScope tutorial documents, so that - # the agent can answer questions related to AgentScope! - if "data_path" not in self.config: - self.config[" data_path"] = "./data" - logger.warning( - "No data_path provided in RAG agent config," - "use default path `./data`", - ) - docs = self.rag.load_data( - loader=DirectoryLoader(self.config["data_path"]), - ) - self.rag.store_and_index(docs) diff --git a/examples/rag/rag_example.py b/examples/rag/rag_example.py index 600bfcfec..4a2c280ba 100644 --- a/examples/rag/rag_example.py +++ b/examples/rag/rag_example.py @@ -4,17 +4,15 @@ an agent with RAG capability. """ import os -import argparse -from loguru import logger import agentscope from agentscope.agents import UserAgent +from agentscope.message import Msg def main() -> None: - """A conversation demo""" - - agentscope.init( + """A RAG multi-agent demo""" + agents = agentscope.init( model_configs=[ { "model_type": "dashscope_chat", @@ -29,39 +27,11 @@ def main() -> None: "api_key": f"{os.environ.get('DASHSCOPE_API_KEY')}", }, ], + agent_configs="./agent_config.json", ) - if args.module == "llamaindex": - try: - from rag_agents import LlamaIndexAgent - except ImportError as exc: - raise ImportError( - "Please install llamaindex packages.", - ) from exc - AgentClass = LlamaIndexAgent - else: - try: - from rag_agents import LangChainRAGAgent - except ImportError as exc: - raise ImportError( - "Please install LangChain RAG packages.", - ) from exc - AgentClass = LangChainRAGAgent - rag_agent = AgentClass( - name="Assistant", - sys_prompt="You're a helpful assistant. You need to generate" - " answers based on the provided context.", - model_config_name="qwen_config", # your model config name - emb_model_config_name="qwen_emb_config", - config={ - "data_path": args.data_path, - "chunk_size": 2048, - "chunk_overlap": 40, - "similarity_top_k": 10, - "log_retrieval": True, - "recent_n_mem": 1, - }, - ) + tutorial_agent, code_explain_agent, summarize_agent = agents + user_agent = UserAgent() # start the conversation between user and assistant while True: @@ -69,27 +39,19 @@ def main() -> None: x.role = "user" # to enforce dashscope requirement on roles if len(x["content"]) == 0 or str(x["content"]).startswith("exit"): break - rag_agent(x) + tutorial_response = tutorial_agent(x) + code_explain = code_explain_agent(x) + msg = Msg( + name="user", + role="user", + content=tutorial_response["content"] + + "\n" + + code_explain["content"] + + "\n" + + x["content"], + ) + summarize_agent(msg) if __name__ == "__main__": - # The default parameters can set a AgentScope consultant to - # answer question about agentscope based on tutorial. - parser = argparse.ArgumentParser() - parser.add_argument( - "--module", - choices=["llamaindex", "langchain"], - default="llamaindex", - ) - parser.add_argument( - "--data_path", - type=str, - default="../../docs/sphinx_doc/en/source/tutorial", - ) - args = parser.parse_args() - if args.module == "langchain": - logger.warning( - "LangChain RAG Chosen. May require install pandoc in advanced.", - "For example, run ` brew install pandoc` on MacOS.", - ) main() diff --git a/src/agentscope/agents/__init__.py b/src/agentscope/agents/__init__.py index ee0fdbb33..ec7c15bab 100644 --- a/src/agentscope/agents/__init__.py +++ b/src/agentscope/agents/__init__.py @@ -7,6 +7,7 @@ from .user_agent import UserAgent from .text_to_image_agent import TextToImageAgent from .rpc_agent import RpcAgentServerLauncher +from .rag_agents import RAGAgentBase, LlamaIndexAgent __all__ = [ @@ -17,4 +18,6 @@ "TextToImageAgent", "UserAgent", "RpcAgentServerLauncher", + "RAGAgentBase", + "LlamaIndexAgent", ] diff --git a/src/agentscope/agents/rag_agents.py b/src/agentscope/agents/rag_agents.py new file mode 100644 index 000000000..9f70e50ff --- /dev/null +++ b/src/agentscope/agents/rag_agents.py @@ -0,0 +1,330 @@ +# -*- coding: utf-8 -*- +""" +This example shows how to build an agent with RAG +(with LlamaIndex or Langchain) +""" + +from abc import ABC, abstractmethod +from typing import Optional, Any +import importlib +from loguru import logger + + +from agentscope.agents.agent import AgentBase +from agentscope.message import Msg +from agentscope.prompt import PromptEngine +from agentscope.models import load_model_by_config_name +from agentscope.rag import RAGBase, LlamaIndexRAG + + +class RAGAgentBase(AgentBase, ABC): + """ + Base class for RAG agents + """ + + def __init__( + self, + name: str, + sys_prompt: str, + model_config_name: str, + emb_model_config_name: str, + memory_config: Optional[dict] = None, + rag_config: Optional[dict] = None, + ) -> None: + """ + Initialize the RAG base agent + Args: + name (str): + the name for the agent. + sys_prompt (str): + system prompt for the RAG agent. + model_config_name (str): + language model for the agent. + emb_model_config_name (str): + embedding model for the agent. + memory_config (dict): + memory configuration. + rag_config (dict): + config for RAG. It contains most of the + important parameters for RAG modules. If not provided, + the default setting will be used. + Examples can refer to children classes. + """ + super().__init__( + name=name, + sys_prompt=sys_prompt, + model_config_name=model_config_name, + use_memory=True, + memory_config=memory_config, + ) + # init prompt engine + self.engine = PromptEngine(self.model) + # setup embedding model used in RAG + self.emb_model = load_model_by_config_name(emb_model_config_name) + + self.rag_config = rag_config or {} + if "log_retrieval" not in self.rag_config: + self.rag_config["log_retrieval"] = True + + # use LlamaIndexAgent OR LangChainAgent + self.rag = self.init_rag() + + @abstractmethod + def init_rag(self) -> RAGBase: + """initialize RAG with configuration""" + + def _prepare_args_from_config( + self, + config: dict, + ) -> Any: + """ + Helper function to build args for the two functions: + rag.load_data(...) and rag.store_and_index(docs, ...) + in RAG classes. + Args: + config (dict): a dictionary containing configurations + + Returns: + Any: an object that is parsed/built to be an element + of input to the function of RAG module. + """ + if not isinstance(config, dict): + return config + + if "create_object" in config: + # if a term in args is a object, + # recursively create object with args from config + module_name = config.get("module", "") + class_name = config.get("class", "") + init_args = config.get("init_args", {}) + try: + cur_module = importlib.import_module(module_name) + cur_class = getattr(cur_module, class_name) + init_args = self._prepare_args_from_config(init_args) + logger.info( + f"load and build object{cur_module, cur_class, init_args}", + ) + return cur_class(**init_args) + except ImportError as exc: + logger.error( + f"Fail to load class {class_name} " + f"from module {module_name}", + ) + raise ImportError from exc + else: + prepared_args = {} + for key, value in config.items(): + if isinstance(value, list): + prepared_args[key] = [] + for c in value: + prepared_args[key].append( + self._prepare_args_from_config(c), + ) + elif isinstance(value, dict): + prepared_args[key] = self._prepare_args_from_config(value) + else: + prepared_args[key] = value + return prepared_args + + def reply( + self, + x: dict = None, + ) -> dict: + """ + Reply function of the RAG agent. + Processes the input data, + 1) use the input data to retrieve with RAG function; + 2) generates a prompt using the current memory and system + prompt; + 3) invokes the language model to produce a response. The + response is then formatted and added to the dialogue memory. + + Args: + x (`dict`, defaults to `None`): + A dictionary representing the user's input to the agent. This + input is added to the memory if provided. Defaults to + None. + Returns: + A dictionary representing the message generated by the agent in + response to the user's input. + """ + retrieved_docs_to_string = "" + # record the input if needed + if self.memory: + self.memory.add(x) + # in case no input is provided (e.g., in msghub), + # use the memory as query + history = self.engine.join( + self.memory.get_memory( + recent_n=self.rag_config.get("recent_n_mem", 1), + ), + ) + query = ( + "/n".join( + [msg["content"] for msg in history], + ) + if isinstance(history, list) + else str(history) + ) + elif x is not None: + query = x["content"] + else: + query = "" + + if len(query) > 0: + # when content has information, do retrieval + retrieved_docs = self.rag.retrieve(query, to_list_strs=True) + for content in retrieved_docs: + retrieved_docs_to_string += "\n>>>> " + content + + if self.rag_config["log_retrieval"]: + self.speak("[retrieved]:" + retrieved_docs_to_string) + + # prepare prompt + prompt = self.model.format( + Msg( + name="system", + role="system", + content=self.sys_prompt, + ), + # {"role": "system", "content": retrieved_docs_to_string}, + self.memory.get_memory( + recent_n=self.rag_config.get("recent_n_mem", 1), + ), + Msg( + name="user", + role="user", + content="Context: " + retrieved_docs_to_string, + ), + ) + + # call llm and generate response + response = self.model(prompt).text + msg = Msg(self.name, response) + + # Print/speak the message in this agent's voice + self.speak(msg) + + if self.memory: + # Record the message in memory + self.memory.add(msg) + + return msg + + +class LlamaIndexAgent(RAGAgentBase): + """ + A LlamaIndex agent build on LlamaIndex. + """ + + def __init__( + self, + name: str, + sys_prompt: str, + model_config_name: str, + emb_model_config_name: str = None, + memory_config: Optional[dict] = None, + rag_config: Optional[dict] = None, + ) -> None: + """ + Initialize the RAG LlamaIndexAgent + Args: + name (str): + the name for the agent + sys_prompt (str): + system prompt for the RAG agent + model_config_name (str): + language model for the agent + emb_model_config_name (str): + embedding model for the agent + memory_config (dict): + memory configuration + rag_config (dict): + config for RAG. It contains the parameters for + RAG modules functions: + rag.load_data(...) and rag.store_and_index(docs, ...) + If not provided, the default setting will be used. + An example of the config for retrieving code files + is as following: + "rag_config": { + "load_data": { + "loader": { + "create_object": true, + "module": "llama_index.core", + "class": "SimpleDirectoryReader", + "init_args": { + "input_dir": "../../src/agentscope/models", + "recursive": true + } + } + }, + "store_and_index": { + "transformations": [ + { + "create_object": true, + "module": "llama_index.core.node_parser", + "class": "CodeSplitter", + "init_args": { + "language": "python", + "chunk_lines": 100 + } + } + ] + }, + "chunk_size": 2048, + "chunk_overlap": 40, + "similarity_top_k": 10, + "log_retrieval": true, + "recent_n_mem": 1 + } + """ + super().__init__( + name=name, + sys_prompt=sys_prompt, + model_config_name=model_config_name, + emb_model_config_name=emb_model_config_name, + memory_config=memory_config, + rag_config=rag_config, + ) + + def init_rag(self) -> LlamaIndexRAG: + # dynamic loading loader + # init rag related attributes + rag = LlamaIndexRAG( + model=self.model, + emb_model=self.emb_model, + config=self.rag_config, + ) + # load the document to memory + # Feed the AgentScope tutorial documents, so that + # the agent can answer questions related to AgentScope! + if "load_data" in self.rag_config: + load_data_args = self._prepare_args_from_config( + self.rag_config["load_data"], + ) + else: + try: + from llama_index.core import SimpleDirectoryReader + except ImportError as exc: + raise ImportError( + " LlamaIndexAgent requires llama-index to be install." + "Please run `pip install llama-index`", + ) from exc + load_data_args = { + "loader": SimpleDirectoryReader(self.config["data_path"]), + } + logger.info(f"rag.load_data args: {load_data_args}") + docs = rag.load_data(**load_data_args) + + # store and indexing + if "store_and_index" in self.rag_config: + store_and_index_args = self._prepare_args_from_config( + self.rag_config["store_and_index"], + ) + else: + store_and_index_args = {} + + logger.info(f"store_and_index_args args: {store_and_index_args}") + rag.store_and_index(docs, **store_and_index_args) + + return rag From c41eeaf7acedd8bd767351d7e13e21d652ce9b10 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Thu, 28 Mar 2024 20:00:14 +0800 Subject: [PATCH 16/21] update to avoid test import error --- src/agentscope/agents/rag_agents.py | 19 ++++++++++++++----- src/agentscope/rag/llama_index_rag.py | 5 ++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/agentscope/agents/rag_agents.py b/src/agentscope/agents/rag_agents.py index 9f70e50ff..98629e0d2 100644 --- a/src/agentscope/agents/rag_agents.py +++ b/src/agentscope/agents/rag_agents.py @@ -14,7 +14,13 @@ from agentscope.message import Msg from agentscope.prompt import PromptEngine from agentscope.models import load_model_by_config_name -from agentscope.rag import RAGBase, LlamaIndexRAG + +try: + from agentscope.rag import RAGBase, LlamaIndexRAG +except ImportError as exc: + raise ImportError( + "Need to install llama index for RAG agents.", + ) from exc class RAGAgentBase(AgentBase, ABC): @@ -105,12 +111,15 @@ def _prepare_args_from_config( f"load and build object{cur_module, cur_class, init_args}", ) return cur_class(**init_args) - except ImportError as exc: + except ImportError as exc_inner: logger.error( f"Fail to load class {class_name} " f"from module {module_name}", ) - raise ImportError from exc + raise ImportError( + f"Fail to load class {class_name} " + f"from module {module_name}", + ) from exc_inner else: prepared_args = {} for key, value in config.items(): @@ -305,11 +314,11 @@ def init_rag(self) -> LlamaIndexRAG: else: try: from llama_index.core import SimpleDirectoryReader - except ImportError as exc: + except ImportError as exc_inner: raise ImportError( " LlamaIndexAgent requires llama-index to be install." "Please run `pip install llama-index`", - ) from exc + ) from exc_inner load_data_args = { "loader": SimpleDirectoryReader(self.config["data_path"]), } diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index 60b0e1838..da4f687a1 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -21,7 +21,6 @@ from llama_index.core.node_parser import SentenceSplitter from llama_index.core import ( VectorStoreIndex, - SimpleDirectoryReader, ) except ImportError: BaseReader, BaseRetriever, BaseEmbedding, Embedding = ( @@ -32,7 +31,7 @@ ) IngestionPipeline, BasePydanticVectorStore, VectorStore = None, None, None PrivateAttr, NodeParser, SentenceSplitter = None, None, None - VectorStoreIndex, SimpleDirectoryReader = None, None + VectorStoreIndex = None from agentscope.rag import RAGBase from agentscope.rag.rag import ( @@ -156,7 +155,7 @@ def __init__( def load_data( self, - loader: BaseReader = SimpleDirectoryReader("./data/"), + loader: BaseReader, query: Optional[str] = None, **kwargs: Any, ) -> Any: From ca3eaa136f82814fb20af17314ab988d68e4028e Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Thu, 28 Mar 2024 20:53:46 +0800 Subject: [PATCH 17/21] update again --- src/agentscope/agents/__init__.py | 6 +++++- src/agentscope/rag/llama_index_rag.py | 11 ++++------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/agentscope/agents/__init__.py b/src/agentscope/agents/__init__.py index 2fd67009a..4a62c153c 100644 --- a/src/agentscope/agents/__init__.py +++ b/src/agentscope/agents/__init__.py @@ -8,7 +8,11 @@ from .text_to_image_agent import TextToImageAgent from .rpc_agent import RpcAgentServerLauncher from .react_agent import ReActAgent -from .rag_agents import RAGAgentBase, LlamaIndexAgent + +try: + from .rag_agents import RAGAgentBase, LlamaIndexAgent +except Exception: + RAGAgentBase, LlamaIndexAgent = None, None # type: ignore # NOQA __all__ = [ diff --git a/src/agentscope/rag/llama_index_rag.py b/src/agentscope/rag/llama_index_rag.py index da4f687a1..c46cc61cb 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/src/agentscope/rag/llama_index_rag.py @@ -23,15 +23,12 @@ VectorStoreIndex, ) except ImportError: - BaseReader, BaseRetriever, BaseEmbedding, Embedding = ( - None, - None, - None, - None, - ) + BaseReader, BaseRetriever = None, None + BaseEmbedding, Embedding = None, None IngestionPipeline, BasePydanticVectorStore, VectorStore = None, None, None - PrivateAttr, NodeParser, SentenceSplitter = None, None, None + NodeParser, SentenceSplitter = None, None VectorStoreIndex = None + PrivateAttr = None from agentscope.rag import RAGBase from agentscope.rag.rag import ( From dce475a39a1d355128dfaa7b78a10112670ee124 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Fri, 29 Mar 2024 09:03:29 +0800 Subject: [PATCH 18/21] update docs --- README.md | 1 + README_ZH.md | 1 + .../en/source/agentscope.agents.rst | 9 +++++ .../zh_CN/source/agentscope.agents.rst | 9 +++++ .../conversation_with_RAG_agents/README.md | 40 +++++++++++++++++++ .../agent_config.json | 4 +- .../rag_example.py | 0 src/agentscope/rag/__init__.py | 13 +++++- 8 files changed, 73 insertions(+), 4 deletions(-) create mode 100644 examples/conversation_with_RAG_agents/README.md rename examples/{rag => conversation_with_RAG_agents}/agent_config.json (97%) rename examples/{rag => conversation_with_RAG_agents}/rag_example.py (100%) diff --git a/README.md b/README.md index dee94f5d7..2c68bb324 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,7 @@ the following libraries. - [Self-Organizing Conversation](./examples/conversation_self_organizing) - [Basic Conversation with LangChain library](./examples/conversation_with_langchain) - [Conversation with ReAct Agent](./examples/conversation_with_react_agent) + - [Conversation with RAG Agent](./examples/conversation_with_RAG_agents) - Game - [Gomoku](./examples/game_gomoku) diff --git a/README_ZH.md b/README_ZH.md index 125c684b4..53a542c25 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -95,6 +95,7 @@ AgentScope支持使用以下库快速部署本地模型服务。 - [智能体自组织的对话](./examples/conversation_self_organizing) - [兼容LangChain的基础对话](./examples/conversation_with_langchain) - [与ReAct智能体对话](./examples/conversation_with_react_agent) + - [与RAG智能体对话](./examples/conversation_with_RAG_agents) - 游戏 - [五子棋](./examples/game_gomoku) diff --git a/docs/sphinx_doc/en/source/agentscope.agents.rst b/docs/sphinx_doc/en/source/agentscope.agents.rst index a27e688ea..f2e077cbc 100644 --- a/docs/sphinx_doc/en/source/agentscope.agents.rst +++ b/docs/sphinx_doc/en/source/agentscope.agents.rst @@ -66,3 +66,12 @@ react_agent module :members: :undoc-members: :show-inheritance: + + +rag_agent module +------------------------------- + +.. automodule:: agentscope.agents.rag_agents + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst b/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst index a27e688ea..f2e077cbc 100644 --- a/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst +++ b/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst @@ -66,3 +66,12 @@ react_agent module :members: :undoc-members: :show-inheritance: + + +rag_agent module +------------------------------- + +.. automodule:: agentscope.agents.rag_agents + :members: + :undoc-members: + :show-inheritance: diff --git a/examples/conversation_with_RAG_agents/README.md b/examples/conversation_with_RAG_agents/README.md new file mode 100644 index 000000000..92f9e3360 --- /dev/null +++ b/examples/conversation_with_RAG_agents/README.md @@ -0,0 +1,40 @@ +# AgentScope Consultants: a Multi-Agent RAG Application + +* **What is this example about?** +With the provided implementation and configuration, +you will obtain three different agents who can help you answer different questions about AgentScope. + +* **What is this example for?** By this example, we want to show how the agent with retrieval augmented generation (RAG) +capability can be used to build easily. + +## Prerequisites +* **Cloning repo:** This example requires cloning the whole AgentScope repo to local. +* **Packages:** This example is built on the LlamaIndex package. Thus, some packages need to be installed before running the example. + ```bash + pip install llama-index tree_sitter tree-sitter-languages + ``` +* **Model APIs:** This example uses Dashscope APIs. Thus, we also need an API key for DashScope. + ```bash + export DASH_SCOPE_API='YOUR_API_KEY' + ``` + However, you are welcome to replace the Dashscope language and embedding models with other models you like. + +## Start AgentScope Consultants +* **Terminal:** The most simple way to execute the AgentScope Consultants is running in terminal. + ```bash + python ./rag_example.py + ``` + Setting `log_retrieval` to `false` in `agent_config.json` can hide the retrieved information and provide only answers of agents. + +* **AS studio:** If you want to have more organized, clean UI, you can also run with our `as_studio`. + ```bash + as_studio ./rag_example.py + ``` + +### Customize AgentScope Consultants to other consultants +After you run the example, you may notice that this example consists of three RAG agents: +* `AgentScope Tutorial Assistant`: responsible for answering questions based on AgentScope tutorials (markdown files). +* `AgentScope Framework Code Assistant`: responsible for answering questions based on AgentScope code base (python files). +* `Summarize Assistant`: responsible for summarize the questions from the above two agents. + +These agents can be configured to answering questions based on other GitHub repo, by simply modifying the `input_dir` fields in the `agent_config.json`. \ No newline at end of file diff --git a/examples/rag/agent_config.json b/examples/conversation_with_RAG_agents/agent_config.json similarity index 97% rename from examples/rag/agent_config.json rename to examples/conversation_with_RAG_agents/agent_config.json index 742cf5a67..fc0a23c12 100644 --- a/examples/rag/agent_config.json +++ b/examples/conversation_with_RAG_agents/agent_config.json @@ -21,7 +21,7 @@ "chunk_size": 2048, "chunk_overlap": 40, "similarity_top_k": 10, - "log_retrieval": true, + "log_retrieval": false, "recent_n_mem": 1 } } @@ -62,7 +62,7 @@ "chunk_size": 2048, "chunk_overlap": 40, "similarity_top_k": 10, - "log_retrieval": true, + "log_retrieval": false, "recent_n_mem": 1 } } diff --git a/examples/rag/rag_example.py b/examples/conversation_with_RAG_agents/rag_example.py similarity index 100% rename from examples/rag/rag_example.py rename to examples/conversation_with_RAG_agents/rag_example.py diff --git a/src/agentscope/rag/__init__.py b/src/agentscope/rag/__init__.py index 2a1d3de61..b94975be1 100644 --- a/src/agentscope/rag/__init__.py +++ b/src/agentscope/rag/__init__.py @@ -1,8 +1,17 @@ # -*- coding: utf-8 -*- """ Import all pipeline related modules in the package. """ from .rag import RAGBase -from .llama_index_rag import LlamaIndexRAG -from .langchain_rag import LangChainRAG + +try: + from .llama_index_rag import LlamaIndexRAG +except Exception: + LlamaIndexRAG = None # type: ignore # NOQA + +try: + from .langchain_rag import LangChainRAG +except Exception: + LangChainRAG = None # type: ignore # NOQA + __all__ = [ "RAGBase", From e669c39067ff3edd4859e78132a9d5a71eda49c7 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Fri, 29 Mar 2024 15:23:33 +0800 Subject: [PATCH 19/21] update docs and move rag agents to example --- .../en/source/agentscope.agents.rst | 8 ------- .../zh_CN/source/agentscope.agents.rst | 9 ------- .../conversation_with_RAG_agents/README.md | 21 ++++++++++++++-- .../rag_agents.py | 24 +++++++------------ .../rag_example.py | 13 +++++++--- src/agentscope/agents/__init__.py | 7 ------ 6 files changed, 38 insertions(+), 44 deletions(-) rename {src/agentscope/agents => examples/conversation_with_RAG_agents}/rag_agents.py (95%) diff --git a/docs/sphinx_doc/en/source/agentscope.agents.rst b/docs/sphinx_doc/en/source/agentscope.agents.rst index f2e077cbc..dc77765c0 100644 --- a/docs/sphinx_doc/en/source/agentscope.agents.rst +++ b/docs/sphinx_doc/en/source/agentscope.agents.rst @@ -67,11 +67,3 @@ react_agent module :undoc-members: :show-inheritance: - -rag_agent module -------------------------------- - -.. automodule:: agentscope.agents.rag_agents - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst b/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst index f2e077cbc..a27e688ea 100644 --- a/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst +++ b/docs/sphinx_doc/zh_CN/source/agentscope.agents.rst @@ -66,12 +66,3 @@ react_agent module :members: :undoc-members: :show-inheritance: - - -rag_agent module -------------------------------- - -.. automodule:: agentscope.agents.rag_agents - :members: - :undoc-members: - :show-inheritance: diff --git a/examples/conversation_with_RAG_agents/README.md b/examples/conversation_with_RAG_agents/README.md index 92f9e3360..5b08379a4 100644 --- a/examples/conversation_with_RAG_agents/README.md +++ b/examples/conversation_with_RAG_agents/README.md @@ -7,6 +7,8 @@ you will obtain three different agents who can help you answer different questio * **What is this example for?** By this example, we want to show how the agent with retrieval augmented generation (RAG) capability can be used to build easily. +**Notice:** This example is a Beta version of the AgentScope RAG agent. A formal version will soon be added to `src/agentscope/agents`, but it may be subject to changes. + ## Prerequisites * **Cloning repo:** This example requires cloning the whole AgentScope repo to local. * **Packages:** This example is built on the LlamaIndex package. Thus, some packages need to be installed before running the example. @@ -17,7 +19,9 @@ capability can be used to build easily. ```bash export DASH_SCOPE_API='YOUR_API_KEY' ``` - However, you are welcome to replace the Dashscope language and embedding models with other models you like. + +**Note:** This example has been tested with `dashscope_chat` and `dashscope_text_embedding` model wrapper, with `qwen-max` and `text-embedding-v2` models. +However, you are welcome to replace the Dashscope language and embedding model wrappers or models with other models you like to test. ## Start AgentScope Consultants * **Terminal:** The most simple way to execute the AgentScope Consultants is running in terminal. @@ -37,4 +41,17 @@ After you run the example, you may notice that this example consists of three RA * `AgentScope Framework Code Assistant`: responsible for answering questions based on AgentScope code base (python files). * `Summarize Assistant`: responsible for summarize the questions from the above two agents. -These agents can be configured to answering questions based on other GitHub repo, by simply modifying the `input_dir` fields in the `agent_config.json`. \ No newline at end of file +These agents can be configured to answering questions based on other GitHub repo, by simply modifying the `input_dir` fields in the `agent_config.json`. + +For more advanced customization, we may need to learn a little bit from the following. + +**RAG modules:** In AgentScope, RAG modules are abstract to provide three basic functions: `load_data`, `store_and_index` and `retrieve`. Refer to `src/agentscope/rag` for more details. + +**RAG configs:** In the example configuration (the `rag_config` field), all parameters are optional. But if you want to customize them, you may want to learn the following: +* `load_data`: contains all parameters for the the `rag.load_data` function. +Since the `load_data` accepts a dataloader object `loader`, the `loader` in the config need to have `"create_object": true` to let a internal parse create a LlamaIndex data loader object. +The loader object is an instance of `class` in module `module`, with initialization parameters in `init_args`. + +* `store_and_index`: contains all parameters for the the `rag.store_and_index` function. +For example, you can pass `vector_store` and `retriever` configurations in a similar way as the `loader` mentioned above. +For the `transformations` parameter, you can pass a list of dicts, each of which corresponds to building a `NodeParser`-kind of preprocessor in Llamaindex. \ No newline at end of file diff --git a/src/agentscope/agents/rag_agents.py b/examples/conversation_with_RAG_agents/rag_agents.py similarity index 95% rename from src/agentscope/agents/rag_agents.py rename to examples/conversation_with_RAG_agents/rag_agents.py index 98629e0d2..c2a5406e8 100644 --- a/src/agentscope/agents/rag_agents.py +++ b/examples/conversation_with_RAG_agents/rag_agents.py @@ -1,7 +1,9 @@ # -*- coding: utf-8 -*- """ This example shows how to build an agent with RAG -(with LlamaIndex or Langchain) +with LlamaIndex. + +Notice, this is a Beta version of RAG agent. """ from abc import ABC, abstractmethod @@ -12,15 +14,9 @@ from agentscope.agents.agent import AgentBase from agentscope.message import Msg -from agentscope.prompt import PromptEngine from agentscope.models import load_model_by_config_name -try: - from agentscope.rag import RAGBase, LlamaIndexRAG -except ImportError as exc: - raise ImportError( - "Need to install llama index for RAG agents.", - ) from exc +from agentscope.rag import RAGBase, LlamaIndexRAG class RAGAgentBase(AgentBase, ABC): @@ -63,8 +59,6 @@ def __init__( use_memory=True, memory_config=memory_config, ) - # init prompt engine - self.engine = PromptEngine(self.model) # setup embedding model used in RAG self.emb_model = load_model_by_config_name(emb_model_config_name) @@ -163,10 +157,8 @@ def reply( self.memory.add(x) # in case no input is provided (e.g., in msghub), # use the memory as query - history = self.engine.join( - self.memory.get_memory( - recent_n=self.rag_config.get("recent_n_mem", 1), - ), + history = self.memory.get_memory( + recent_n=self.rag_config.get("recent_n_mem", 1), ) query = ( "/n".join( @@ -255,6 +247,7 @@ def __init__( If not provided, the default setting will be used. An example of the config for retrieving code files is as following: + "rag_config": { "load_data": { "loader": { @@ -262,8 +255,9 @@ def __init__( "module": "llama_index.core", "class": "SimpleDirectoryReader", "init_args": { - "input_dir": "../../src/agentscope/models", + "input_dir": "path/to/data", "recursive": true + ... } } }, diff --git a/examples/conversation_with_RAG_agents/rag_example.py b/examples/conversation_with_RAG_agents/rag_example.py index 4a2c280ba..1e2b4c6d3 100644 --- a/examples/conversation_with_RAG_agents/rag_example.py +++ b/examples/conversation_with_RAG_agents/rag_example.py @@ -3,16 +3,20 @@ A simple example for conversation between user and an agent with RAG capability. """ +import json import os +from rag_agents import LlamaIndexAgent + import agentscope from agentscope.agents import UserAgent from agentscope.message import Msg +from agentscope.agents import DialogAgent def main() -> None: """A RAG multi-agent demo""" - agents = agentscope.init( + agentscope.init( model_configs=[ { "model_type": "dashscope_chat", @@ -27,10 +31,13 @@ def main() -> None: "api_key": f"{os.environ.get('DASHSCOPE_API_KEY')}", }, ], - agent_configs="./agent_config.json", ) - tutorial_agent, code_explain_agent, summarize_agent = agents + with open("./agent_config.json", "r", encoding="utf-8") as f: + agent_configs = json.load(f) + tutorial_agent = LlamaIndexAgent(**agent_configs[0]["args"]) + code_explain_agent = LlamaIndexAgent(**agent_configs[1]["args"]) + summarize_agent = DialogAgent(**agent_configs[2]["args"]) user_agent = UserAgent() # start the conversation between user and assistant diff --git a/src/agentscope/agents/__init__.py b/src/agentscope/agents/__init__.py index 4a62c153c..92ead2676 100644 --- a/src/agentscope/agents/__init__.py +++ b/src/agentscope/agents/__init__.py @@ -9,11 +9,6 @@ from .rpc_agent import RpcAgentServerLauncher from .react_agent import ReActAgent -try: - from .rag_agents import RAGAgentBase, LlamaIndexAgent -except Exception: - RAGAgentBase, LlamaIndexAgent = None, None # type: ignore # NOQA - __all__ = [ "AgentBase", @@ -24,6 +19,4 @@ "UserAgent", "RpcAgentServerLauncher", "ReActAgent", - "RAGAgentBase", - "LlamaIndexAgent", ] From 0f95847c04a5deeb9620badfadf0170b92a02fda Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Fri, 29 Mar 2024 16:23:07 +0800 Subject: [PATCH 20/21] move rag module to example --- .../conversation_with_RAG_agents}/rag/__init__.py | 6 ++---- .../conversation_with_RAG_agents}/rag/langchain_rag.py | 8 ++++++-- .../conversation_with_RAG_agents}/rag/llama_index_rag.py | 6 +++--- .../conversation_with_RAG_agents}/rag/rag.py | 0 examples/conversation_with_RAG_agents/rag_agents.py | 3 +-- 5 files changed, 12 insertions(+), 11 deletions(-) rename {src/agentscope => examples/conversation_with_RAG_agents}/rag/__init__.py (71%) rename {src/agentscope => examples/conversation_with_RAG_agents}/rag/langchain_rag.py (96%) rename {src/agentscope => examples/conversation_with_RAG_agents}/rag/llama_index_rag.py (98%) rename {src/agentscope => examples/conversation_with_RAG_agents}/rag/rag.py (100%) diff --git a/src/agentscope/rag/__init__.py b/examples/conversation_with_RAG_agents/rag/__init__.py similarity index 71% rename from src/agentscope/rag/__init__.py rename to examples/conversation_with_RAG_agents/rag/__init__.py index b94975be1..3c8f48882 100644 --- a/src/agentscope/rag/__init__.py +++ b/examples/conversation_with_RAG_agents/rag/__init__.py @@ -2,10 +2,8 @@ """ Import all pipeline related modules in the package. """ from .rag import RAGBase -try: - from .llama_index_rag import LlamaIndexRAG -except Exception: - LlamaIndexRAG = None # type: ignore # NOQA +from .llama_index_rag import LlamaIndexRAG + try: from .langchain_rag import LangChainRAG diff --git a/src/agentscope/rag/langchain_rag.py b/examples/conversation_with_RAG_agents/rag/langchain_rag.py similarity index 96% rename from src/agentscope/rag/langchain_rag.py rename to examples/conversation_with_RAG_agents/rag/langchain_rag.py index d9b176691..36a329547 100644 --- a/src/agentscope/rag/langchain_rag.py +++ b/examples/conversation_with_RAG_agents/rag/langchain_rag.py @@ -23,8 +23,11 @@ TextSplitter = None CharacterTextSplitter = None -from agentscope.rag import RAGBase -from agentscope.rag.rag import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE +from examples.conversation_with_RAG_agents.rag import RAGBase +from examples.conversation_with_RAG_agents.rag.rag import ( + DEFAULT_CHUNK_OVERLAP, + DEFAULT_CHUNK_SIZE, +) from agentscope.models import ModelWrapperBase @@ -135,6 +138,7 @@ def store_and_index( splitter: Optional[TextSplitter] = None, **kwargs: Any, ) -> Any: + # pylint: disable=unused-argument """ Preprocessing the loaded documents. Args: diff --git a/src/agentscope/rag/llama_index_rag.py b/examples/conversation_with_RAG_agents/rag/llama_index_rag.py similarity index 98% rename from src/agentscope/rag/llama_index_rag.py rename to examples/conversation_with_RAG_agents/rag/llama_index_rag.py index c46cc61cb..8756856ff 100644 --- a/src/agentscope/rag/llama_index_rag.py +++ b/examples/conversation_with_RAG_agents/rag/llama_index_rag.py @@ -30,8 +30,8 @@ VectorStoreIndex = None PrivateAttr = None -from agentscope.rag import RAGBase -from agentscope.rag.rag import ( +from rag import RAGBase +from rag.rag import ( DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, DEFAULT_TOP_K, @@ -136,7 +136,7 @@ def __init__( super().__init__(model, emb_model, config, **kwargs) self.retriever = None self.index = None - self.persist_dir = kwargs.get("persist_dir", "./") + self.persist_dir = kwargs.get("persist_dir", "/") self.emb_model = emb_model print(self.config) diff --git a/src/agentscope/rag/rag.py b/examples/conversation_with_RAG_agents/rag/rag.py similarity index 100% rename from src/agentscope/rag/rag.py rename to examples/conversation_with_RAG_agents/rag/rag.py diff --git a/examples/conversation_with_RAG_agents/rag_agents.py b/examples/conversation_with_RAG_agents/rag_agents.py index c2a5406e8..101b2e305 100644 --- a/examples/conversation_with_RAG_agents/rag_agents.py +++ b/examples/conversation_with_RAG_agents/rag_agents.py @@ -11,13 +11,12 @@ import importlib from loguru import logger +from rag import RAGBase, LlamaIndexRAG from agentscope.agents.agent import AgentBase from agentscope.message import Msg from agentscope.models import load_model_by_config_name -from agentscope.rag import RAGBase, LlamaIndexRAG - class RAGAgentBase(AgentBase, ABC): """ From d8417d57f585a94e584bb1b68017e3080b9a4f4a Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Fri, 29 Mar 2024 16:25:56 +0800 Subject: [PATCH 21/21] remove one empty line --- docs/sphinx_doc/en/source/agentscope.agents.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/sphinx_doc/en/source/agentscope.agents.rst b/docs/sphinx_doc/en/source/agentscope.agents.rst index dc77765c0..a27e688ea 100644 --- a/docs/sphinx_doc/en/source/agentscope.agents.rst +++ b/docs/sphinx_doc/en/source/agentscope.agents.rst @@ -66,4 +66,3 @@ react_agent module :members: :undoc-members: :show-inheritance: -