48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
import pytest
|
|
|
|
from app.core.errors import ForbiddenAgentError
|
|
from app.infrastructure.milvus_knowledge_adapter import MilvusKnowledgeClient
|
|
|
|
|
|
class FakeMilvus:
|
|
def __init__(self) -> None:
|
|
self.kwargs = None
|
|
|
|
async def search(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
return [[{
|
|
"distance": 0.91,
|
|
"entity": {
|
|
"knowledge_id": "101",
|
|
"snippet": "开户说明",
|
|
"title": "基金开户",
|
|
"tags": ["开户"],
|
|
"version": "v1",
|
|
},
|
|
}]]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_knowledge_adapter_uses_cosine_and_minimal_public_projection() -> None:
|
|
client = MilvusKnowledgeClient("http://unused")
|
|
fake = FakeMilvus()
|
|
client._client = fake
|
|
|
|
hits = await client.search("fin_faq_collection", [0.1] * 1024, 3)
|
|
|
|
assert hits[0]["knowledge_id"] == "101"
|
|
assert hits[0]["snippet"] == "开户说明"
|
|
assert hits[0]["score"] == 0.91
|
|
assert fake.kwargs["collection_name"] == "fin_faq_collection"
|
|
assert fake.kwargs["limit"] == 3
|
|
assert fake.kwargs["search_params"] == {"metric_type": "COSINE"}
|
|
assert fake.kwargs["output_fields"] == ["knowledge_id", "title", "snippet", "tags", "version"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_knowledge_adapter_rejects_non_public_collection() -> None:
|
|
client = MilvusKnowledgeClient("http://unused")
|
|
|
|
with pytest.raises(ForbiddenAgentError):
|
|
await client.search("customer_vectors", [0.1] * 1024, 3)
|