fix: fix long-term-mem and knowledge base

Author: zakahan

config.yaml.full

@@ -17,7 +17,7 @@ model:
   embedding:
     name: doubao-embedding-text-240715
     dim: 2560
-    api_base: https://ark.cn-beijing.volces.com/api/v3/embeddings
+    api_base: https://ark.cn-beijing.volces.com/api/v3/
     api_key:
   video:
     name: doubao-seedance-1-0-pro-250528

docs/docs/agent.md

@@ -129,12 +129,12 @@ root_agent:
     backend: local
   knowledgebase:
     backend: opensearch
+    index: test
   tools:
     - module: demo_tool   # tool 所在的模块
       func: greeting      # tool 的函数名称
     - module: tools.tool
       func: count
-  sub_agents:
   sub_agents:
     - ${sub_agent_1}
 

docs/docs/installation.md

@@ -70,7 +70,7 @@ model:
   embedding:
     name: doubao-embedding-text-240715
     dim: 2560
-    api_base: https://ark.cn-beijing.volces.com/api/v3/embeddings
+    api_base: https://ark.cn-beijing.volces.com/api/v3/
     api_key:
 
 volcengine:

pyproject.toml

@@ -28,14 +28,19 @@ dependencies = [
     "cookiecutter>=2.6.0", # For cloud deploy # For OpenSearch database
     "opensearch-py==2.8.0",
     "omegaconf>=2.3.0", # For agent builder
+    "llama-index>=0.14.0",
+    "llama-index-embeddings-openai-like>=0.2.2",
+    "llama-index-llms-openai-like>=0.5.1",
+    "llama-index-vector-stores-opensearch>=0.6.1",
+    "llama-index-vector-stores-redis>=0.6.1",
 ]
 
 [project.scripts]
 veadk = "veadk.cli.cli:veadk"
 
 [project.optional-dependencies]
 database = [
-    "redis>=6.2.0",                 # For Redis database
+    "redis>=5.0",                 # For Redis database
     "pymysql>=1.1.1",               # For MySQL database
     "volcengine>=1.0.193",          # For Viking DB
     "tos>=2.8.4",                   # For TOS storage and Viking DB

tests/test_agent.py

@@ -33,7 +33,7 @@
 
 
 def test_agent():
-    knowledgebase = KnowledgeBase()
+    knowledgebase = KnowledgeBase(index="test_index", backend="local")
     long_term_memory = LongTermMemory(backend="local")
     tracer = OpentelemetryTracer()
 

tests/test_knowledgebase.py

@@ -21,15 +21,14 @@
 async def test_knowledgebase():
     app_name = "kb_test_app"
     key = "Supercalifragilisticexpialidocious"
-    kb = KnowledgeBase(backend="local")
+    kb = KnowledgeBase(backend="local", app_name=app_name)
     # Attempt to delete any existing data for the app_name before adding new data
-    kb.add(
-        data=[f"knowledgebase_id is {key}"],
-        app_name=app_name,
+    kb.add_from_text(
+        text=[f"knowledgebase_id is {key}"],
     )
     res_list = kb.search(
         query="knowledgebase_id",
-        app_name=app_name,
+        top_k=1,
     )
     res = "".join(res_list)
     assert key in res, f"Test failed for backend local res is {res}"

tests/test_long_term_memory.py

@@ -27,7 +27,9 @@
 
 @pytest.mark.asyncio
 async def test_long_term_memory():
-    long_term_memory = LongTermMemory(backend="local")
+    long_term_memory = LongTermMemory(
+        backend="local", app_name=app_name, user_id=user_id
+    )
     agent = Agent(
         name="all_name",
         model_name="test_model_name",

veadk/knowledgebase/backends/base_backend.py

@@ -21,6 +21,10 @@ class BaseKnowledgebaseBackend(ABC, BaseModel):
     index: str
     """Index or collection name of the vector storage."""
 
+    @abstractmethod
+    def precheck_index_naming(self):
+        """Check the index name is valid or not"""
+
     @abstractmethod
     def add_from_directory(self, directory: str, **kwargs) -> bool:
         """Add knowledge from file path to knowledgebase"""

veadk/knowledgebase/backends/in_memory_backend.py

@@ -27,14 +27,17 @@ class InMemoryKnowledgeBackend(BaseKnowledgebaseBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(
             model_name=self.embedding_config.name,
             api_key=self.embedding_config.api_key,
             api_base=self.embedding_config.api_base,
         )
         self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
-        self._retriever = self._vector_index.as_retriever()
 
     @override
     def add_from_directory(self, directory: str) -> bool:
@@ -62,7 +65,8 @@ def add_from_text(self, text: str | list[str]) -> bool:
 
     @override
     def search(self, query: str, top_k: int = 5) -> list[str]:
-        retrieved_nodes = self._retriever.retrieve(query, top_k=top_k)
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]
 
     def _split_documents(self, documents: list[Document]) -> list[BaseNode]:

veadk/knowledgebase/backends/opensearch_backend.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import re
 
 from llama_index.core import (
     Document,
@@ -40,7 +41,19 @@ class OpensearchKnowledgeBackend(BaseKnowledgebaseBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
+    def precheck_index_naming(self):
+        if not (
+            isinstance(self.index, str)
+            and not self.index.startswith(("_", "-"))
+            and self.index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", self.index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the naming rules of OpenSearch"
+            )
+
     def model_post_init(self, __context: Any) -> None:
+        self.precheck_index_naming()
         self._opensearch_client = OpensearchVectorClient(
             endpoint=self.opensearch_config.host,
             port=self.opensearch_config.port,
@@ -71,7 +84,6 @@ def model_post_init(self, __context: Any) -> None:
             storage_context=self._storage_context,
             embed_model=self._embed_model,
         )
-        self._retriever = self._vector_index.as_retriever()
 
     @override
     def add_from_directory(self, directory: str) -> bool:
@@ -99,12 +111,8 @@ def add_from_text(self, text: str | list[str]) -> bool:
 
     @override
     def search(self, query: str, top_k: int = 5) -> list[str]:
-        _original_top_k = self._retriever.similarity_top_k  # type: ignore
-        self._retriever.similarity_top_k = top_k  # type: ignore
-
-        retrieved_nodes = self._retriever.retrieve(query)
-
-        self._retriever.similarity_top_k = _original_top_k  # type: ignore
+        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+        retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]
 
     def _split_documents(self, documents: list[Document]) -> list[BaseNode]: