forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
from typing import TYPE_CHECKING, List
|
|
|
|
import pytest
|
|
from pytest_mock import MockerFixture
|
|
|
|
from langchain.retrievers import ZepRetriever
|
|
from langchain.schema import Document
|
|
|
|
if TYPE_CHECKING:
|
|
from zep_python import MemorySearchResult, ZepClient
|
|
|
|
|
|
@pytest.fixture
|
|
def search_results() -> List[MemorySearchResult]:
|
|
from zep_python import MemorySearchResult, Message
|
|
|
|
search_result = [
|
|
{
|
|
"message": {
|
|
"uuid": "66830914-19f5-490b-8677-1ba06bcd556b",
|
|
"created_at": "2023-05-18T20:40:42.743773Z",
|
|
"role": "user",
|
|
"content": "I'm looking to plan a trip to Iceland. Can you help me?",
|
|
"token_count": 17,
|
|
},
|
|
"summary": None,
|
|
"dist": 0.8734284910450115,
|
|
},
|
|
{
|
|
"message": {
|
|
"uuid": "015e618c-ba9d-45b6-95c3-77a8e611570b",
|
|
"created_at": "2023-05-18T20:40:42.743773Z",
|
|
"role": "user",
|
|
"content": "How much does a trip to Iceland typically cost?",
|
|
"token_count": 12,
|
|
},
|
|
"summary": None,
|
|
"dist": 0.8554048017463456,
|
|
},
|
|
]
|
|
|
|
return [
|
|
MemorySearchResult(
|
|
message=Message.parse_obj(result["message"]),
|
|
summary=result["summary"],
|
|
dist=result["dist"],
|
|
)
|
|
for result in search_result
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
@pytest.mark.requires("zep_python")
|
|
def zep_retriever(
|
|
mocker: MockerFixture, search_results: List[MemorySearchResult]
|
|
) -> ZepRetriever:
|
|
mock_zep_client: ZepClient = mocker.patch("zep_python.ZepClient", autospec=True)
|
|
mock_zep_client.search_memory.return_value = copy.deepcopy( # type: ignore
|
|
search_results
|
|
)
|
|
mock_zep_client.asearch_memory.return_value = copy.deepcopy( # type: ignore
|
|
search_results
|
|
)
|
|
zep = ZepRetriever(session_id="123", url="http://localhost:8000")
|
|
zep.zep_client = mock_zep_client
|
|
return zep
|
|
|
|
|
|
@pytest.mark.requires("zep_python")
|
|
def test_zep_retriever_get_relevant_documents(
|
|
zep_retriever: ZepRetriever, search_results: List[MemorySearchResult]
|
|
) -> None:
|
|
documents: List[Document] = zep_retriever.get_relevant_documents(
|
|
query="My trip to Iceland"
|
|
)
|
|
_test_documents(documents, search_results)
|
|
|
|
|
|
@pytest.mark.requires("zep_python")
|
|
@pytest.mark.asyncio
|
|
async def test_zep_retriever_aget_relevant_documents(
|
|
zep_retriever: ZepRetriever, search_results: List[MemorySearchResult]
|
|
) -> None:
|
|
documents: List[Document] = await zep_retriever.aget_relevant_documents(
|
|
query="My trip to Iceland"
|
|
)
|
|
_test_documents(documents, search_results)
|
|
|
|
|
|
def _test_documents(
|
|
documents: List[Document], search_results: List[MemorySearchResult]
|
|
) -> None:
|
|
assert len(documents) == 2
|
|
for i, document in enumerate(documents):
|
|
assert document.page_content == search_results[i].message.get( # type: ignore
|
|
"content"
|
|
)
|
|
assert document.metadata.get("uuid") == search_results[
|
|
i
|
|
].message.get( # type: ignore
|
|
"uuid"
|
|
)
|
|
assert document.metadata.get("role") == search_results[
|
|
i
|
|
].message.get( # type: ignore
|
|
"role"
|
|
)
|
|
assert document.metadata.get("score") == search_results[i].dist
|