mirror of
https://github.com/langgenius/dify.git
synced 2026-06-18 07:41:09 +08:00
249 lines
9.2 KiB
Python
249 lines
9.2 KiB
Python
import json
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from dify_agent.layers.knowledge.client import DifyKnowledgeBaseClient, DifyKnowledgeBaseClientError
|
|
from dify_agent.layers.knowledge.configs import (
|
|
DifyKnowledgeMetadataFilteringConfig,
|
|
DifyKnowledgeRetrievalConfig,
|
|
)
|
|
|
|
|
|
def _retrieval_config() -> DifyKnowledgeRetrievalConfig:
|
|
return DifyKnowledgeRetrievalConfig(mode="multiple", top_k=4, score_threshold=0.2)
|
|
|
|
|
|
def _metadata_filtering() -> DifyKnowledgeMetadataFilteringConfig:
|
|
return DifyKnowledgeMetadataFilteringConfig(mode="disabled")
|
|
|
|
|
|
def test_knowledge_client_posts_inner_api_request_with_static_controls() -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
assert str(request.url) == "http://dify-api/inner/api/knowledge/retrieve"
|
|
assert request.headers["X-Inner-Api-Key"] == "inner-secret"
|
|
payload = json.loads(request.content.decode("utf-8"))
|
|
assert payload == {
|
|
"caller": {
|
|
"tenant_id": "tenant-1",
|
|
"user_id": "user-1",
|
|
"app_id": "app-1",
|
|
"user_from": "account",
|
|
"invoke_from": "agent_app",
|
|
},
|
|
"dataset_ids": ["dataset-1"],
|
|
"query": "reset password",
|
|
"retrieval": {
|
|
"mode": "multiple",
|
|
"top_k": 4,
|
|
"score_threshold": 0.2,
|
|
"reranking_mode": "reranking_model",
|
|
"reranking_enable": True,
|
|
"reranking_model": None,
|
|
"weights": None,
|
|
},
|
|
"metadata_filtering": {"mode": "disabled"},
|
|
"attachment_ids": [],
|
|
}
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"results": [
|
|
{
|
|
"metadata": {
|
|
"_source": "knowledge",
|
|
"dataset_name": "Docs",
|
|
"document_name": "FAQ.md",
|
|
"score": 0.9,
|
|
},
|
|
"title": "FAQ",
|
|
"files": [],
|
|
"content": "Use the reset link.",
|
|
"summary": None,
|
|
}
|
|
],
|
|
"usage": {},
|
|
},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api",
|
|
api_key="inner-secret",
|
|
http_client=http_client,
|
|
)
|
|
response = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert response.results[0].metadata.dataset_name == "Docs"
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_knowledge_client_marks_retryable_http_failures() -> None:
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(
|
|
502, json={"code": "external_knowledge_failed", "message": "bad gateway"}
|
|
)
|
|
)
|
|
) as http_client:
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api", api_key="inner-secret", http_client=http_client
|
|
)
|
|
with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
|
|
_ = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert exc_info.value.status_code == 502
|
|
assert exc_info.value.error_code == "external_knowledge_failed"
|
|
assert exc_info.value.retryable is True
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_knowledge_client_marks_non_retryable_http_failures() -> None:
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(
|
|
lambda _request: httpx.Response(
|
|
403,
|
|
json={"code": "dataset_tenant_mismatch", "message": "forbidden"},
|
|
)
|
|
)
|
|
) as http_client:
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api", api_key="inner-secret", http_client=http_client
|
|
)
|
|
with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
|
|
_ = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert exc_info.value.status_code == 403
|
|
assert exc_info.value.error_code == "dataset_tenant_mismatch"
|
|
assert exc_info.value.retryable is False
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_knowledge_client_rejects_malformed_success_response() -> None:
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(lambda _request: httpx.Response(200, json={"bad": []}))
|
|
) as http_client:
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api", api_key="inner-secret", http_client=http_client
|
|
)
|
|
with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
|
|
_ = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert exc_info.value.error_code == "invalid_response"
|
|
assert exc_info.value.retryable is False
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"error_factory",
|
|
[
|
|
lambda request: httpx.ReadTimeout("timed out", request=request),
|
|
lambda request: httpx.ConnectError("connection failed", request=request),
|
|
],
|
|
)
|
|
def test_knowledge_client_marks_transport_failures_retryable(error_factory) -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
raise error_factory(request)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api", api_key="inner-secret", http_client=http_client
|
|
)
|
|
with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
|
|
_ = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert exc_info.value.retryable is True
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_knowledge_client_treats_invalid_url_errors_as_non_retryable_configuration_error() -> None:
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient() as http_client:
|
|
http_client.post = AsyncMock(side_effect=httpx.UnsupportedProtocol("unsupported protocol"))
|
|
client = DifyKnowledgeBaseClient(
|
|
base_url="http://dify-api", api_key="inner-secret", http_client=http_client
|
|
)
|
|
with pytest.raises(DifyKnowledgeBaseClientError) as exc_info:
|
|
_ = await client.retrieve(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
app_id="app-1",
|
|
user_from="account",
|
|
invoke_from="agent_app",
|
|
dataset_ids=["dataset-1"],
|
|
query="reset password",
|
|
retrieval=_retrieval_config(),
|
|
metadata_filtering=_metadata_filtering(),
|
|
)
|
|
assert exc_info.value.retryable is False
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(scenario())
|