From 12fc46f902c96133e8d3532f3060e591dbdf9dc9 Mon Sep 17 00:00:00 2001 From: shanchunhua Date: Thu, 9 Jul 2026 20:34:36 +0800 Subject: [PATCH] feat(knowledgebase): add openviking backend --- pyproject.toml | 1 + tests/test_openviking_knowledgebase.py | 319 ++++++++++++++++++ .../builtin_tools/test_load_knowledgebase.py | 124 +++++++ .../backends/openviking_backend.py | 295 ++++++++++++++++ veadk/knowledgebase/knowledgebase.py | 23 +- .../tools/builtin_tools/load_knowledgebase.py | 3 +- veadk/tools/load_knowledgebase_tool.py | 3 +- 7 files changed, 761 insertions(+), 7 deletions(-) create mode 100644 tests/test_openviking_knowledgebase.py create mode 100644 tests/tools/builtin_tools/test_load_knowledgebase.py create mode 100644 veadk/knowledgebase/backends/openviking_backend.py diff --git a/pyproject.toml b/pyproject.toml index 17d78b8f..398634e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "filetype==1.2.0", "vikingdb-python-sdk>=0.1.3", # For Viking DB "agentkit-sdk-python>=0.2.0", + "openviking-sdk>=0.1.3", "volcengine>=1.0.0", # For AgentKit Runtime API access "python-frontmatter==1.1.0", "tos>=2.8.4", # For TOS storage and Viking DB diff --git a/tests/test_openviking_knowledgebase.py b/tests/test_openviking_knowledgebase.py new file mode 100644 index 00000000..e67efb3f --- /dev/null +++ b/tests/test_openviking_knowledgebase.py @@ -0,0 +1,319 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +from __future__ import annotations + +from pathlib import Path +from typing import Any, ClassVar + +import pytest + +from veadk.knowledgebase import KnowledgeBase +from veadk.knowledgebase.backends.openviking_backend import OpenVikingKnowledgeBackend +from veadk.knowledgebase.entry import KnowledgebaseEntry + + +class FakeOpenVikingClient: + def __init__(self) -> None: + self.add_resource_calls: list[dict[str, Any]] = [] + self.added_texts: list[str] = [] + self.find_calls: list[dict[str, Any]] = [] + self.search_calls: list[dict[str, Any]] = [] + self.read_calls: list[dict[str, Any]] = [] + self.overview_calls: list[str] = [] + self.find_response: Any = { + "status": "ok", + "result": { + "resources": [], + "memories": [], + "skills": [], + }, + } + self.search_response: Any = self.find_response + self.read_response = "read body" + self.overview_response = "overview body" + self.fail_read = False + self.fail_overview = False + + def add_resource(self, **kwargs): + self.add_resource_calls.append(kwargs) + path = kwargs.get("path") + if path and Path(path).is_file(): + self.added_texts.append(Path(path).read_text(encoding="utf-8")) + return { + "status": "ok", + "result": {"root_uri": kwargs.get("to") or kwargs.get("parent")}, + } + + def find(self, **kwargs): + self.find_calls.append(kwargs) + return self.find_response + + def search(self, **kwargs): + self.search_calls.append(kwargs) + return self.search_response + + def read(self, uri: str, offset: int = 0, limit: int = -1): + self.read_calls.append({"uri": uri, "offset": offset, "limit": limit}) + if self.fail_read: + raise RuntimeError("read failed") + return self.read_response + + def overview(self, uri: str): + self.overview_calls.append(uri) + if self.fail_overview: + raise RuntimeError("overview failed") + return self.overview_response + + +class FakeOpenVikingKnowledgeBackend(OpenVikingKnowledgeBackend): + fake_client: ClassVar[FakeOpenVikingClient | None] = None + + def _build_client(self): + return self.fake_client or FakeOpenVikingClient() + + +def make_backend( + client: FakeOpenVikingClient | None = None, + **kwargs, +) -> FakeOpenVikingKnowledgeBackend: + FakeOpenVikingKnowledgeBackend.fake_client = client + try: + return FakeOpenVikingKnowledgeBackend(**kwargs) + finally: + FakeOpenVikingKnowledgeBackend.fake_client = None + + +def test_knowledgebase_openviking_instantiates(): + kb = KnowledgeBase( + backend="openviking", + app_name="demo", + backend_config={ + "index": "demo", + "url": "http://127.0.0.1:1933", + "api_key": "test-key", + }, + ) + + assert isinstance(kb._backend, OpenVikingKnowledgeBackend) + assert kb._backend.target_uri == "viking://resources/demo/" + + +def test_missing_index_and_app_name_keeps_existing_error(): + with pytest.raises(ValueError, match="Either `index` or `app_name`"): + KnowledgeBase(backend="openviking") + + +def test_default_target_uri_from_index(): + backend = make_backend(index="demo") + + assert backend.target_uri == "viking://resources/demo/" + + +def test_add_from_files_imports_each_file_under_target_uri(): + client = FakeOpenVikingClient() + backend = make_backend(client, index="demo") + + assert backend.add_from_files(["a.md", "b.md"]) + + assert [call["path"] for call in client.add_resource_calls] == ["a.md", "b.md"] + assert all( + call["parent"] == "viking://resources/demo/" + for call in client.add_resource_calls + ) + assert all(call["wait"] is True for call in client.add_resource_calls) + assert all(call["timeout"] == 300 for call in client.add_resource_calls) + + +def test_add_from_directory_imports_to_target_uri_with_structure(): + client = FakeOpenVikingClient() + backend = make_backend(client, index="demo") + + assert backend.add_from_directory("./docs") + + assert client.add_resource_calls == [ + { + "path": "./docs", + "to": "viking://resources/demo/", + "wait": True, + "timeout": 300, + "strict": False, + "ignore_dirs": None, + "include": None, + "exclude": None, + "directly_upload_media": True, + "preserve_structure": True, + "watch_interval": 0, + "args": None, + "telemetry": False, + } + ] + + +def test_add_from_text_writes_temp_files_and_reuses_file_import(): + client = FakeOpenVikingClient() + backend = make_backend(client, index="demo") + + assert backend.add_from_text(["alpha", "beta"]) + + assert len(client.add_resource_calls) == 2 + assert client.added_texts == ["alpha", "beta"] + + +def test_search_converts_resources_to_entries_without_memories_or_skills(): + client = FakeOpenVikingClient() + client.find_response = { + "status": "ok", + "result": { + "resources": [ + { + "uri": "viking://resources/demo/a.md", + "abstract": "resource abstract", + "score": 0.9, + "match_reason": "semantic", + "context_type": "resource", + "is_leaf": True, + } + ], + "memories": [{"abstract": "memory abstract"}], + "skills": [{"abstract": "skill abstract"}], + }, + } + backend = make_backend( + client, + index="demo", + hydrate_results=False, + ) + + entries = backend.search("policy", top_k=3) + + assert len(entries) == 1 + assert isinstance(entries[0], KnowledgebaseEntry) + assert entries[0].content == "resource abstract" + assert entries[0].metadata == { + "uri": "viking://resources/demo/a.md", + "score": 0.9, + "match_reason": "semantic", + "context_type": "resource", + "is_leaf": True, + } + assert client.find_calls[0]["target_uri"] == "viking://resources/demo/" + assert client.find_calls[0]["limit"] == 3 + + +def test_hydrate_file_resource_reads_content(): + client = FakeOpenVikingClient() + client.find_response = { + "result": { + "resources": [ + { + "uri": "viking://resources/demo/a.md", + "abstract": "abstract", + "is_leaf": True, + } + ] + } + } + backend = make_backend( + client, + index="demo", + hydrate_results=True, + read_limit=123, + ) + + entries = backend.search("policy") + + assert entries[0].content == "read body" + assert client.read_calls == [ + { + "uri": "viking://resources/demo/a.md", + "offset": 0, + "limit": 123, + } + ] + + +def test_hydrate_directory_resource_reads_overview(): + client = FakeOpenVikingClient() + client.find_response = { + "result": { + "resources": [ + { + "uri": "viking://resources/demo/", + "abstract": "abstract", + "is_leaf": False, + } + ] + } + } + backend = make_backend( + client, + index="demo", + hydrate_results=True, + ) + + entries = backend.search("policy") + + assert entries[0].content == "overview body" + assert client.overview_calls == ["viking://resources/demo/"] + + +def test_hydrate_failure_falls_back_to_abstract(): + client = FakeOpenVikingClient() + client.fail_read = True + client.find_response = { + "result": { + "resources": [ + { + "uri": "viking://resources/demo/a.md", + "abstract": "fallback abstract", + "is_leaf": True, + } + ] + } + } + backend = make_backend(client, index="demo") + + entries = backend.search("policy") + + assert entries[0].content == "fallback abstract" + + +def test_score_threshold_and_context_search_are_forwarded(): + client = FakeOpenVikingClient() + client.search_response = { + "result": { + "resources": [ + { + "uri": "viking://resources/demo/a.md", + "abstract": "abstract", + "is_leaf": True, + } + ] + } + } + backend = make_backend( + client, + index="demo", + hydrate_results=False, + score_threshold=0.3, + use_context_search=True, + ) + + backend.search("policy", top_k=7, session_id="s1") + + assert client.find_calls == [] + assert client.search_calls[0]["score_threshold"] == 0.3 + assert client.search_calls[0]["limit"] == 7 + assert client.search_calls[0]["session_id"] == "s1" diff --git a/tests/tools/builtin_tools/test_load_knowledgebase.py b/tests/tools/builtin_tools/test_load_knowledgebase.py new file mode 100644 index 00000000..76c74796 --- /dev/null +++ b/tests/tools/builtin_tools/test_load_knowledgebase.py @@ -0,0 +1,124 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 threading + +import pytest + +from veadk.knowledgebase import KnowledgeBase +from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend +from veadk.knowledgebase.entry import KnowledgebaseEntry +from veadk.tools import load_knowledgebase_tool as legacy_tool +from veadk.tools.builtin_tools.load_knowledgebase import LoadKnowledgebaseTool + + +class FakeSyncBackend(BaseKnowledgebaseBackend): + search_thread_id: int | None = None + + def precheck_index_naming(self) -> None: + pass + + def add_from_directory(self, directory: str, *args, **kwargs) -> bool: + return True + + def add_from_files(self, files: list[str], *args, **kwargs) -> bool: + return True + + def add_from_text(self, text: str | list[str], *args, **kwargs) -> bool: + return True + + def search(self, query: str, top_k: int = 5) -> list[str]: + self.search_thread_id = threading.get_ident() + return [f"generic result for {query}"] + + +class FakeKnowledgeBase: + name = "fake_kb" + backend = "fake" + + def __init__(self): + self.search_thread_id = None + + def search(self, query: str): + self.search_thread_id = threading.get_ident() + return [KnowledgebaseEntry(content=f"result for {query}")] + + +@pytest.mark.asyncio +async def test_load_knowledgebase_runs_sync_search_in_worker_thread(): + event_loop_thread_id = threading.get_ident() + knowledgebase = FakeKnowledgeBase() + tool = LoadKnowledgebaseTool(knowledgebase=knowledgebase) # type: ignore[arg-type] + + response = await tool.load_knowledgebase("openviking", tool_context=None) # type: ignore[arg-type] + + assert response.knowledges == [KnowledgebaseEntry(content="result for openviking")] + assert knowledgebase.search_thread_id is not None + assert knowledgebase.search_thread_id != event_loop_thread_id + + +@pytest.mark.asyncio +async def test_legacy_load_knowledgebase_runs_sync_search_in_worker_thread(): + event_loop_thread_id = threading.get_ident() + knowledgebase = FakeKnowledgeBase() + original_knowledgebase = legacy_tool.knowledgebase + original_knowledgebase_cls = legacy_tool.KnowledgeBase + legacy_tool.knowledgebase = knowledgebase # type: ignore[assignment] + legacy_tool.KnowledgeBase = FakeKnowledgeBase # type: ignore[misc] + + try: + response = await legacy_tool.search_knowledgebase(None, "legacy", "app") + finally: + legacy_tool.knowledgebase = original_knowledgebase + legacy_tool.KnowledgeBase = original_knowledgebase_cls + + assert response.knowledges == [KnowledgebaseEntry(content="result for legacy")] + assert knowledgebase.search_thread_id is not None + assert knowledgebase.search_thread_id != event_loop_thread_id + + +@pytest.mark.asyncio +async def test_load_knowledgebase_to_thread_works_for_generic_knowledgebase_backend(): + event_loop_thread_id = threading.get_ident() + backend = FakeSyncBackend(index="generic") + knowledgebase = KnowledgeBase(backend=backend) + tool = LoadKnowledgebaseTool(knowledgebase=knowledgebase) + + response = await tool.load_knowledgebase("generic", tool_context=None) # type: ignore[arg-type] + + assert response.knowledges == [ + KnowledgebaseEntry(content="generic result for generic") + ] + assert backend.search_thread_id is not None + assert backend.search_thread_id != event_loop_thread_id + + +@pytest.mark.asyncio +async def test_legacy_to_thread_works_for_generic_knowledgebase_backend(): + event_loop_thread_id = threading.get_ident() + backend = FakeSyncBackend(index="generic") + knowledgebase = KnowledgeBase(backend=backend) + original_knowledgebase = legacy_tool.knowledgebase + legacy_tool.knowledgebase = knowledgebase + + try: + response = await legacy_tool.search_knowledgebase(None, "legacy", "app") + finally: + legacy_tool.knowledgebase = original_knowledgebase + + assert response.knowledges == [ + KnowledgebaseEntry(content="generic result for legacy") + ] + assert backend.search_thread_id is not None + assert backend.search_thread_id != event_loop_thread_id diff --git a/veadk/knowledgebase/backends/openviking_backend.py b/veadk/knowledgebase/backends/openviking_backend.py new file mode 100644 index 00000000..9509d509 --- /dev/null +++ b/veadk/knowledgebase/backends/openviking_backend.py @@ -0,0 +1,295 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path +from typing import Any + +from pydantic import Field, PrivateAttr +from typing_extensions import override + +from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend +from veadk.knowledgebase.entry import KnowledgebaseEntry +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +def _getenv(*names: str) -> str | None: + for name in names: + value = os.getenv(name) + if value: + return value + return None + + +def _parse_bool(value: str | None, default: bool) -> bool: + if value is None or value == "": + return default + return value.strip().lower() in {"1", "true", "t", "yes", "y", "on"} + + +def _parse_int(value: str | None, default: int) -> int: + if value is None or value == "": + return default + return int(value) + + +def _parse_float(value: str | None, default: float | None) -> float | None: + if value is None or value == "": + return default + return float(value) + + +def _default_target_uri(index: str) -> str: + return f"viking://resources/{index.strip('/')}/" + + +class OpenVikingKnowledgeBackend(BaseKnowledgebaseBackend): + """OpenViking SDK backend for VeADK knowledge base operations. + + Each VeADK knowledgebase index maps to an OpenViking resources directory, + e.g. ``viking://resources/{index}/``. Imported resources are written under + that directory and searches are scoped with ``target_uri``. OpenViking + performs resource parsing and indexing, so this backend does not require + local embedding configuration. + """ + + url: str | None = Field( + default_factory=lambda: _getenv("DATABASE_OPENVIKING_URL", "OPENVIKING_URL") + ) + api_key: str | None = Field( + default_factory=lambda: _getenv( + "DATABASE_OPENVIKING_API_KEY", "OPENVIKING_API_KEY" + ) + ) + account: str | None = Field( + default_factory=lambda: _getenv( + "DATABASE_OPENVIKING_ACCOUNT", "OPENVIKING_ACCOUNT" + ) + ) + user: str | None = Field( + default_factory=lambda: _getenv("DATABASE_OPENVIKING_USER", "OPENVIKING_USER") + ) + actor_peer_id: str | None = Field( + default_factory=lambda: _getenv( + "DATABASE_OPENVIKING_ACTOR_PEER_ID", "OPENVIKING_ACTOR_PEER_ID" + ) + ) + target_uri: str | None = Field( + default_factory=lambda: _getenv( + "DATABASE_OPENVIKING_TARGET_URI", "OPENVIKING_TARGET_URI" + ) + ) + wait: bool = Field( + default_factory=lambda: _parse_bool( + os.getenv("DATABASE_OPENVIKING_WAIT"), default=True + ) + ) + import_timeout: float | None = Field( + default_factory=lambda: _parse_float( + os.getenv("DATABASE_OPENVIKING_IMPORT_TIMEOUT"), default=300 + ) + ) + hydrate_results: bool = Field( + default_factory=lambda: _parse_bool( + os.getenv("DATABASE_OPENVIKING_HYDRATE_RESULTS"), default=True + ) + ) + read_limit: int = Field( + default_factory=lambda: _parse_int( + os.getenv("DATABASE_OPENVIKING_READ_LIMIT"), default=200 + ) + ) + score_threshold: float | None = Field( + default_factory=lambda: _parse_float( + os.getenv("DATABASE_OPENVIKING_SCORE_THRESHOLD"), default=None + ) + ) + use_context_search: bool = Field( + default_factory=lambda: _parse_bool( + os.getenv("DATABASE_OPENVIKING_USE_CONTEXT_SEARCH"), default=False + ) + ) + + _client: Any = PrivateAttr(default=None) + + def model_post_init(self, __context: Any) -> None: + self.precheck_index_naming() + self.target_uri = self._normalize_target_uri( + self.target_uri or _default_target_uri(self.index) + ) + self._client = self._build_client() + + @override + def precheck_index_naming(self) -> None: + if not isinstance(self.index, str) or not self.index.strip(): + raise ValueError( + "OpenViking knowledgebase index must be a non-empty string." + ) + + def _build_client(self) -> Any: + try: + from openviking_sdk import SyncHTTPClient + except ImportError as first_error: + try: + from openviking.client import SyncHTTPClient + except ImportError as second_error: + raise ImportError( + "OpenViking knowledgebase backend requires 'openviking-sdk'. " + "Please install it via `pip install openviking-sdk>=0.1.3`." + ) from second_error + logger.debug(f"Fallback import used after error: {first_error}") + + return SyncHTTPClient( + url=self.url, + api_key=self.api_key, + account=self.account, + user=self.user, + actor_peer_id=self.actor_peer_id, + ) + + def _normalize_target_uri(self, target_uri: str) -> str: + if not target_uri: + return _default_target_uri(self.index) + if target_uri.startswith("viking://resources/") and not target_uri.endswith( + "/" + ): + return f"{target_uri}/" + return target_uri + + @override + def add_from_directory(self, directory: str, **kwargs) -> bool: + self._client.add_resource( + path=directory, + to=kwargs.get("target_uri", self.target_uri), + wait=kwargs.get("wait", self.wait), + timeout=kwargs.get("timeout", self.import_timeout), + strict=kwargs.get("strict", False), + ignore_dirs=kwargs.get("ignore_dirs"), + include=kwargs.get("include"), + exclude=kwargs.get("exclude"), + directly_upload_media=kwargs.get("directly_upload_media", True), + preserve_structure=kwargs.get("preserve_structure", True), + watch_interval=kwargs.get("watch_interval", 0), + args=kwargs.get("args"), + telemetry=kwargs.get("telemetry", False), + ) + return True + + @override + def add_from_files(self, files: list[str], **kwargs) -> bool: + for file in files: + self._client.add_resource( + path=file, + parent=kwargs.get("target_uri", self.target_uri), + wait=kwargs.get("wait", self.wait), + timeout=kwargs.get("timeout", self.import_timeout), + strict=kwargs.get("strict", False), + reason=kwargs.get("reason", ""), + instruction=kwargs.get("instruction", ""), + directly_upload_media=kwargs.get("directly_upload_media", True), + telemetry=kwargs.get("telemetry", False), + ) + return True + + @override + def add_from_text(self, text: str | list[str], **kwargs) -> bool: + texts = [text] if isinstance(text, str) else text + with tempfile.TemporaryDirectory() as tmpdir: + files = [] + for index, item in enumerate(texts): + path = Path(tmpdir) / f"text_{index}.txt" + path.write_text(item, encoding="utf-8") + files.append(str(path)) + return self.add_from_files(files, **kwargs) + + @override + def search(self, query: str, top_k: int = 5, **kwargs) -> list[KnowledgebaseEntry]: + target_uri = kwargs.get("target_uri", self.target_uri) + score_threshold = kwargs.get("score_threshold", self.score_threshold) + use_context_search = kwargs.get("use_context_search", self.use_context_search) + method = self._client.search if use_context_search else self._client.find + + call_kwargs = { + "query": query, + "target_uri": target_uri, + "limit": top_k, + "score_threshold": score_threshold, + "filter": kwargs.get("filter"), + "context_type": kwargs.get("context_type"), + "tags": kwargs.get("tags"), + "telemetry": kwargs.get("telemetry", False), + } + if use_context_search: + call_kwargs["session"] = kwargs.get("session") + call_kwargs["session_id"] = kwargs.get("session_id") + + result = method(**call_kwargs) + return self._to_entries(result) + + def _to_entries(self, response: Any) -> list[KnowledgebaseEntry]: + result = self._unwrap_result(response) + resources = self._extract_resources(result) + entries: list[KnowledgebaseEntry] = [] + + for item in resources: + if not isinstance(item, dict): + continue + uri = str(item.get("uri") or "") + abstract = item.get("abstract") or "" + content = self._hydrate(uri, item) if self.hydrate_results else abstract + metadata = { + "uri": uri, + "score": item.get("score"), + "match_reason": item.get("match_reason"), + "context_type": item.get("context_type"), + "is_leaf": item.get("is_leaf"), + } + entries.append( + KnowledgebaseEntry( + content=str(content or abstract or ""), + metadata=metadata, + ) + ) + + return entries + + def _unwrap_result(self, response: Any) -> Any: + if isinstance(response, dict) and isinstance(response.get("result"), dict): + return response["result"] + return response + + def _extract_resources(self, result: Any) -> list[Any]: + if isinstance(result, dict): + resources = result.get("resources", []) + return resources if isinstance(resources, list) else [] + if isinstance(result, list): + return result + return [] + + def _hydrate(self, uri: str, item: dict[str, Any]) -> str: + if not uri: + return str(item.get("abstract") or "") + + try: + if item.get("is_leaf"): + return self._client.read(uri, offset=0, limit=self.read_limit) + return self._client.overview(uri) + except Exception as e: + logger.debug(f"Failed to hydrate OpenViking resource {uri}: {e}") + return str(item.get("abstract") or "") diff --git a/veadk/knowledgebase/knowledgebase.py b/veadk/knowledgebase/knowledgebase.py index 93c8eeca..4b00ddc7 100644 --- a/veadk/knowledgebase/knowledgebase.py +++ b/veadk/knowledgebase/knowledgebase.py @@ -71,6 +71,12 @@ def _get_backend_cls(backend: str) -> type[BaseKnowledgebaseBackend]: ) return ContextSearchBackend + case "openviking": + from veadk.knowledgebase.backends.openviking_backend import ( + OpenVikingKnowledgeBackend, + ) + + return OpenVikingKnowledgeBackend case _: raise ValueError(f"Unsupported knowledgebase backend: {backend}") except ImportError as e: @@ -86,20 +92,24 @@ class KnowledgeBase(BaseModel): """A knowledge base for storing user-related information. This class represents a knowledge base used to store and retrieve user-specific data. - It supports multiple backend options, including in-memory, OpenSearch, Redis, and Volcengine's - VikingDB. The knowledge base allows for efficient document retrieval based on similarity, - with the ability to configure backend-specific settings. + It supports multiple backend options, including in-memory, OpenSearch, Redis, + Volcengine VikingDB, Context Search, and OpenViking resources. The knowledge + base allows document retrieval through the selected backend, with backend- + specific settings. Attributes: name (str): The name of the knowledge base. Default is "user_knowledgebase". description (str): A description of the knowledge base. Default is "This knowledgebase stores some user-related information." - backend (Union[Literal["local", "opensearch", "viking", "redis", "milvus"], BaseKnowledgebaseBackend]): + backend (Union[Literal["local", "opensearch", "viking", "redis", "milvus", "tos_vector", "context_search", "openviking"], BaseKnowledgebaseBackend]): The type of backend to use for storing and querying the knowledge base. Supported options include: - 'local' for in-memory storage (data is lost when the program exits). - 'opensearch' for OpenSearch (requires OpenSearch cluster). - 'viking' for Volcengine VikingDB (requires VikingDB service). - 'redis' for Redis with vector search capability (requires Redis). - 'milvus' for Milvus vector database (requires Milvus). + - 'tos_vector' for TOS-backed vector retrieval. + - 'context_search' for Volcengine Context Search. + - 'openviking' for OpenViking resources-based retrieval. Default is 'local'. backend_config (dict): Configuration dictionary for the selected backend. top_k (int): The number of top similar documents to retrieve during a search. Default is 10. @@ -107,7 +117,9 @@ class KnowledgeBase(BaseModel): index (str): The name of the knowledge base index. Notes: - Please ensure that you have set the embedding-related configurations in environment variables. + Vector backends such as 'local', 'opensearch', and 'redis' require + embedding-related environment variables. Resource/search-service backends + such as 'openviking' and 'context_search' use their own service configs. """ name: str = "user_knowledgebase" @@ -123,6 +135,7 @@ class KnowledgeBase(BaseModel): "milvus", "tos_vector", "context_search", + "openviking", ], BaseKnowledgebaseBackend, ] = "local" diff --git a/veadk/tools/builtin_tools/load_knowledgebase.py b/veadk/tools/builtin_tools/load_knowledgebase.py index 51b80944..6ae6e2c5 100644 --- a/veadk/tools/builtin_tools/load_knowledgebase.py +++ b/veadk/tools/builtin_tools/load_knowledgebase.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio import json from google.adk.models.llm_request import LlmRequest @@ -161,6 +162,6 @@ async def load_knowledgebase( A list of knowledgebase results. """ logger.info(f"Search knowledgebase: {self.knowledgebase.name}") - response = self.knowledgebase.search(query) + response = await asyncio.to_thread(self.knowledgebase.search, query) logger.info(f"Loaded {len(response)} knowledgebase entries for query: {query}") return LoadKnowledgebaseResponse(knowledges=response) diff --git a/veadk/tools/load_knowledgebase_tool.py b/veadk/tools/load_knowledgebase_tool.py index b087957b..d5204da7 100644 --- a/veadk/tools/load_knowledgebase_tool.py +++ b/veadk/tools/load_knowledgebase_tool.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio from typing import TYPE_CHECKING import wrapt @@ -55,7 +56,7 @@ async def search_knowledgebase( ) -> SearchKnowledgebaseResponse: """Searches the knowledgebase of the current user.""" if isinstance(knowledgebase, KnowledgeBase): - entry_list = knowledgebase.search(query) + entry_list = await asyncio.to_thread(knowledgebase.search, query) return SearchKnowledgebaseResponse(knowledges=entry_list) else: return SearchKnowledgebaseResponse(knowledges=[])