From 61398cf565654f2e2d2da5247af488ec098b2d2e Mon Sep 17 00:00:00 2001 From: thepetk Date: Fri, 6 Feb 2026 15:50:37 +0000 Subject: [PATCH 1/3] Ensure unique documents in vector stores --- src/exceptions.py | 8 ++ src/ingest.py | 169 +++++++++++++++++++++++++++++++++---------- src/utils.py | 5 ++ src/workflow.py | 5 +- streamlit_app.py | 10 +-- tests/test_ingest.py | 3 +- 6 files changed, 151 insertions(+), 49 deletions(-) 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..dbd062c 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, @@ -579,64 +580,154 @@ def _insert_single_document( ) return False + async def _get_existing_filenames(self, vector_store_id: "str") -> "set[str]": + """ + retrieves the set of filenames already present in a vector store + and normalizes them for comparison with documents to be inserted + """ + existing_filenames: "set[str]" = set() + try: + # get all file IDs in the vector store + store_files = await asyncio.to_thread( + self.client.vector_stores.files.list, + vector_store_id=vector_store_id, + ) + store_file_ids: "set[str]" = set() + for file_info in store_files or []: + file_id = getattr(file_info, "id", None) or getattr( + file_info, "file_id", None + ) + if file_id: + store_file_ids.add(file_id) + + if not store_file_ids: + return existing_filenames + + # list all files to get filenames + all_files = await asyncio.to_thread(self.client.files.list) + for f in all_files or []: + file_id = getattr(f, "id", None) + if file_id and file_id in store_file_ids: + filename = getattr(f, "filename", None) + if filename: + base = os.path.splitext(str(filename))[0].lower() + existing_filenames.add(base) + except Exception as e: + logger.debug( + f"Could not list files for vector store {vector_store_id}: {e}" + ) + return existing_filenames + + def _normalize_doc_filename(self, doc: "File") -> "str": + """ + returns the normalized filename for a document. + """ + meta = self.file_metadata.get(doc.id, {}) + original = meta.get("original_filename", "") + if original: + return os.path.splitext(original)[0].lower() + return "" + async def create_vector_db( self, vector_store_name: "str", documents: "list[File]" ) -> "bool": """ creates vector database and inserts documents. uses async (concurrent) or sync (sequential) mode based on config. + + before inserting any docs, it checks whether each document already + exists in the vector store (by filename). """ if not documents: logger.warning(f"No documents to insert for {vector_store_name}") return False - logger.info(f"Creating vector database: {vector_store_name}") - + # check if a vector store with this name already exists vector_store = None + existing_filenames: "set[str]" = set() try: - vector_store = await asyncio.to_thread( - self.client.vector_stores.create, name=vector_store_name - ) - self.vector_store_ids.append(vector_store.id) + all_stores = await asyncio.to_thread(self.client.vector_stores.list) + for vs in all_stores or []: + if vs.name == vector_store_name: + vector_store = vs + existing_filenames = await self._get_existing_filenames(vs.id) + break except Exception as e: - error_msg = str(e) - if "already exists" in error_msg.lower(): + logger.debug(f"Could not check existing vector stores: {e}") + + # filter out documents whose filenames are already in the store + if existing_filenames: + new_documents = [] + for doc in documents: + doc_base = self._normalize_doc_filename(doc) + if doc_base and doc_base in existing_filenames: + logger.info( + f"Document '{ + self.file_metadata.get(doc.id, {}).get( + 'original_filename', doc.id + ) + }' " + f"already exists in '{vector_store_name}', " + f"skipping (declined duplicate)" + ) + else: + new_documents.append(doc) + + if not new_documents: logger.info( - f"Vector DB '{vector_store_name}' already exists, continuing..." + f"All documents already exist in '{vector_store_name}', " + f"skipping ingestion" ) + # even if we skip ingestion, we still + # have to add the vector store ID + # to vector_store_ids + if vector_store: + self.vector_store_ids.append(vector_store.id) + return True + documents = new_documents - vector_stores = await asyncio.to_thread(self.client.vector_stores.list) - for vs in vector_stores or []: - if vs.name == vector_store_name: - vector_store = vs - break - - if vector_store is None: - logger.error( - f"Could not find existing vector store '{vector_store_name}'" + # create vector store if it doesn't exist yet + if vector_store is None: + logger.info(f"Creating vector database: {vector_store_name}") + try: + vector_store = await asyncio.to_thread( + self.client.vector_stores.create, name=vector_store_name + ) + self.vector_store_ids.append(vector_store.id) + except Exception as e: + error_msg = str(e) + if "already exists" in error_msg.lower(): + logger.info( + f"Vector DB '{vector_store_name}' already exists " + f"(created by another instance), continuing..." ) - return False - - # check if the store already has files — skip insertion to - # avoid duplicate documents from parallel sessions - try: - existing_files = await asyncio.to_thread( - self.client.vector_stores.files.list, - vector_store_id=vector_store.id, + # find the existing store + vector_stores = await asyncio.to_thread( + self.client.vector_stores.list ) - if existing_files and len(list(existing_files)) > 0: - logger.info( - f"Vector store '{vector_store_name}' already has files, " - f"skipping document insertion" + for vs in vector_stores or []: + if vs.name == vector_store_name: + vector_store = vs + break + if vector_store is None: + logger.error( + f"Could not find existing vector store " + f"'{vector_store_name}'" ) - return True - except Exception as list_err: - logger.debug( - f"Could not list files for '{vector_store_name}': {list_err}" + return False + self.vector_store_ids.append(vector_store.id) + else: + logger.error( + f"Failed to register vector DB '{vector_store_name}': {e}" ) - else: - logger.error(f"Failed to register vector DB '{vector_store_name}': {e}") - return False + return False + else: + if vector_store.id not in self.vector_store_ids: + self.vector_store_ids.append(vector_store.id) + + if vector_store is None: + logger.error(f"Vector store '{vector_store_name}' is unexpectedly None") + return False try: if self.ingestion_mode == "async": @@ -833,7 +924,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..b29700c 100644 --- a/src/utils.py +++ b/src/utils.py @@ -31,6 +31,11 @@ def __setitem__(self, key, value): submission_states: "ObservableDict[str, WorkflowState]" = ObservableDict() +# lock shared across all Streamlit sessions in this process, +# shared through the sys.modules accross Streamlit reruns. +# 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 From 745f2acff250bf9cd18ef6e4a79c7a1bfe39a175 Mon Sep 17 00:00:00 2001 From: thepetk Date: Fri, 6 Feb 2026 16:08:44 +0000 Subject: [PATCH 2/3] Revert changes in ingest.py --- src/ingest.py | 164 ++++++++++++-------------------------------------- 1 file changed, 37 insertions(+), 127 deletions(-) diff --git a/src/ingest.py b/src/ingest.py index dbd062c..dc949cf 100644 --- a/src/ingest.py +++ b/src/ingest.py @@ -580,154 +580,64 @@ def _insert_single_document( ) return False - async def _get_existing_filenames(self, vector_store_id: "str") -> "set[str]": - """ - retrieves the set of filenames already present in a vector store - and normalizes them for comparison with documents to be inserted - """ - existing_filenames: "set[str]" = set() - try: - # get all file IDs in the vector store - store_files = await asyncio.to_thread( - self.client.vector_stores.files.list, - vector_store_id=vector_store_id, - ) - store_file_ids: "set[str]" = set() - for file_info in store_files or []: - file_id = getattr(file_info, "id", None) or getattr( - file_info, "file_id", None - ) - if file_id: - store_file_ids.add(file_id) - - if not store_file_ids: - return existing_filenames - - # list all files to get filenames - all_files = await asyncio.to_thread(self.client.files.list) - for f in all_files or []: - file_id = getattr(f, "id", None) - if file_id and file_id in store_file_ids: - filename = getattr(f, "filename", None) - if filename: - base = os.path.splitext(str(filename))[0].lower() - existing_filenames.add(base) - except Exception as e: - logger.debug( - f"Could not list files for vector store {vector_store_id}: {e}" - ) - return existing_filenames - - def _normalize_doc_filename(self, doc: "File") -> "str": - """ - returns the normalized filename for a document. - """ - meta = self.file_metadata.get(doc.id, {}) - original = meta.get("original_filename", "") - if original: - return os.path.splitext(original)[0].lower() - return "" - async def create_vector_db( self, vector_store_name: "str", documents: "list[File]" ) -> "bool": """ creates vector database and inserts documents. uses async (concurrent) or sync (sequential) mode based on config. - - before inserting any docs, it checks whether each document already - exists in the vector store (by filename). """ if not documents: logger.warning(f"No documents to insert for {vector_store_name}") return False - # check if a vector store with this name already exists + logger.info(f"Creating vector database: {vector_store_name}") + vector_store = None - existing_filenames: "set[str]" = set() try: - all_stores = await asyncio.to_thread(self.client.vector_stores.list) - for vs in all_stores or []: - if vs.name == vector_store_name: - vector_store = vs - existing_filenames = await self._get_existing_filenames(vs.id) - break + vector_store = await asyncio.to_thread( + self.client.vector_stores.create, name=vector_store_name + ) + self.vector_store_ids.append(vector_store.id) except Exception as e: - logger.debug(f"Could not check existing vector stores: {e}") - - # filter out documents whose filenames are already in the store - if existing_filenames: - new_documents = [] - for doc in documents: - doc_base = self._normalize_doc_filename(doc) - if doc_base and doc_base in existing_filenames: - logger.info( - f"Document '{ - self.file_metadata.get(doc.id, {}).get( - 'original_filename', doc.id - ) - }' " - f"already exists in '{vector_store_name}', " - f"skipping (declined duplicate)" - ) - else: - new_documents.append(doc) - - if not new_documents: + error_msg = str(e) + if "already exists" in error_msg.lower(): logger.info( - f"All documents already exist in '{vector_store_name}', " - f"skipping ingestion" + f"Vector DB '{vector_store_name}' already exists, continuing..." ) - # even if we skip ingestion, we still - # have to add the vector store ID - # to vector_store_ids - if vector_store: - self.vector_store_ids.append(vector_store.id) - return True - documents = new_documents - # create vector store if it doesn't exist yet - if vector_store is None: - logger.info(f"Creating vector database: {vector_store_name}") - try: - vector_store = await asyncio.to_thread( - self.client.vector_stores.create, name=vector_store_name - ) - self.vector_store_ids.append(vector_store.id) - except Exception as e: - error_msg = str(e) - if "already exists" in error_msg.lower(): - logger.info( - f"Vector DB '{vector_store_name}' already exists " - f"(created by another instance), continuing..." - ) - # find the existing store - vector_stores = await asyncio.to_thread( - self.client.vector_stores.list - ) - for vs in vector_stores or []: - if vs.name == vector_store_name: - vector_store = vs - break - if vector_store is None: - logger.error( - f"Could not find existing vector store " - f"'{vector_store_name}'" - ) - return False - self.vector_store_ids.append(vector_store.id) - else: + vector_stores = await asyncio.to_thread(self.client.vector_stores.list) + for vs in vector_stores or []: + if vs.name == vector_store_name: + vector_store = vs + break + + if vector_store is None: logger.error( - f"Failed to register vector DB '{vector_store_name}': {e}" + f"Could not find existing vector store '{vector_store_name}'" ) return False - else: - if vector_store.id not in self.vector_store_ids: - self.vector_store_ids.append(vector_store.id) - if vector_store is None: - logger.error(f"Vector store '{vector_store_name}' is unexpectedly None") - return False + # check if the store already has files — skip insertion to + # avoid duplicate documents from parallel sessions + try: + existing_files = await asyncio.to_thread( + self.client.vector_stores.files.list, + vector_store_id=vector_store.id, + ) + if existing_files and len(list(existing_files)) > 0: + logger.info( + f"Vector store '{vector_store_name}' already has files, " + f"skipping document insertion" + ) + return True + except Exception as list_err: + logger.debug( + f"Could not list files for '{vector_store_name}': {list_err}" + ) + else: + logger.error(f"Failed to register vector DB '{vector_store_name}': {e}") + return False try: if self.ingestion_mode == "async": From d3f91a7980c8043d32848c3ad9db29da64b2d7f0 Mon Sep 17 00:00:00 2001 From: thepetk Date: Mon, 9 Feb 2026 10:23:02 +0000 Subject: [PATCH 3/3] Add further comments --- src/ingest.py | 18 ++++++++++++++++-- src/utils.py | 16 ++++++++++++++-- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/src/ingest.py b/src/ingest.py index dc949cf..0d56278 100644 --- a/src/ingest.py +++ b/src/ingest.py @@ -586,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}") @@ -595,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 ) @@ -606,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: @@ -618,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, @@ -639,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( diff --git a/src/utils.py b/src/utils.py index b29700c..9757afa 100644 --- a/src/utils.py +++ b/src/utils.py @@ -31,8 +31,20 @@ def __setitem__(self, key, value): submission_states: "ObservableDict[str, WorkflowState]" = ObservableDict() -# lock shared across all Streamlit sessions in this process, -# shared through the sys.modules accross Streamlit reruns. +# 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()