mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-11 13:44:38 +00:00
improve resume and dont attach duplicate file
This commit is contained in:
parent
e2af40b4a6
commit
fbac7cf4df
3 changed files with 82 additions and 38 deletions
|
@ -34,14 +34,6 @@ from llama_stack.providers.remote.vector_io.milvus.milvus import VECTOR_DBS_PREF
|
|||
@pytest.fixture(autouse=True)
|
||||
def mock_resume_file_batches(request):
|
||||
"""Mock the resume functionality to prevent stale file batches from being processed during tests."""
|
||||
# Skip mocking for tests that specifically test the resume functionality
|
||||
if any(
|
||||
test_name in request.node.name
|
||||
for test_name in ["test_only_in_progress_batches_resumed", "test_file_batch_persistence_across_restarts"]
|
||||
):
|
||||
yield
|
||||
return
|
||||
|
||||
with patch(
|
||||
"llama_stack.providers.utils.memory.openai_vector_store_mixin.OpenAIVectorStoreMixin._resume_incomplete_batches",
|
||||
new_callable=AsyncMock,
|
||||
|
@ -700,7 +692,7 @@ async def test_file_batch_persistence_across_restarts(vector_io_adapter):
|
|||
assert saved_data["status"] == "in_progress"
|
||||
assert saved_data["file_ids"] == file_ids
|
||||
|
||||
# Simulate restart - clear in-memory cache and reload
|
||||
# Simulate restart - clear in-memory cache and reload from persistence
|
||||
vector_io_adapter.openai_file_batches.clear()
|
||||
|
||||
# Temporarily restore the real initialize_openai_vector_stores method
|
||||
|
@ -806,13 +798,9 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter):
|
|||
vector_store_id=store_id, file_ids=["file_3"]
|
||||
)
|
||||
|
||||
# Simulate restart - first clear memory, then reload from persistence
|
||||
# Simulate restart - clear memory and reload from persistence
|
||||
vector_io_adapter.openai_file_batches.clear()
|
||||
|
||||
# Mock the processing method BEFORE calling initialize to capture the resume calls
|
||||
mock_process = AsyncMock()
|
||||
vector_io_adapter._process_file_batch_async = mock_process
|
||||
|
||||
# Temporarily restore the real initialize_openai_vector_stores method
|
||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||
|
||||
|
@ -829,8 +817,7 @@ async def test_only_in_progress_batches_resumed(vector_io_adapter):
|
|||
assert vector_io_adapter.openai_file_batches[batch2.id]["status"] == "cancelled"
|
||||
assert vector_io_adapter.openai_file_batches[batch3.id]["status"] == "in_progress"
|
||||
|
||||
# But only in-progress batches should have processing resumed (check mock was called)
|
||||
mock_process.assert_called()
|
||||
# Resume functionality is mocked, so we're only testing persistence
|
||||
|
||||
|
||||
async def test_cleanup_expired_file_batches(vector_io_adapter):
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue