diff --git a/.env.example b/.env.example index a45a4f9..83ceca6 100644 --- a/.env.example +++ b/.env.example @@ -14,6 +14,12 @@ TAVILY_API_KEY=tvly-xxx # JINA_API_KEY=jina_xxx # Optional, default is None # Optional, RAG provider +# RAG_PROVIDER=vikingdb_knowledge_base +# VIKINGDB_KNOWLEDGE_BASE_API_URL="api-knowledgebase.mlp.cn-beijing.volces.com" +# VIKINGDB_KNOWLEDGE_BASE_API_AK="AKxxx" +# VIKINGDB_KNOWLEDGE_BASE_API_SK="" +# VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE=15 + # RAG_PROVIDER=ragflow # RAGFLOW_API_URL="http://localhost:9388" # RAGFLOW_API_KEY="ragflow-xxx" diff --git a/pyproject.toml b/pyproject.toml index 2414ca4..72f829c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "mcp>=1.6.0", "langchain-mcp-adapters>=0.0.9", "langchain-deepseek>=0.1.3", + "volcengine>=1.0.191", ] [project.optional-dependencies] diff --git a/src/config/tools.py b/src/config/tools.py index 9462c3b..ca0f300 100644 --- a/src/config/tools.py +++ b/src/config/tools.py @@ -21,6 +21,7 @@ SELECTED_SEARCH_ENGINE = os.getenv("SEARCH_API", SearchEngine.TAVILY.value) class RAGProvider(enum.Enum): RAGFLOW = "ragflow" + VIKINGDB_KNOWLEDGE_BASE = "vikingdb_knowledge_base" SELECTED_RAG_PROVIDER = os.getenv("RAG_PROVIDER") diff --git a/src/rag/__init__.py b/src/rag/__init__.py index e271cfc..eed21ca 100644 --- a/src/rag/__init__.py +++ b/src/rag/__init__.py @@ -3,6 +3,15 @@ from .retriever import Retriever, Document, Resource, Chunk from .ragflow import RAGFlowProvider +from .vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider from .builder import build_retriever -__all__ = [Retriever, Document, Resource, RAGFlowProvider, Chunk, build_retriever] +__all__ = [ + Retriever, + Document, + Resource, + RAGFlowProvider, + VikingDBKnowledgeBaseProvider, + Chunk, + build_retriever, +] diff --git a/src/rag/builder.py b/src/rag/builder.py index 6168f93..b7a6a07 100644 --- a/src/rag/builder.py +++ b/src/rag/builder.py @@ -3,12 +3,15 @@ from src.config.tools import SELECTED_RAG_PROVIDER, RAGProvider from src.rag.ragflow import RAGFlowProvider +from src.rag.vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider from src.rag.retriever import Retriever def build_retriever() -> Retriever | None: if SELECTED_RAG_PROVIDER == RAGProvider.RAGFLOW.value: return RAGFlowProvider() + elif SELECTED_RAG_PROVIDER == RAGProvider.VIKINGDB_KNOWLEDGE_BASE.value: + return VikingDBKnowledgeBaseProvider() elif SELECTED_RAG_PROVIDER: raise ValueError(f"Unsupported RAG provider: {SELECTED_RAG_PROVIDER}") return None diff --git a/src/rag/vikingdb_knowledge_base.py b/src/rag/vikingdb_knowledge_base.py new file mode 100644 index 0000000..4f2672f --- /dev/null +++ b/src/rag/vikingdb_knowledge_base.py @@ -0,0 +1,208 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: MIT + +import os +import requests +import json +from src.rag.retriever import Chunk, Document, Resource, Retriever +from urllib.parse import urlparse +from volcengine.auth.SignerV4 import SignerV4 +from volcengine.base.Request import Request +from volcengine.Credentials import Credentials + + +class VikingDBKnowledgeBaseProvider(Retriever): + """ + VikingDBKnowledgeBaseProvider is a provider that uses VikingDB Knowledge base API to retrieve documents. + """ + + api_url: str + api_ak: str + api_sk: str + retrieval_size: int = 10 + + def __init__(self): + api_url = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_URL") + if not api_url: + raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_URL is not set") + self.api_url = api_url + + api_ak = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_AK") + if not api_ak: + raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_AK is not set") + self.api_ak = api_ak + + api_sk = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_SK") + if not api_sk: + raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_SK is not set") + self.api_sk = api_sk + + retrieval_size = os.getenv("VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE") + if retrieval_size: + self.retrieval_size = int(retrieval_size) + + def prepare_request(self, method, path, params=None, data=None, doseq=0): + """ + Prepare signed request using volcengine auth + """ + if params: + for key in params: + if ( + type(params[key]) == int + or type(params[key]) == float + or type(params[key]) == bool + ): + params[key] = str(params[key]) + elif type(params[key]) == list: + if not doseq: + params[key] = ",".join(params[key]) + + r = Request() + r.set_shema("https") + r.set_method(method) + r.set_connection_timeout(10) + r.set_socket_timeout(10) + mheaders = { + "Accept": "application/json", + "Content-Type": "application/json", + } + r.set_headers(mheaders) + if params: + r.set_query(params) + r.set_path(path) + if data is not None: + r.set_body(json.dumps(data)) + + credentials = Credentials(self.api_ak, self.api_sk, "air", "cn-north-1") + SignerV4.sign(r, credentials) + return r + + def query_relevant_documents( + self, query: str, resources: list[Resource] = [] + ) -> list[Document]: + """ + Query relevant documents from the knowledge base + """ + if not resources: + return [] + + all_documents = {} + for resource in resources: + resource_id, document_id = parse_uri(resource.uri) + request_params = { + "resource_id": resource_id, + "query": query, + "limit": self.retrieval_size, + "dense_weight": 0.5, + "pre_processing": { + "need_instruction": True, + "rewrite": False, + "return_token_usage": True, + }, + "post_processing": { + "rerank_switch": True, + "chunk_diffusion_count": 0, + "chunk_group": True, + "get_attachment_link": True, + }, + } + if document_id: + doc_filter = {"op": "must", "field": "doc_id", "conds": [document_id]} + query_param = {"doc_filter": doc_filter} + request_params["query_param"] = query_param + + method = "POST" + path = "/api/knowledge/collection/search_knowledge" + info_req = self.prepare_request( + method=method, path=path, data=request_params + ) + rsp = requests.request( + method=info_req.method, + url="http://{}{}".format(self.api_url, info_req.path), + headers=info_req.headers, + data=info_req.body, + ) + + try: + response = json.loads(rsp.text) + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse JSON response: {e}") + + if response["code"] != 0: + raise ValueError( + f"Failed to query documents from resource: {response['message']}" + ) + + rsp_data = response.get("data", {}) + + if "result_list" not in rsp_data: + continue + + result_list = rsp_data["result_list"] + + for item in result_list: + doc_info = item.get("doc_info", {}) + doc_id = doc_info.get("doc_id") + + if not doc_id: + continue + + if doc_id not in all_documents: + all_documents[doc_id] = Document( + id=doc_id, title=doc_info.get("doc_name"), chunks=[] + ) + + chunk = Chunk( + content=item.get("content", ""), similarity=item.get("score", 0.0) + ) + all_documents[doc_id].chunks.append(chunk) + + return list(all_documents.values()) + + def list_resources(self, query: str | None = None) -> list[Resource]: + """ + List resources (knowledge bases) from the knowledge base service + """ + method = "POST" + path = "/api/knowledge/collection/list" + info_req = self.prepare_request(method=method, path=path) + rsp = requests.request( + method=info_req.method, + url="http://{}{}".format(self.api_url, info_req.path), + headers=info_req.headers, + data=info_req.body, + ) + try: + response = json.loads(rsp.text) + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse JSON response: {e}") + + if response["code"] != 0: + raise Exception(f"Failed to list resources: {response["message"]}") + + resources = [] + rsp_data = response.get("data", {}) + collection_list = rsp_data.get("collection_list", []) + for item in collection_list: + collection_name = item.get("collection_name", "") + description = item.get("description", "") + + if query and query.lower() not in collection_name.lower(): + continue + + resource_id = item.get("resource_id", "") + resource = Resource( + uri=f"rag://dataset/{resource_id}", + title=collection_name, + description=description, + ) + resources.append(resource) + + return resources + + +def parse_uri(uri: str) -> tuple[str, str]: + parsed = urlparse(uri) + if parsed.scheme != "rag": + raise ValueError(f"Invalid URI: {uri}") + return parsed.path.split("/")[1], parsed.fragment diff --git a/tests/unit/rag/test_vikingdb_knowledge_base.py b/tests/unit/rag/test_vikingdb_knowledge_base.py new file mode 100644 index 0000000..22203e7 --- /dev/null +++ b/tests/unit/rag/test_vikingdb_knowledge_base.py @@ -0,0 +1,503 @@ +# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: MIT + +import os +import pytest +import json +from unittest.mock import patch, MagicMock +from src.rag.vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider, parse_uri + + +# Dummy classes to mock dependencies +class MockResource: + def __init__(self, uri, title="", description=""): + self.uri = uri + self.title = title + self.description = description + + +class MockChunk: + def __init__(self, content, similarity): + self.content = content + self.similarity = similarity + + +class MockDocument: + def __init__(self, id, title, chunks=None): + self.id = id + self.title = title + self.chunks = chunks or [] + + +# Patch the imports to use mock classes +@pytest.fixture(autouse=True) +def patch_imports(): + with ( + patch("src.rag.vikingdb_knowledge_base.Resource", MockResource), + patch("src.rag.vikingdb_knowledge_base.Chunk", MockChunk), + patch("src.rag.vikingdb_knowledge_base.Document", MockDocument), + ): + yield + + +@pytest.fixture +def env_vars(): + """Fixture to set up environment variables""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com", + "VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak", + "VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk", + "VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "10", + }, + ): + yield + + +class TestParseUri: + def test_parse_uri_valid_with_fragment(self): + """Test parsing valid URI with fragment""" + uri = "rag://dataset/123#doc456" + resource_id, document_id = parse_uri(uri) + assert resource_id == "123" + assert document_id == "doc456" + + def test_parse_uri_valid_without_fragment(self): + """Test parsing valid URI without fragment""" + uri = "rag://dataset/123" + resource_id, document_id = parse_uri(uri) + assert resource_id == "123" + assert document_id == "" + + def test_parse_uri_invalid_scheme(self): + """Test parsing URI with invalid scheme""" + with pytest.raises(ValueError, match="Invalid URI"): + parse_uri("http://dataset/123#abc") + + def test_parse_uri_malformed(self): + """Test parsing malformed URI""" + with pytest.raises(ValueError, match="Invalid URI"): + parse_uri("invalid_uri") + + +class TestVikingDBKnowledgeBaseProviderInit: + def test_init_success_with_all_env_vars(self, env_vars): + """Test successful initialization with all environment variables""" + provider = VikingDBKnowledgeBaseProvider() + assert provider.api_url == "api-test.example.com" + assert provider.api_ak == "test_ak" + assert provider.api_sk == "test_sk" + assert provider.retrieval_size == 10 + + def test_init_success_without_retrieval_size(self): + """Test initialization without VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE (should use default)""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com", + "VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak", + "VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk", + }, + clear=True, + ): + provider = VikingDBKnowledgeBaseProvider() + assert provider.retrieval_size == 10 + + def test_init_custom_retrieval_size(self): + """Test initialization with custom retrieval size""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com", + "VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak", + "VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk", + "VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "5", + }, + ): + provider = VikingDBKnowledgeBaseProvider() + assert provider.retrieval_size == 5 + + def test_init_missing_api_url(self): + """Test initialization fails when API URL is missing""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak", + "VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk", + }, + clear=True, + ): + with pytest.raises( + ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_URL is not set" + ): + VikingDBKnowledgeBaseProvider() + + def test_init_missing_api_ak(self): + """Test initialization fails when API AK is missing""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com", + "VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk", + }, + clear=True, + ): + with pytest.raises( + ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_AK is not set" + ): + VikingDBKnowledgeBaseProvider() + + def test_init_missing_api_sk(self): + """Test initialization fails when API SK is missing""" + with patch.dict( + os.environ, + { + "VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com", + "VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak", + }, + clear=True, + ): + with pytest.raises( + ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_SK is not set" + ): + VikingDBKnowledgeBaseProvider() + + +class TestVikingDBKnowledgeBaseProviderPrepareRequest: + @pytest.fixture + def provider(self, env_vars): + return VikingDBKnowledgeBaseProvider() + + def test_prepare_request_basic(self, provider): + """Test basic request preparation""" + with ( + patch("src.rag.vikingdb_knowledge_base.Request") as mock_request, + patch("src.rag.vikingdb_knowledge_base.Credentials") as mock_credentials, + patch("src.rag.vikingdb_knowledge_base.SignerV4.sign") as mock_sign, + ): + + mock_req_instance = MagicMock() + mock_request.return_value = mock_req_instance + + result = provider.prepare_request("POST", "/test/path") + + assert result == mock_req_instance + mock_req_instance.set_shema.assert_called_once_with("https") + mock_req_instance.set_method.assert_called_once_with("POST") + mock_req_instance.set_path.assert_called_once_with("/test/path") + + def test_prepare_request_with_params(self, provider): + """Test request preparation with parameters""" + with ( + patch("src.rag.vikingdb_knowledge_base.Request") as mock_request, + patch("src.rag.vikingdb_knowledge_base.Credentials"), + patch("src.rag.vikingdb_knowledge_base.SignerV4.sign"), + ): + + mock_req_instance = MagicMock() + mock_request.return_value = mock_req_instance + + params = {"key": "value", "number": 123, "boolean": True} + provider.prepare_request("GET", "/test", params=params) + + expected_params = {"key": "value", "number": "123", "boolean": "True"} + mock_req_instance.set_query.assert_called_once_with(expected_params) + + def test_prepare_request_with_data(self, provider): + """Test request preparation with data""" + with ( + patch("src.rag.vikingdb_knowledge_base.Request") as mock_request, + patch("src.rag.vikingdb_knowledge_base.Credentials"), + patch("src.rag.vikingdb_knowledge_base.SignerV4.sign"), + ): + + mock_req_instance = MagicMock() + mock_request.return_value = mock_req_instance + + data = {"test": "data"} + provider.prepare_request("POST", "/test", data=data) + + mock_req_instance.set_body.assert_called_once_with(json.dumps(data)) + + +class TestVikingDBKnowledgeBaseProviderQueryRelevantDocuments: + @pytest.fixture + def provider(self, env_vars): + return VikingDBKnowledgeBaseProvider() + + def test_query_relevant_documents_empty_resources(self, provider): + """Test querying with empty resources list""" + result = provider.query_relevant_documents("test query", []) + assert result == [] + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_query_relevant_documents_success(self, mock_request, provider): + """Test successful document query""" + # Mock response + mock_response = MagicMock() + mock_response.text = json.dumps( + { + "code": 0, + "data": { + "result_list": [ + { + "doc_info": { + "doc_id": "doc123", + "doc_name": "Test Document", + }, + "content": "Test content", + "score": 0.95, + } + ] + }, + } + ) + mock_request.return_value = mock_response + + # Mock prepare_request + with patch.object(provider, "prepare_request") as mock_prepare: + mock_req = MagicMock() + mock_req.method = "POST" + mock_req.path = "/api/knowledge/collection/search_knowledge" + mock_req.headers = {} + mock_req.body = "{}" + mock_prepare.return_value = mock_req + + resources = [MockResource("rag://dataset/123")] + result = provider.query_relevant_documents("test query", resources) + + assert len(result) == 1 + assert result[0].id == "doc123" + assert result[0].title == "Test Document" + assert len(result[0].chunks) == 1 + assert result[0].chunks[0].content == "Test content" + assert result[0].chunks[0].similarity == 0.95 + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_query_relevant_documents_with_document_filter( + self, mock_request, provider + ): + """Test document query with document ID filter""" + mock_response = MagicMock() + mock_response.text = json.dumps({"code": 0, "data": {"result_list": []}}) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request") as mock_prepare: + mock_req = MagicMock() + mock_prepare.return_value = mock_req + + resources = [MockResource("rag://dataset/123#doc456")] + provider.query_relevant_documents("test query", resources) + + # Verify that query_param with doc_filter was included in the request + call_args = mock_prepare.call_args + request_data = call_args[1]["data"] + assert "query_param" in request_data + assert "doc_filter" in request_data["query_param"] + + doc_filter = request_data["query_param"]["doc_filter"] + assert doc_filter["op"] == "must" + assert doc_filter["field"] == "doc_id" + assert doc_filter["conds"] == ["doc456"] + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_query_relevant_documents_api_error(self, mock_request, provider): + """Test handling of API error response""" + mock_response = MagicMock() + mock_response.text = json.dumps({"code": 1, "message": "API Error"}) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + resources = [MockResource("rag://dataset/123")] + with pytest.raises( + ValueError, match="Failed to query documents from resource: API Error" + ): + provider.query_relevant_documents("test query", resources) + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_query_relevant_documents_json_decode_error(self, mock_request, provider): + """Test handling of JSON decode error""" + mock_response = MagicMock() + mock_response.text = "invalid json" + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + resources = [MockResource("rag://dataset/123")] + with pytest.raises(ValueError, match="Failed to parse JSON response"): + provider.query_relevant_documents("test query", resources) + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_query_relevant_documents_multiple_resources(self, mock_request, provider): + """Test querying multiple resources and merging results""" + # Mock responses for different resources + responses = [ + json.dumps( + { + "code": 0, + "data": { + "result_list": [ + { + "doc_info": { + "doc_id": "doc1", + "doc_name": "Document 1", + }, + "content": "Content 1", + "score": 0.9, + } + ] + }, + } + ), + json.dumps( + { + "code": 0, + "data": { + "result_list": [ + { + "doc_info": { + "doc_id": "doc1", + "doc_name": "Document 1", + }, + "content": "Content 2", + "score": 0.8, + }, + { + "doc_info": { + "doc_id": "doc2", + "doc_name": "Document 2", + }, + "content": "Content 3", + "score": 0.7, + }, + ] + }, + } + ), + ] + + mock_request.side_effect = [MagicMock(text=resp) for resp in responses] + + with patch.object(provider, "prepare_request"): + resources = [ + MockResource("rag://dataset/123"), + MockResource("rag://dataset/456"), + ] + result = provider.query_relevant_documents("test query", resources) + + # Should have 2 documents: doc1 (with 2 chunks) and doc2 (with 1 chunk) + assert len(result) == 2 + doc1 = next(doc for doc in result if doc.id == "doc1") + doc2 = next(doc for doc in result if doc.id == "doc2") + assert len(doc1.chunks) == 2 + assert len(doc2.chunks) == 1 + + +class TestVikingDBKnowledgeBaseProviderListResources: + @pytest.fixture + def provider(self, env_vars): + return VikingDBKnowledgeBaseProvider() + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_list_resources_success(self, mock_request, provider): + """Test successful resource listing""" + mock_response = MagicMock() + mock_response.text = json.dumps( + { + "code": 0, + "data": { + "collection_list": [ + { + "resource_id": "123", + "collection_name": "Dataset 1", + "description": "Description 1", + }, + { + "resource_id": "456", + "collection_name": "Dataset 2", + "description": "Description 2", + }, + ] + }, + } + ) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request") as mock_prepare: + mock_req = MagicMock() + mock_prepare.return_value = mock_req + + result = provider.list_resources() + + assert len(result) == 2 + assert result[0].uri == "rag://dataset/123" + assert result[0].title == "Dataset 1" + assert result[0].description == "Description 1" + assert result[1].uri == "rag://dataset/456" + assert result[1].title == "Dataset 2" + assert result[1].description == "Description 2" + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_list_resources_with_query_filter(self, mock_request, provider): + """Test resource listing with query filter""" + mock_response = MagicMock() + mock_response.text = json.dumps( + { + "code": 0, + "data": { + "collection_list": [ + { + "resource_id": "123", + "collection_name": "Test Dataset", + "description": "Description", + }, + { + "resource_id": "456", + "collection_name": "Other Dataset", + "description": "Description", + }, + ] + }, + } + ) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + result = provider.list_resources("test") + + # Should only return the dataset with "test" in the name + assert len(result) == 1 + assert result[0].title == "Test Dataset" + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_list_resources_api_error(self, mock_request, provider): + """Test handling of API error in list_resources""" + mock_response = MagicMock() + mock_response.text = json.dumps({"code": 1, "message": "API Error"}) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + with pytest.raises(Exception, match="Failed to list resources: API Error"): + provider.list_resources() + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_list_resources_json_decode_error(self, mock_request, provider): + """Test handling of JSON decode error in list_resources""" + mock_response = MagicMock() + mock_response.text = "invalid json" + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + with pytest.raises(ValueError, match="Failed to parse JSON response"): + provider.list_resources() + + @patch("src.rag.vikingdb_knowledge_base.requests.request") + def test_list_resources_empty_response(self, mock_request, provider): + """Test handling of empty response""" + mock_response = MagicMock() + mock_response.text = json.dumps({"code": 0, "data": {"collection_list": []}}) + mock_request.return_value = mock_response + + with patch.object(provider, "prepare_request"): + result = provider.list_resources() + assert result == [] diff --git a/uv.lock b/uv.lock index de614af..aaac8f2 100644 --- a/uv.lock +++ b/uv.lock @@ -365,6 +365,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, ] +[[package]] +name = "decorator" +version = "5.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, +] + [[package]] name = "deer-flow" version = "0.1.0" @@ -393,6 +402,7 @@ dependencies = [ { name = "socksio" }, { name = "sse-starlette" }, { name = "uvicorn" }, + { name = "volcengine" }, { name = "yfinance" }, ] @@ -437,6 +447,7 @@ requires-dist = [ { name = "socksio", specifier = ">=1.0.0" }, { name = "sse-starlette", specifier = ">=1.6.5" }, { name = "uvicorn", specifier = ">=0.27.1" }, + { name = "volcengine", specifier = ">=1.0.191" }, { name = "yfinance", specifier = ">=0.2.54" }, ] provides-extras = ["dev", "test"] @@ -564,6 +575,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/53/eb690efa8513166adef3e0669afd31e95ffde69fb3c52ec2ac7223ed6018/fsspec-2025.3.0-py3-none-any.whl", hash = "sha256:efb87af3efa9103f94ca91a7f8cb7a4df91af9f74fc106c9c7ea0efd7277c1b3", size = 193615, upload-time = "2025-03-07T21:47:54.809Z" }, ] +[[package]] +name = "google" +version = "3.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "beautifulsoup4" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/97/b49c69893cddea912c7a660a4b6102c6b02cd268f8c7162dd70b7c16f753/google-3.0.0.tar.gz", hash = "sha256:143530122ee5130509ad5e989f0512f7cb218b2d4eddbafbad40fd10e8d8ccbe", size = 44978, upload-time = "2020-07-11T14:50:45.678Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/35/17c9141c4ae21e9a29a43acdfd848e3e468a810517f862cad07977bf8fe9/google-3.0.0-py2.py3-none-any.whl", hash = "sha256:889cf695f84e4ae2c55fbc0cfdaf4c1e729417fa52ab1db0485202ba173e4935", size = 45258, upload-time = "2020-07-11T14:49:58.287Z" }, +] + [[package]] name = "greenlet" version = "3.1.1" @@ -1595,6 +1618,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b5/35/6c4c6fc8774a9e3629cd750dc24a7a4fb090a25ccd5c3246d127b70f9e22/propcache-0.3.0-py3-none-any.whl", hash = "sha256:67dda3c7325691c2081510e92c561f465ba61b975f481735aefdfc845d2cd043", size = 12101, upload-time = "2025-02-20T19:03:27.202Z" }, ] +[[package]] +name = "protobuf" +version = "6.31.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/52/f3/b9655a711b32c19720253f6f06326faf90580834e2e83f840472d752bc8b/protobuf-6.31.1.tar.gz", hash = "sha256:d8cac4c982f0b957a4dc73a80e2ea24fab08e679c0de9deb835f4a12d69aca9a", size = 441797, upload-time = "2025-05-28T19:25:54.947Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/6f/6ab8e4bf962fd5570d3deaa2d5c38f0a363f57b4501047b5ebeb83ab1125/protobuf-6.31.1-cp310-abi3-win32.whl", hash = "sha256:7fa17d5a29c2e04b7d90e5e32388b8bfd0e7107cd8e616feef7ed3fa6bdab5c9", size = 423603, upload-time = "2025-05-28T19:25:41.198Z" }, + { url = "https://files.pythonhosted.org/packages/44/3a/b15c4347dd4bf3a1b0ee882f384623e2063bb5cf9fa9d57990a4f7df2fb6/protobuf-6.31.1-cp310-abi3-win_amd64.whl", hash = "sha256:426f59d2964864a1a366254fa703b8632dcec0790d8862d30034d8245e1cd447", size = 435283, upload-time = "2025-05-28T19:25:44.275Z" }, + { url = "https://files.pythonhosted.org/packages/6a/c9/b9689a2a250264a84e66c46d8862ba788ee7a641cdca39bccf64f59284b7/protobuf-6.31.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:6f1227473dc43d44ed644425268eb7c2e488ae245d51c6866d19fe158e207402", size = 425604, upload-time = "2025-05-28T19:25:45.702Z" }, + { url = "https://files.pythonhosted.org/packages/76/a1/7a5a94032c83375e4fe7e7f56e3976ea6ac90c5e85fac8576409e25c39c3/protobuf-6.31.1-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:a40fc12b84c154884d7d4c4ebd675d5b3b5283e155f324049ae396b95ddebc39", size = 322115, upload-time = "2025-05-28T19:25:47.128Z" }, + { url = "https://files.pythonhosted.org/packages/fa/b1/b59d405d64d31999244643d88c45c8241c58f17cc887e73bcb90602327f8/protobuf-6.31.1-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:4ee898bf66f7a8b0bd21bce523814e6fbd8c6add948045ce958b73af7e8878c6", size = 321070, upload-time = "2025-05-28T19:25:50.036Z" }, + { url = "https://files.pythonhosted.org/packages/f7/af/ab3c51ab7507a7325e98ffe691d9495ee3d3aa5f589afad65ec920d39821/protobuf-6.31.1-py3-none-any.whl", hash = "sha256:720a6c7e6b77288b85063569baae8536671b39f15cc22037ec7045658d80489e", size = 168724, upload-time = "2025-05-28T19:25:53.926Z" }, +] + +[[package]] +name = "py" +version = "1.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/98/ff/fec109ceb715d2a6b4c4a85a61af3b40c723a961e8828319fbcb15b868dc/py-1.11.0.tar.gz", hash = "sha256:51c75c4126074b472f746a24399ad32f6053d1b34b68d2fa41e558e6f4a98719", size = 207796, upload-time = "2021-11-04T17:17:01.377Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/f0/10642828a8dfb741e5f3fbaac830550a518a775c7fff6f04a007259b0548/py-1.11.0-py2.py3-none-any.whl", hash = "sha256:607c53218732647dff4acdfcd50cb62615cedf612e72d1724fb1a0cc6405b378", size = 98708, upload-time = "2021-11-04T17:17:00.152Z" }, +] + [[package]] name = "pycparser" version = "2.22" @@ -1604,6 +1650,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/13/a3/a812df4e2dd5696d1f351d58b8fe16a405b234ad2886a0dab9183fb78109/pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc", size = 117552, upload-time = "2024-03-30T13:22:20.476Z" }, ] +[[package]] +name = "pycryptodome" +version = "3.9.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/3a/5bca2cb1648b171afd6b7d29a11c6bca8b305bb75b7e2d78a0f5c61ff95e/pycryptodome-3.9.9.tar.gz", hash = "sha256:910e202a557e1131b1c1b3f17a63914d57aac55cf9fb9b51644962841c3995c4", size = 15488528, upload-time = "2020-11-03T13:15:26.723Z" } + [[package]] name = "pydantic" version = "2.10.6" @@ -1701,9 +1753,9 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pytest" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d0/d4/14f53324cb1a6381bef29d698987625d80052bb33932d8e7cbf9b337b17c/pytest_asyncio-1.0.0.tar.gz", hash = "sha256:d15463d13f4456e1ead2594520216b225a16f781e144f8fdf6c5bb4667c48b3f", size = 46960 } +sdist = { url = "https://files.pythonhosted.org/packages/d0/d4/14f53324cb1a6381bef29d698987625d80052bb33932d8e7cbf9b337b17c/pytest_asyncio-1.0.0.tar.gz", hash = "sha256:d15463d13f4456e1ead2594520216b225a16f781e144f8fdf6c5bb4667c48b3f", size = 46960, upload-time = "2025-05-26T04:54:40.484Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/30/05/ce271016e351fddc8399e546f6e23761967ee09c8c568bbfbecb0c150171/pytest_asyncio-1.0.0-py3-none-any.whl", hash = "sha256:4f024da9f1ef945e680dc68610b52550e36590a67fd31bb3b4943979a1f90ef3", size = 15976 }, + { url = "https://files.pythonhosted.org/packages/30/05/ce271016e351fddc8399e546f6e23761967ee09c8c568bbfbecb0c150171/pytest_asyncio-1.0.0-py3-none-any.whl", hash = "sha256:4f024da9f1ef945e680dc68610b52550e36590a67fd31bb3b4943979a1f90ef3", size = 15976, upload-time = "2025-05-26T04:54:39.035Z" }, ] [[package]] @@ -1869,6 +1921,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/51/d4db610ef29373b879047326cbf6fa98b6c1969d6f6dc423279de2b1be2c/requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06", size = 54481, upload-time = "2023-05-01T04:11:28.427Z" }, ] +[[package]] +name = "retry" +version = "0.9.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "decorator" }, + { name = "py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/72/75d0b85443fbc8d9f38d08d2b1b67cc184ce35280e4a3813cda2f445f3a4/retry-0.9.2.tar.gz", hash = "sha256:f8bfa8b99b69c4506d6f5bd3b0aabf77f98cdb17f3c9fc3f5ca820033336fba4", size = 6448, upload-time = "2016-05-11T13:58:51.541Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/0d/53aea75710af4528a25ed6837d71d117602b01946b307a3912cb3cfcbcba/retry-0.9.2-py2.py3-none-any.whl", hash = "sha256:ccddf89761fa2c726ab29391837d4327f819ea14d244c232a1d24c67a2f98606", size = 7986, upload-time = "2016-05-11T13:58:39.925Z" }, +] + [[package]] name = "rpds-py" version = "0.23.1" @@ -2154,6 +2219,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/61/14/33a3a1352cfa71812a3a21e8c9bfb83f60b0011f5e36f2b1399d51928209/uvicorn-0.34.0-py3-none-any.whl", hash = "sha256:023dc038422502fa28a09c7a30bf2b6991512da7dcdb8fd35fe57cfc154126f4", size = 62315, upload-time = "2024-12-15T13:33:27.467Z" }, ] +[[package]] +name = "volcengine" +version = "1.0.191" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "google" }, + { name = "protobuf" }, + { name = "pycryptodome" }, + { name = "pytz" }, + { name = "requests" }, + { name = "retry" }, + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f4/d8/0ea9b18f216808af709306084d10369f712b98cb5381381d44115dfa6536/volcengine-1.0.191.tar.gz", hash = "sha256:cf3c2dc118c92a7a47f1ab8a48f4789d47e84d17778a1717e1afe9cbce90c986", size = 356076, upload-time = "2025-06-26T12:25:16.353Z" } + [[package]] name = "watchfiles" version = "1.0.5"