From 716cb1a0c5434d4bd1ce20bffc5d7cbc80919039 Mon Sep 17 00:00:00 2001 From: betterwang Date: Thu, 7 Mar 2024 22:07:04 +0800 Subject: [PATCH] memory_storage use rag_pipeline --- metagpt/memory/longterm_memory.py | 5 +++- metagpt/memory/memory_storage.py | 6 +++- metagpt/rag/schema.py | 1 + .../document_store/test_faiss_store.py | 25 +++++++++++----- tests/metagpt/memory/mock_text_embed.py | 30 ++++++++++++------- tests/metagpt/memory/test_longterm_memory.py | 9 ++++-- tests/metagpt/memory/test_memory_storage.py | 10 +++++-- 7 files changed, 63 insertions(+), 23 deletions(-) diff --git a/metagpt/memory/longterm_memory.py b/metagpt/memory/longterm_memory.py index e90413085..27a737e6c 100644 --- a/metagpt/memory/longterm_memory.py +++ b/metagpt/memory/longterm_memory.py @@ -32,7 +32,7 @@ class LongTermMemory(Memory): self.memory_storage.recover_memory(role_id) self.rc = rc if not self.memory_storage.is_initialized: - logger.warning(f"It may the first time to run Agent {role_id}, the long-term memory is empty") + logger.warning(f"It may the first time to run Role {role_id}, the long-term memory is empty") else: logger.warning(f"Role {role_id} has existing memory storage and has recovered them.") self.msg_from_recover = True @@ -66,6 +66,9 @@ class LongTermMemory(Memory): ltm_news.append(mem) return ltm_news[-k:] + def persit(self): + self.memory_storage.persit() + def delete(self, message: Message): super().delete(message) # TODO delete message in memory_storage diff --git a/metagpt/memory/memory_storage.py b/metagpt/memory/memory_storage.py index b7d49e1c3..706e75c5a 100644 --- a/metagpt/memory/memory_storage.py +++ b/metagpt/memory/memory_storage.py @@ -43,7 +43,7 @@ class MemoryStorage(object): if self.role_mem_path.joinpath("default__vector_store.json").exists(): self.faiss_engine = SimpleEngine.from_index( - index_config=[FAISSIndexConfig(persist_path=self.cache_dir)], + index_config=FAISSIndexConfig(persist_path=self.cache_dir), retriever_configs=[FAISSRetrieverConfig()], embed_model=self.embedding, ) @@ -73,3 +73,7 @@ class MemoryStorage(object): def clean(self): shutil.rmtree(self.cache_dir, ignore_errors=True) self._initialized = False + + def persit(self): + if self.faiss_engine: + self.faiss_engine.index.storage_context.persist(self.cache_dir) diff --git a/metagpt/rag/schema.py b/metagpt/rag/schema.py index 9657ae846..8f5828233 100644 --- a/metagpt/rag/schema.py +++ b/metagpt/rag/schema.py @@ -104,6 +104,7 @@ class ObjectNode(TextNode): def __init__(self, **kwargs): super().__init__(**kwargs) self.excluded_llm_metadata_keys = list(ObjectNodeMetadata.model_fields.keys()) + self.excluded_embed_metadata_keys = self.excluded_llm_metadata_keys @staticmethod def get_obj_metadata(obj: RAGObject) -> dict: diff --git a/tests/metagpt/document_store/test_faiss_store.py b/tests/metagpt/document_store/test_faiss_store.py index 7c712294e..f5a479d35 100644 --- a/tests/metagpt/document_store/test_faiss_store.py +++ b/tests/metagpt/document_store/test_faiss_store.py @@ -6,7 +6,10 @@ @File : test_faiss_store.py """ +<<<<<<< HEAD from typing import Optional +======= +>>>>>>> f14fee9b (memory_storage use rag_pipeline) import numpy as np import pytest @@ -17,16 +20,22 @@ from metagpt.logs import logger from metagpt.roles import Sales -def mock_openai_embed_documents(self, texts: list[str], chunk_size: Optional[int] = 0) -> list[list[float]]: +def mock_openai_embed_documents(self, texts: list[str], show_progress: bool = False) -> list[list[float]]: num = len(texts) embeds = np.random.randint(1, 100, size=(num, 1536)) # 1536: openai embedding dim - embeds = (embeds - embeds.mean(axis=0)) / (embeds.std(axis=0)) - return embeds + embeds = (embeds - embeds.mean(axis=0)) / embeds.std(axis=0) + return embeds.tolist() + + +def mock_openai_embed_document(self, text: str) -> list[float]: + embeds = mock_openai_embed_documents(self, [text]) + return embeds[0] @pytest.mark.asyncio async def test_search_json(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.json") role = Sales(profile="Sales", store=store) @@ -37,9 +46,10 @@ async def test_search_json(mocker): @pytest.mark.asyncio async def test_search_xlsx(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) - store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.xlsx") + store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.xlsx", meta_col="Answer", content_col="Question") role = Sales(profile="Sales", store=store) query = "Which facial cleanser is good for oily skin?" result = await role.run(query) @@ -48,7 +58,8 @@ async def test_search_xlsx(mocker): @pytest.mark.asyncio async def test_write(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) store = FaissStore(EXAMPLE_PATH / "data/search_kb/example.xlsx", meta_col="Answer", content_col="Question") _faiss_store = store.write() diff --git a/tests/metagpt/memory/mock_text_embed.py b/tests/metagpt/memory/mock_text_embed.py index 897c7cf10..af5f960ae 100644 --- a/tests/metagpt/memory/mock_text_embed.py +++ b/tests/metagpt/memory/mock_text_embed.py @@ -2,32 +2,42 @@ # -*- coding: utf-8 -*- # @Desc : -from typing import Optional - import numpy as np dim = 1536 # openai embedding dim +embed_zeros_arrr = np.zeros(shape=[1, dim]).tolist() +embed_ones_arrr = np.ones(shape=[1, dim]).tolist() text_embed_arr = [ - {"text": "Write a cli snake game", "embed": np.zeros(shape=[1, dim])}, # mock data, same as below - {"text": "Write a game of cli snake", "embed": np.zeros(shape=[1, dim])}, - {"text": "Write a 2048 web game", "embed": np.ones(shape=[1, dim])}, - {"text": "Write a Battle City", "embed": np.ones(shape=[1, dim])}, + {"text": "Write a cli snake game", "embed": embed_zeros_arrr}, # mock data, same as below + {"text": "Write a game of cli snake", "embed": embed_zeros_arrr}, + {"text": "Write a 2048 web game", "embed": embed_ones_arrr}, + {"text": "Write a Battle City", "embed": embed_ones_arrr}, { "text": "The user has requested the creation of a command-line interface (CLI) snake game", - "embed": np.zeros(shape=[1, dim]), + "embed": embed_zeros_arrr, }, - {"text": "The request is command-line interface (CLI) snake game", "embed": np.zeros(shape=[1, dim])}, + {"text": "The request is command-line interface (CLI) snake game", "embed": embed_zeros_arrr}, { "text": "Incorporate basic features of a snake game such as scoring and increasing difficulty", - "embed": np.ones(shape=[1, dim]), + "embed": embed_ones_arrr, }, ] text_idx_dict = {item["text"]: idx for idx, item in enumerate(text_embed_arr)} -def mock_openai_embed_documents(self, texts: list[str], chunk_size: Optional[int] = 0) -> list[list[float]]: + +def mock_openai_embed_documents(self, texts: list[str], show_progress: bool = False) -> list[list[float]]: idx = text_idx_dict.get(texts[0]) embed = text_embed_arr[idx].get("embed") return embed + + +def mock_openai_embed_document(self, text: str) -> list[float]: + embeds = mock_openai_embed_documents(self, [text]) + return embeds[0] + + +async def mock_openai_aembed_document(self, text: str) -> list[float]: + return mock_openai_embed_document(self, text) diff --git a/tests/metagpt/memory/test_longterm_memory.py b/tests/metagpt/memory/test_longterm_memory.py index d9eb5e67f..398b48c5d 100644 --- a/tests/metagpt/memory/test_longterm_memory.py +++ b/tests/metagpt/memory/test_longterm_memory.py @@ -13,13 +13,17 @@ from metagpt.roles.role import RoleContext from metagpt.schema import Message from tests.metagpt.memory.mock_text_embed import ( mock_openai_embed_documents, + mock_openai_embed_document, + mock_openai_aembed_document, text_embed_arr, ) @pytest.mark.asyncio async def test_ltm_search(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._aget_query_embedding", mock_openai_aembed_document) role_id = "UTUserLtm(Product Manager)" from metagpt.environment import Environment @@ -33,7 +37,7 @@ async def test_ltm_search(mocker): idea = text_embed_arr[0].get("text", "Write a cli snake game") message = Message(role="User", content=idea, cause_by=UserRequirement) news = await ltm.find_news([message]) - assert len(news) == 1 + assert len(news) == 0 ltm.add(message) sim_idea = text_embed_arr[1].get("text", "Write a game of cli snake") @@ -48,6 +52,7 @@ async def test_ltm_search(mocker): news = await ltm.find_news([new_message]) assert len(news) == 1 ltm.add(new_message) + ltm.persit() # restore from local index ltm_new = LongTermMemory() diff --git a/tests/metagpt/memory/test_memory_storage.py b/tests/metagpt/memory/test_memory_storage.py index 35f2309c5..eb96120d0 100644 --- a/tests/metagpt/memory/test_memory_storage.py +++ b/tests/metagpt/memory/test_memory_storage.py @@ -17,13 +17,17 @@ from metagpt.memory.memory_storage import MemoryStorage from metagpt.schema import Message from tests.metagpt.memory.mock_text_embed import ( mock_openai_embed_documents, + mock_openai_embed_document, + mock_openai_aembed_document, text_embed_arr, ) @pytest.mark.asyncio async def test_idea_message(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._aget_query_embedding", mock_openai_aembed_document) idea = text_embed_arr[0].get("text", "Write a cli snake game") role_id = "UTUser1(Product Manager)" @@ -53,7 +57,9 @@ async def test_idea_message(mocker): @pytest.mark.asyncio async def test_actionout_message(mocker): - mocker.patch("langchain_community.embeddings.openai.OpenAIEmbeddings.embed_documents", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embeddings", mock_openai_embed_documents) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._get_text_embedding", mock_openai_embed_document) + mocker.patch("llama_index.embeddings.openai.base.OpenAIEmbedding._aget_query_embedding", mock_openai_aembed_document) out_mapping = {"field1": (str, ...), "field2": (List[str], ...)} out_data = {"field1": "field1 value", "field2": ["field2 value1", "field2 value2"]}