diff --git a/src/exceptions.py b/src/exceptions.py index c6c5e76..e86ee74 100644 --- a/src/exceptions.py +++ b/src/exceptions.py @@ -13,3 +13,11 @@ class NoVectorStoresFoundError(Exception): """ pass + + +class IngestionPipelineError(Exception): + """ + Raised when all ingestion pipelines fail. + """ + + pass diff --git a/src/ingest.py b/src/ingest.py index c07496e..0d56278 100644 --- a/src/ingest.py +++ b/src/ingest.py @@ -28,6 +28,7 @@ DEFAULT_LLAMA_STACK_RETRY_DELAY, DEFAULT_LLAMA_STACK_WAITING_RETRIES, ) +from src.exceptions import IngestionPipelineError from src.types import ( RAW_PIPELINES_TYPE, Pipeline, @@ -585,6 +586,9 @@ async def create_vector_db( """ creates vector database and inserts documents. uses async (concurrent) or sync (sequential) mode based on config. + + all (LlamaStack) client calls are wrapped in thread (asyncio.to_thread) + to avoid blocking the concurrent execution. """ if not documents: logger.warning(f"No documents to insert for {vector_store_name}") @@ -594,6 +598,9 @@ async def create_vector_db( vector_store = None try: + # asyncio.to_thread moves the synchronous llama-stack-client + # call to a separate thread so the asyncio (event) loop keeps + # running. vector_store = await asyncio.to_thread( self.client.vector_stores.create, name=vector_store_name ) @@ -605,6 +612,8 @@ async def create_vector_db( f"Vector DB '{vector_store_name}' already exists, continuing..." ) + # vector store exists, find it by name to get its ID + # to continue ingestion vector_stores = await asyncio.to_thread(self.client.vector_stores.list) for vs in vector_stores or []: if vs.name == vector_store_name: @@ -617,8 +626,10 @@ async def create_vector_db( ) return False - # check if the store already has files — skip insertion to - # avoid duplicate documents from parallel sessions + # if the existing store already contains files, we skip + # ingestion entirely. Simple rule to prevent duplicate documents + # when multiple parallel sessions race to populate the same store. + # TODO: investigate more robust approaches try: existing_files = await asyncio.to_thread( self.client.vector_stores.files.list, @@ -638,6 +649,10 @@ async def create_vector_db( logger.error(f"Failed to register vector DB '{vector_store_name}': {e}") return False + # ingest documents into the vector store. + # - async mode: all documents are inserted concurrently using + # asyncio.gather (each ingestion action runs in its own thread). + # - sync mode: documents are inserted one at a time sequentially. try: if self.ingestion_mode == "async": logger.info( @@ -833,7 +848,9 @@ async def process_with_error_handling(pipeline: Pipeline) -> bool: if successful == 0: logger.warning("all pipeline(s) failed. Check logs for details.") - sys.exit(1) + raise IngestionPipelineError( + "All ingestion pipelines failed. Check logs for details." + ) elif failed > 0: logger.warning(f"{failed} pipeline(s) failed. Check logs for details.") else: diff --git a/src/utils.py b/src/utils.py index 6948634..9757afa 100644 --- a/src/utils.py +++ b/src/utils.py @@ -31,6 +31,23 @@ def __setitem__(self, key, value): submission_states: "ObservableDict[str, WorkflowState]" = ObservableDict() +# lock/mutex shared across all Streamlit sessions in this process. +# +# Fix: commit 61398cf565654f2e2d2da5247af488ec098b2d2e +# +# the lock was moved in utils.py (from streamlit_app.py) in order to +# be cached similar to all imported modules (sys.modules) of Python: +# - https://docs.python.org/3/reference/import.html#the-module-cache +# +# this allows us to keep the same lock between different streamlit +# sessions (of the same streamlit process). With the previous approach, +# (lock object was located in streamlit_app.py), each streamlit session +# was generating again the lock object (since streamlit_app.py was +# re-run for each session) + +# see: https://docs.streamlit.io/develop/concepts/design/multithreading +ingestion_lock = threading.Lock() + def clean_text(text: "str") -> "str": """ diff --git a/src/workflow.py b/src/workflow.py index 56127b8..27df3c1 100644 --- a/src/workflow.py +++ b/src/workflow.py @@ -71,7 +71,8 @@ def _call_openai_llm(self, state: "WorkflowState") -> "str": messages = self._convert_messages_to_openai_format(state) completion = self.rag_service.openai_client.chat.completions.create( - model=INFERENCE_MODEL, messages=messages + model=INFERENCE_MODEL, + messages=messages, # type: ignore[invalid-argument-type] ) return completion.choices[0].message.content or "" @@ -205,7 +206,7 @@ def llm_node(state: "WorkflowState") -> "WorkflowState": rag_response = client_to_use.responses.create( model=INFERENCE_MODEL, input=rag_prompt, - tools=[file_search_tool], + tools=[file_search_tool], # type: ignore[invalid-argument-type] ) rag_end_time = time.time() state["rag_query_time"] = rag_end_time - rag_start_time diff --git a/streamlit_app.py b/streamlit_app.py index af26206..f441169 100644 --- a/streamlit_app.py +++ b/streamlit_app.py @@ -1,7 +1,6 @@ import asyncio import json import os -import threading import time import uuid from concurrent.futures import Future, ThreadPoolExecutor @@ -24,17 +23,12 @@ from src.types import Pipeline, WorkflowState from src.utils import ( check_llama_stack_availability, + ingestion_lock, logger, submission_states, ) from src.workflow import Workflow -# lock shared across all Streamlit sessions in this process. Streamlit runs -# each session in a separate thread but within the same process, so this -# lock will handle parallel ingestion attempts. -# see: https://docs.streamlit.io/develop/concepts/design/multithreading -_ingestion_lock = threading.Lock() - # API_KEY: OpenAI API key (not used directly but may be needed API_KEY = os.getenv("OPENAI_API_KEY", "not applicable") @@ -930,7 +924,7 @@ def main() -> "None": # vector stores. All other sessions will get the lock only after # that's completed, check again (stores now exist), and finally # will skip ingestion. - with _ingestion_lock: + with ingestion_lock: loop = get_or_create_event_loop() loop.run_until_complete(check_and_run_ingestion_if_needed()) st.rerun() diff --git a/tests/test_ingest.py b/tests/test_ingest.py index 733ccfd..115b367 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -5,6 +5,7 @@ import pytest import yaml +from src.exceptions import IngestionPipelineError from src.ingest import IngestionService from src.types import Pipeline, SourceConfig, SourceTypes @@ -591,7 +592,7 @@ async def test_run_with_failed_pipelines( with patch.object(service, "process_pipeline", return_value=False): with patch.object(service, "save_file_metadata"): - with pytest.raises(SystemExit): + with pytest.raises(IngestionPipelineError): await service.run() @pytest.mark.anyio