fix: typechecking and extensions

Author: zakahan

pyproject.toml

@@ -28,6 +28,17 @@ dependencies = [
     "cookiecutter>=2.6.0", # For cloud deploy # For OpenSearch database
     "opensearch-py==2.8.0",
     "omegaconf>=2.3.0", # For agent builder
+]
+
+[project.scripts]
+veadk = "veadk.cli.cli:veadk"
+
+[project.optional-dependencies]
+extensions = [
+    "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
     "llama-index>=0.14.0",
     "llama-index-embeddings-openai-like>=0.2.2",
     "llama-index-llms-openai-like>=0.5.1",
@@ -36,11 +47,6 @@ dependencies = [
     "mcp-server-vikingdb-memory",
     "psycopg2-binary>=2.9.10",
 ]
-
-[project.scripts]
-veadk = "veadk.cli.cli:veadk"
-
-[project.optional-dependencies]
 database = [
     "redis>=5.0",                 # For Redis database
     "pymysql>=1.1.1",               # For MySQL database

veadk/agent.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 from google.adk.agents import LlmAgent, RunConfig
 from google.adk.agents.base_agent import BaseAgent
@@ -32,15 +32,17 @@
     DEFAULT_MODEL_EXTRA_CONFIG,
 )
 from veadk.evaluation import EvalSetRecorder
-from veadk.knowledgebase import KnowledgeBase
-from veadk.memory.long_term_memory import LongTermMemory
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.prompts.agent_default_prompt import DEFAULT_DESCRIPTION, DEFAULT_INSTRUCTION
 from veadk.tracing.base_tracer import BaseTracer
 from veadk.utils.logger import get_logger
 from veadk.utils.patches import patch_asyncio
 from veadk.version import VERSION
 
+if TYPE_CHECKING:
+    from veadk.knowledgebase import KnowledgeBase
+    from veadk.memory.long_term_memory import LongTermMemory
+
 patch_asyncio()
 logger = get_logger(__name__)
 

veadk/knowledgebase/backends/vikingdb_knowledge_backend.py

@@ -90,11 +90,10 @@ def precheck_index_naming(self):
     def model_post_init(self, __context: Any) -> None:
         self.precheck_index_naming()
         # check whether collection exist, if not, create it
-        if not self._collection_exist():
-            logger.info(
-                f"VikingDB knowledgebase collection {self.index} does not exist, creating it..."
+        if not self.collection_status()["existed"]:
+            logger.warning(
+                f"VikingDB knowledgebase collection {self.index} does not exist, please create it first..."
             )
-            self._create_collection()
 
     @override
     def add_from_directory(self, directory: str, **kwargs) -> bool:
@@ -175,7 +174,7 @@ def add_from_text(self, text: str | list[str], **kwargs) -> bool:
     def search(self, query: str, top_k: int = 5) -> list:
         return self._search_knowledge(query=query, top_k=top_k)
 
-    def delete_collection(self):
+    def delete_collection(self) -> bool:
         DELETE_COLLECTION_PATH = "/api/knowledge/collection/delete"
 
         response = self._do_request(
@@ -188,21 +187,25 @@ def delete_collection(self):
         )
 
         if response.get("code") != 0:
-            raise ValueError(f"Error during collection deletion: {response}")
+            logger.error(f"Error during collection deletion: {response}")
+            return False
+        return True
 
-    def delete_doc_by_id(self, doc_id: str):
+    def delete_doc_by_id(self, id: str) -> bool:
         DELETE_DOC_PATH = "/api/knowledge/doc/delete"
         response = self._do_request(
             body={
                 "collection_name": self.index,
                 "project": self.volcengine_project,
+                "doc_id": id,
             },
             path=DELETE_DOC_PATH,
             method="POST",
         )
 
         if response.get("code") != 0:
-            raise ValueError(f"Error during document deletion: {response}")
+            return False
+        return True
 
     def list_docs(self, offset: int = 0, limit: int = -1):
         """List documents in collection.
@@ -272,30 +275,21 @@ def collection_status(self):
             path=COLLECTION_INFO_PATH,
             method="POST",
         )
-        return response
-
-    def _collection_exist(self) -> bool:
-        """List docs in collection, error code 0 means collection exist, error code 1000005 means collection not exist"""
-        LIST_DOC_PATH = "/api/knowledge/doc/list"
-
-        response = self._do_request(
-            body={
-                "collection_name": self.index,
-            },
-            path=LIST_DOC_PATH,
-            method="POST",
-        )
-
-        if response.get("code") == 0:
-            return True
-        elif response.get("code") == 1000005:
-            return False
+        if response["code"] == 0:
+            status = response["data"]["pipeline_list"][0]["index_list"][0]["status"]
+            return {
+                "existed": True,
+                "status": status,
+            }
+        elif response["code"] == 1000005:
+            return {
+                "existed": False,
+                "status": None,
+            }
         else:
-            raise ValueError(
-                f"Error during collection existence check: {response.get('code')}"
-            )
+            raise ValueError(f"Error during collection status: {response}")
 
-    def _create_collection(self) -> None:
+    def create_collection(self) -> None:
         CREATE_COLLECTION_PATH = "/api/knowledge/collection/create"
 
         response = self._do_request(