ruff
This commit is contained in:
parent
c5c60ccad0
commit
0d0fe3aa30
1 changed files with 35 additions and 9 deletions
|
|
@ -616,9 +616,15 @@ async def test_extract_time_from_query_relative_path():
|
||||||
with (
|
with (
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
|
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
|
patch(
|
||||||
|
"cognee.modules.retrieval.temporal_retriever.render_prompt",
|
||||||
|
return_value="System prompt",
|
||||||
|
),
|
||||||
patch.object(
|
patch.object(
|
||||||
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
|
LLMGateway,
|
||||||
|
"acreate_structured_output",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=mock_interval,
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
||||||
|
|
@ -632,7 +638,9 @@ async def test_extract_time_from_query_relative_path():
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_extract_time_from_query_absolute_path():
|
async def test_extract_time_from_query_absolute_path():
|
||||||
"""Test extract_time_from_query with absolute prompt path."""
|
"""Test extract_time_from_query with absolute prompt path."""
|
||||||
retriever = TemporalRetriever(time_extraction_prompt_path="/absolute/path/to/extract_query_time.txt")
|
retriever = TemporalRetriever(
|
||||||
|
time_extraction_prompt_path="/absolute/path/to/extract_query_time.txt"
|
||||||
|
)
|
||||||
|
|
||||||
mock_timestamp_from = Timestamp(year=2024, month=1, day=1)
|
mock_timestamp_from = Timestamp(year=2024, month=1, day=1)
|
||||||
mock_timestamp_to = Timestamp(year=2024, month=12, day=31)
|
mock_timestamp_to = Timestamp(year=2024, month=12, day=31)
|
||||||
|
|
@ -640,12 +648,24 @@ async def test_extract_time_from_query_absolute_path():
|
||||||
|
|
||||||
with (
|
with (
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=True),
|
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=True),
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.os.path.dirname", return_value="/absolute/path/to"),
|
patch(
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.os.path.basename", return_value="extract_query_time.txt"),
|
"cognee.modules.retrieval.temporal_retriever.os.path.dirname",
|
||||||
|
return_value="/absolute/path/to",
|
||||||
|
),
|
||||||
|
patch(
|
||||||
|
"cognee.modules.retrieval.temporal_retriever.os.path.basename",
|
||||||
|
return_value="extract_query_time.txt",
|
||||||
|
),
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
|
patch(
|
||||||
|
"cognee.modules.retrieval.temporal_retriever.render_prompt",
|
||||||
|
return_value="System prompt",
|
||||||
|
),
|
||||||
patch.object(
|
patch.object(
|
||||||
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
|
LLMGateway,
|
||||||
|
"acreate_structured_output",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=mock_interval,
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
||||||
|
|
@ -666,9 +686,15 @@ async def test_extract_time_from_query_with_none_values():
|
||||||
with (
|
with (
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
|
patch("cognee.modules.retrieval.temporal_retriever.os.path.isabs", return_value=False),
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
patch("cognee.modules.retrieval.temporal_retriever.datetime") as mock_datetime,
|
||||||
patch("cognee.modules.retrieval.temporal_retriever.render_prompt", return_value="System prompt"),
|
patch(
|
||||||
|
"cognee.modules.retrieval.temporal_retriever.render_prompt",
|
||||||
|
return_value="System prompt",
|
||||||
|
),
|
||||||
patch.object(
|
patch.object(
|
||||||
LLMGateway, "acreate_structured_output", new_callable=AsyncMock, return_value=mock_interval
|
LLMGateway,
|
||||||
|
"acreate_structured_output",
|
||||||
|
new_callable=AsyncMock,
|
||||||
|
return_value=mock_interval,
|
||||||
),
|
),
|
||||||
):
|
):
|
||||||
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
mock_datetime.now.return_value.strftime.return_value = "11-12-2024"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue