From 8fc63241e32ef627608ee1b17bfb76f09584677a Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 12 Jun 2024 20:20:24 -0700 Subject: [PATCH 1/5] add llamaindexknowledge test --- tests/knowledge_test.py | 90 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 tests/knowledge_test.py diff --git a/tests/knowledge_test.py b/tests/knowledge_test.py new file mode 100644 index 000000000..c8ae2914a --- /dev/null +++ b/tests/knowledge_test.py @@ -0,0 +1,90 @@ +# -*- coding: utf-8 -*- +""" +Unit tests for knowledge (RAG module in AgentScope) +""" + +import os +import unittest +from typing import Any +import shutil + +from agentscope.rag import LlamaIndexKnowledge +from agentscope.models import OpenAIEmbeddingWrapper, ModelResponse + + +class DummyModel(OpenAIEmbeddingWrapper): + """ + Dummy model wrapper for testing + """ + + def __init__(self) -> None: + """dummy init""" + + def __call__(self, *args: Any, **kwargs: Any) -> ModelResponse: + """dummy call""" + return ModelResponse(embedding=[[1.0, 2.0]]) + + +class KnowledgeTest(unittest.TestCase): + """ + Test cases for TemporaryMemory + """ + + def setUp(self) -> None: + """set up test data""" + self.data_dir = "tmp_data_dir" + if not os.path.exists(self.data_dir): + os.mkdir(self.data_dir) + self.file_name_1 = "tmp_data_dir/file1.txt" + self.content = "testing file" + with open(self.file_name_1, "w", encoding="utf-8") as f: + f.write(self.content) + + def tearDown(self) -> None: + """Clean up before & after tests.""" + if os.path.exists(self.data_dir): + shutil.rmtree(self.data_dir) + if os.path.exists("./runs"): + shutil.rmtree("./runs") + + def test_llamaindexknowledge(self) -> None: + """test llamaindexknowledge""" + dummy_model = DummyModel() + + knowledge_config = { + "knowledge_id": "", + "data_processing": [], + } + loader_config = { + "load_data": { + "loader": { + "create_object": True, + "module": "llama_index.core", + "class": "SimpleDirectoryReader", + "init_args": {}, + }, + }, + } + loader_init = {"input_dir": self.data_dir, "required_exts": ".txt"} + + loader_config["load_data"]["loader"]["init_args"] = loader_init + knowledge_config["data_processing"].append(loader_config) + + knowledge = LlamaIndexKnowledge( + knowledge_id="test_knowledge", + emb_model=dummy_model, + knowledge_config=knowledge_config, + ) + retrieved = knowledge.retrieve( + query="testing", + similarity_top_k=2, + to_list_strs=True, + ) + self.assertEqual( + retrieved, + [self.content], + ) + + +if __name__ == "__main__": + unittest.main() From df3d0a5cfb0a617a212655e22e4c1d5ac1f9f462 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 12 Jun 2024 20:33:29 -0700 Subject: [PATCH 2/5] prevent PermissionError --- tests/knowledge_test.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/knowledge_test.py b/tests/knowledge_test.py index c8ae2914a..4095f1a28 100644 --- a/tests/knowledge_test.py +++ b/tests/knowledge_test.py @@ -42,10 +42,13 @@ def setUp(self) -> None: def tearDown(self) -> None: """Clean up before & after tests.""" - if os.path.exists(self.data_dir): - shutil.rmtree(self.data_dir) - if os.path.exists("./runs"): - shutil.rmtree("./runs") + try: + if os.path.exists(self.data_dir): + shutil.rmtree(self.data_dir) + if os.path.exists("./runs"): + shutil.rmtree("./runs") + except Exception: + pass def test_llamaindexknowledge(self) -> None: """test llamaindexknowledge""" From 14fc079d74393f76263f94c1f4c1e1b354211cd3 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 12 Jun 2024 20:44:39 -0700 Subject: [PATCH 3/5] fix version of llamaindex --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index e5ea88322..d858cde17 100644 --- a/setup.py +++ b/setup.py @@ -50,7 +50,7 @@ ] rag_requires = [ - "llama-index", + "llama-index==0.10.30", ] studio_requires = [] From bd2e9bdb6cdc66dff1b4e6caefac93a8c9fb07aa Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 12 Jun 2024 21:44:39 -0700 Subject: [PATCH 4/5] remove ImportErrorReporter in .rag module --- src/agentscope/rag/llama_index_knowledge.py | 33 +++++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/src/agentscope/rag/llama_index_knowledge.py b/src/agentscope/rag/llama_index_knowledge.py index 4f4195a96..f3dda27d6 100644 --- a/src/agentscope/rag/llama_index_knowledge.py +++ b/src/agentscope/rag/llama_index_knowledge.py @@ -9,6 +9,7 @@ from loguru import logger try: + import llama_index from llama_index.core.base.base_retriever import BaseRetriever from llama_index.core.base.embeddings.base import ( BaseEmbedding, @@ -28,19 +29,18 @@ TransformComponent, ) except ImportError as import_error: - from agentscope.utils.tools import ImportErrorReporter - - BaseRetriever = ImportErrorReporter(import_error, "full") - BaseEmbedding = ImportErrorReporter(import_error, "full") - Embedding = ImportErrorReporter(import_error, "full") - IngestionPipeline = ImportErrorReporter(import_error, "full") - SentenceSplitter = ImportErrorReporter(import_error, "full") - VectorStoreIndex = ImportErrorReporter(import_error, "full") - StorageContext = ImportErrorReporter(import_error, "full") - load_index_from_storage = ImportErrorReporter(import_error, "full") - PrivateAttr = ImportErrorReporter(import_error, "full") - Document = ImportErrorReporter(import_error, "full") - TransformComponent = ImportErrorReporter(import_error, "full") + llama_index = None + BaseRetriever = None + BaseEmbedding = None + Embedding = None + IngestionPipeline = None + SentenceSplitter = None + VectorStoreIndex = None + StorageContext = None + load_index_from_storage = None + PrivateAttr = None + Document = None + TransformComponent = None from agentscope.file_manager import file_manager from agentscope.models import ModelWrapperBase @@ -195,6 +195,13 @@ def __init__( model=model, **kwargs, ) + if llama_index is None: + raise ImportError( + "LlamaIndexKnowledge require llama-index installed. " + "Try a stable llama-index version, such as " + "`pip install llama-index==0.10.30`" + ) + if persist_root is None: persist_root = file_manager.dir self.persist_dir = os.path.join(persist_root, knowledge_id) From 9344f6a855ebb6a9668d805c5cf3908802908d41 Mon Sep 17 00:00:00 2001 From: ZiTao-Li Date: Wed, 12 Jun 2024 21:47:40 -0700 Subject: [PATCH 5/5] format --- src/agentscope/rag/llama_index_knowledge.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/agentscope/rag/llama_index_knowledge.py b/src/agentscope/rag/llama_index_knowledge.py index f3dda27d6..af89fb94d 100644 --- a/src/agentscope/rag/llama_index_knowledge.py +++ b/src/agentscope/rag/llama_index_knowledge.py @@ -28,7 +28,7 @@ Document, TransformComponent, ) -except ImportError as import_error: +except ImportError: llama_index = None BaseRetriever = None BaseEmbedding = None @@ -199,7 +199,7 @@ def __init__( raise ImportError( "LlamaIndexKnowledge require llama-index installed. " "Try a stable llama-index version, such as " - "`pip install llama-index==0.10.30`" + "`pip install llama-index==0.10.30`", ) if persist_root is None: