Skip to content

fix(collections): Use correct file batch ID #1073

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vprashrex wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix(collections): Use correct file batch ID #1073

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
aea2dfa
fix(collections): Use file batch ID from create() when listing failed…
vprashrex Jul 23, 2026
6ca053f
fix(tests): Simplify exception handling in OpenAI vector store tests
vprashrex Jul 23, 2026
8c02d57
fix(upload): Replace pre-transform validation with upload validation …
vprashrex Jul 26, 2026
d391ad2
fix(openai): Enhance OpenAI vector store handling with improved batch…
vprashrex Jul 26, 2026
3f6fc81
Merge branch 'main' into fix/openai-vector-store-batch-id
vprashrex Jul 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions backend/app/api/routes/documents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from app.services.documents.helpers import (
calculate_file_size,
schedule_transformation,
pre_transform_validation,
validate_upload,
build_document_schema,
build_document_schemas,
)
Expand Down Expand Up @@ -126,8 +126,8 @@ async def upload_doc(
if callback_url:
validate_callback_url(callback_url)

source_format, actual_transformer = pre_transform_validation(
src_filename=src.filename,
source_format, actual_transformer = validate_upload(
src=src,
target_format=target_format,
transformer=transformer,
)
Expand Down
77 changes: 69 additions & 8 deletions backend/app/crud/rag/open_ai.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import json
import logging
import time
import functools as ft

import openai
from openai import OpenAI, OpenAIError
from openai.types import VectorStore
from openai.types.vector_stores import VectorStoreFileBatch
from pydantic import BaseModel

from app.models import Document, ProviderType
Expand All @@ -12,6 +15,15 @@

OPENAI_PROVIDER = ProviderType.openai.value

OPENAI_MAX_RETRIES = 2

# Per socket operation, not per call: a steady upload runs as long as it needs and
# this only fires on a dead connection. Sized so all attempts fit one Celery task.
OPENAI_TIMEOUT_SECONDS = 90

BATCH_POLL_INTERVAL_SECONDS = 2
BATCH_POLL_TIMEOUT_SECONDS = 240


def vs_ls(client: OpenAI, vector_store_id: str):
kwargs = {}
Expand Down Expand Up @@ -85,7 +97,7 @@ def __init__(self, client):


class OpenAIVectorStoreCrud(OpenAICrud):
def create(self):
def create(self) -> VectorStore:
logger.info(
f"[OpenAIVectorStoreCrud.create] Creating vector store | {{'action': 'create'}}"
)
Expand All @@ -101,6 +113,52 @@ def read(self, vector_store_id: str):
)
yield from vs_ls(self.client, vector_store_id)

def _create_file_batch(self, vector_store_id: str, file_ids: list[str]) -> str:
"""Returns the vsfb_ id. poll()'s return deserializes a vector-store body,
so its .id is the vs_ id - take the batch id from create()."""
created = self.client.vector_stores.file_batches.create(
vector_store_id=vector_store_id,
file_ids=file_ids,
)
return created.id

def _retrieve_file_batch(
self, batch_id: str, vector_store_id: str
) -> VectorStoreFileBatch:
return self.client.vector_stores.file_batches.retrieve(
batch_id, vector_store_id=vector_store_id
)

def _poll_file_batch(
self, batch_id: str, vector_store_id: str
) -> VectorStoreFileBatch:
"""Wait for indexing to finish. The SDK's own poll() loops forever."""
deadline = time.monotonic() + BATCH_POLL_TIMEOUT_SECONDS

while True:
batch = self._retrieve_file_batch(batch_id, vector_store_id)
if batch.status != "in_progress":
Comment on lines +136 to +140

@coderabbitai coderabbitai Bot Jul 26, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "## locate file"
fd -a 'open_ai\.py$' . || true

echo "## inspect relevant file outline/section"
if [ -f backend/app/crud/rag/open_ai.py ]; then
  wc -l backend/app/crud/rag/open_ai.py
  ast-grep outline backend/app/crud/rag/open_ai.py || true
  echo "## lines 1-240"
  sed -n '1,240p' backend/app/crud/rag/open_ai.py | nl -ba
  echo "## timeout/retry references"
  rg -n "timeout|retries|deadline|BATCH_POLL_TIMEOUT_SECONDS|_retrieve_file_batch|AsyncOpenAI|OpenAI" backend/app/crud/rag/open_ai.py || true
fi

echo "## dependency/version manifests"
for f in requirements.txt requirements-dev.txt pyproject.toml setup.cfg setup.py uv.lock poetry.lock Pipfile.lock; do
  [ -f "$f" ] && { echo "--- $f"; rg -n "openai|httpx" "$f" || true; }
done

echo "## other timeout/retry usages"
rg -n "AsyncOpenAI|OpenAI\(|timeout\s*=|retries\s*=" backend || true

echo "## git diff stat and relevant diff"
git diff --stat || true
git diff -- backend/app/crud/rag/open_ai.py || true

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 1362

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "## lines 1-240"
sed -n '1,240p' backend/app/crud/rag/open_ai.py

echo "## timeout/retry references"
rg -n "timeout|retries|deadline|BATCH_POLL_TIMEOUT_SECONDS|_retrieve_file_batch|AsyncOpenAI|OpenAI" backend/app/crud/rag/open_ai.py || true

echo "## dependency/version manifests"
for f in requirements.txt requirements-dev.txt pyproject.toml setup.cfg setup.py uv.lock poetry.lock Pipfile.lock; do
  [ -f "$f" ] && { echo "--- $f"; rg -n "openai|httpx" "$f" || true; }
done

echo "## other timeout/retry usages"
rg -n "AsyncOpenAI|OpenAI\(|timeout\s*=|retries\s*=" backend || true

echo "## focused tests for poll/retrieval client"
rg -n "OpenAIVectorStoreCrud|_poll_file_batch|_retrieve_file_batch|OPENAI_MAX_RETRIES|OPENAI_TIMEOUT_SECONDS|vector_store_file_batches|retrieve" backend/app/tests backend/app/crud backend/app/services || true

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 35059

🌐 Web query:

OpenAI Python SDK configuration 90 second timeout max_retries=2 default retry behavior

💡 Result:

To configure the OpenAI Python SDK with a 90-second timeout and 2 retries, you can set these parameters when initializing the OpenAI client [1][2]. Configuration Example: from openai import OpenAI import httpx # Configure timeout and retries client = OpenAI( timeout=90.0, # Set total timeout to 90 seconds [1][3] max_retries=2, # Set maximum number of retries [1][3]) Default Behavior and Details: - Timeout: The default timeout is 10 minutes (600 seconds) [4]. Setting timeout=90.0 overrides this global default [1][2]. You can also pass an httpx.Timeout object for more granular control (e.g., specifying connect, read, or write timeouts) [1][5]. - Retries: The default number of retries is already 2 [1][6]. These retries occur automatically for specific error types, including connection errors, 408 Request Timeout, 409 Conflict, 429 Rate Limit, and >=500 Internal Server Errors [1][3][6]. - Per-Request Overrides: If you need to apply different settings to a specific call rather than the entire client, you can use the with_options method [1][7]: response = client.with_options(timeout=90.0, max_retries=2).chat.completions.create(...) Note: Requests that time out are subject to the same retry logic and will be retried up to the configured max_retries limit [1][2][5].

Citations:

Make the batch-poll deadline enforceable.

_retrieve_file_batch() uses the client’s 90s timeout and default retries, but _poll_file_batch() only checks the 240s deadline after _retrieve_file_batch() returns. If the SDK retries timed-out retrievals, polling can exceed the advertised wall-clock budget. Bound each retrieve/retry by the remaining deadline, or disable/limit retries for polling specifically.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/app/crud/rag/open_ai.py` around lines 136 - 140, Update
_poll_file_batch so each _retrieve_file_batch call cannot run past the remaining
BATCH_POLL_TIMEOUT_SECONDS deadline. Compute the remaining time before retrieval
and pass it through the client/request timeout while disabling or limiting
retries for this polling operation; preserve the existing status handling and
stop polling when the deadline is exhausted.

Source: MCP tools

return batch
Comment on lines +140 to +141

@coderabbitai coderabbitai Bot Jul 26, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "Files matching open_ai.py:"
fd -a 'open_ai\.py$' . || true

file="$(fd 'open_ai\.py$' backend/app/crud/rag 2>/dev/null | head -n1 || true)"
echo "Selected file: $file"
if [ -n "${file:-}" ]; then
  wc -l "$file"
  ast-grep outline "$file" --view compact || true
  echo "---- lines 1-240 ----"
  sed -n '1,240p' "$file" | nl -ba
fi

echo "---- related references to status/file_counts/cancelled in file ----"
if [ -n "${file:-}" ]; then
  rg -n 'status|file_counts|cancel|cancelled|fail|in_progress|completed' "$file"
fi

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 747

🌐 Web query:

OpenAI vector store file batches statuses cancelled failed completed in_progress API documentation

💡 Result:

In the OpenAI API, both Vector Store File Batches and individual Vector Store Files share the same four possible statuses: in_progress, completed, cancelled, and failed [1][2][3][4]. When working with Vector Store File Batches, the status reflects the overall state of the batch operation [1][5]. Additionally, these objects include a file_counts attribute that provides a breakdown of the constituent files' statuses, including the number of files currently in progress, completed, failed, and cancelled, as well as the total count [1][6]. When listing files within a specific batch, you can filter the results using these same four status values [3][7]. The status completed specifically indicates that a file is successfully processed and ready for use [3][4][7].

Citations:

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="backend/app/crud/rag/open_ai.py"

echo "---- ast-grep outline ----"
ast-grep outline "$file" --view names || true

echo "---- lines 1-240 ----"
awk '{printf "%5d\t%s\n", NR, $0}' "$file" | sed -n '1,240p'

echo "---- relevant references ----"
grep -nE 'status|file_counts|cancel|cancelled|fail|in_progress|completed|update\(' "$file" || true

test_file="backend/app/tests/crud/rag/test_open_ai.py"
if [ -f "$test_file" ]; then
  echo "---- tests relevant references ----"
  grep -nE 'status|file_counts|cancel|cancelled|fail|in_progress|completed|update\(' "$test_file" || true
fi

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 14994

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="backend/app/crud/rag/open_ai.py"

echo "---- lines 240-335 ----"
awk '{printf "%5d\t%s\n", NR, $0}' "$file" | sed -n '240,335p'

echo "---- file-count-like usages repository-wide ----"
grep -RInE 'file_counts|cancelled|cancelled_files|list_files\(.*filter=' backend/app tests 2>/dev/null || true

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 8174

Handle non-completed terminal batch statuses explicitly.

_poll_file_batch() returns on any status except in_progress, so cancelled and failed batches can reach the successful-log path when completed == 0 and only failed > 0 is checked. Treat only completed as success and raise/fail the upload for terminal non-completed statuses, including cancelled counts where relevant.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/app/crud/rag/open_ai.py` around lines 140 - 141, Update
_poll_file_batch() so only a batch with status "completed" is returned as
successful; explicitly handle terminal statuses such as "failed" and "cancelled"
by raising or otherwise failing the upload, including cancelled counts in the
failure details where applicable. Do not let non-completed terminal batches
reach the successful-log path.

Source: MCP tools


if time.monotonic() >= deadline:
error_message = (
f"[KAAPI] Timed out (code: batch-poll-timeout) after "
f"{BATCH_POLL_TIMEOUT_SECONDS}s waiting for OpenAI to finish indexing "
f"{batch.file_counts.in_progress} file(s). Retry the "
f"collection, or split it into smaller batches."
)
logger.error(
f"[OpenAIVectorStoreCrud._poll_file_batch] {error_message} | "
f"vector_store_id={vector_store_id}, batch_id={batch_id}, "
f"completed={batch.file_counts.completed}, "
f"in_progress={batch.file_counts.in_progress}"
)
raise InterruptedError(error_message)

time.sleep(
min(BATCH_POLL_INTERVAL_SECONDS, max(0.0, deadline - time.monotonic()))
)

def update(
self,
vector_store_id: str,
Expand All @@ -115,11 +173,10 @@ def update(
)

try:
batch = self.client.vector_stores.file_batches.upload_and_poll(
vector_store_id=vector_store_id,
files=[],
file_ids=[doc.file_id[OPENAI_PROVIDER] for doc in docs],
batch_id = self._create_file_batch(
vector_store_id, [doc.file_id[OPENAI_PROVIDER] for doc in docs]
)
batch = self._poll_file_batch(batch_id, vector_store_id)
except openai.RateLimitError as e:
error_message = (
f"[OPENAI] Rate limit exceeded (code: {e.status_code}): "
Expand Down Expand Up @@ -215,14 +272,14 @@ def update(

logger.info(
f"[OpenAIVectorStoreCrud.update] Batch complete | "
f"{{'vector_store_id': '{vector_store_id}', "
f"{{'vector_store_id': '{vector_store_id}', 'batch_id': '{batch_id}', "
f"'completed': {batch.file_counts.completed}, 'failed': {batch.file_counts.failed}}}"
)
if batch.file_counts.failed > 0:
try:
failed_files = self.client.vector_stores.file_batches.list_files(
vector_store_id=vector_store_id,
batch_id=batch.id,
batch_id=batch_id,
filter="failed",
)
doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs}
Expand All @@ -232,11 +289,15 @@ def update(
label = d.fname if d else f.id
msg = f.last_error.message if f.last_error else "no error detail"
parts.append(f"{label}: {msg}")
logger.error(
f"[OpenAIVectorStoreCrud.update] Files failed to index | "
f"{{'batch_id': '{batch_id}', 'failed_files': '{', '.join(parts)}'}}"
)
raise RuntimeError("; ".join(parts))
except OpenAIError as err:
logger.warning(
f"[OpenAIVectorStoreCrud.update] Could not fetch per-file errors | "
f"{{'batch_id': '{batch.id}', 'error': '{str(err)}'}}"
f"{{'batch_id': '{batch_id}', 'error': '{str(err)}'}}"
)
raise

Expand Down
88 changes: 84 additions & 4 deletions backend/app/services/collections/create_collection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from sqlmodel import Session
from gevent import Timeout
from celery.exceptions import SoftTimeLimitExceeded
from celery.exceptions import Retry, SoftTimeLimitExceeded
from opentelemetry import trace
from asgi_correlation_id import correlation_id

Expand Down Expand Up @@ -40,6 +40,48 @@
tracer = trace.get_tracer(__name__)


def _can_retry_batch(task_instance) -> bool:
"""Whether a timed-out batch has attempts left.

Worth retrying because upload_files persists each file ID as it uploads, so
every attempt resumes where the last stopped. `max_retries=None` means retry
forever, which would never fail the job or fire its callback; treat as none.
"""
if task_instance is None or task_instance.max_retries is None:
return False
return task_instance.request.retries < task_instance.max_retries


def _note_batch_retry(
project_id: int,
job_id: str,
batch_number: int,
attempt: int,
max_attempts: int,
) -> None:
"""Record a pending retry on the job row.

The job stays PROCESSING between attempts and pending_monitor only watches
PENDING, so a retry that never lands would be invisible.
"""
try:
with Session(engine) as session:
CollectionJobCrud(session, project_id).update(
UUID(job_id),
CollectionJobUpdate(
error_message=(
f"Batch {batch_number} timed out; "
f"retry {attempt}/{max_attempts} queued"
)
),
)
except Exception:
logger.warning(
"[create_collection._note_batch_retry] Could not record retry | job_id=%s",
job_id,
)


def start_job(
db: Session,
request: CreationRequest,
Expand Down Expand Up @@ -403,9 +445,13 @@ def execute_batch_job(
collection_job_crud = CollectionJobCrud(session, project_id)
collection_job = collection_job_crud.read_one(job_uuid)
already_uploaded = collection_job.documents_uploaded or []
now_uploaded = already_uploaded + [
str(d) for d in all_doc_ids_this_batch
]
# A re-run past this point re-appends its own IDs, and read_each
# rejects a list whose duplicates collapse in its IN.
now_uploaded = list(
dict.fromkeys(
already_uploaded + [str(d) for d in all_doc_ids_this_batch]
)
)

collection_job = collection_job_crud.update(
job_uuid,
Expand Down Expand Up @@ -509,6 +555,40 @@ def execute_batch_job(

except (Timeout, SoftTimeLimitExceeded) as err:
timeout_err = TimeoutError("Task exceeded soft time limit")

if _can_retry_batch(task_instance):
attempt = task_instance.request.retries + 1
logger.warning(
"[create_collection.execute_batch_job] Batch timed out, retrying | "
"{'collection_job_id': '%s', 'batch_number': %d, 'attempt': '%d/%d'}",
job_id,
batch_number,
attempt,
task_instance.max_retries,
)
span.set_attribute("collection.batch_retry", attempt)
_note_batch_retry(
project_id,
job_id,
batch_number,
attempt,
task_instance.max_retries,
)
try:
# Only Retry means it reached the broker. retry() also raises
# Reject on a failed publish, or exc outside a worker - those
# must fail the job, not strand it in PROCESSING.
raise task_instance.retry(exc=timeout_err)
except Retry:
raise
except Exception:
logger.exception(
"[create_collection.execute_batch_job] Retry could not be scheduled, "
"failing the job | {'collection_job_id': '%s', 'batch_number': %d}",
job_id,
batch_number,
)

logger.warning(
"[create_collection.execute_batch_job] Collection Creation Timed Out | {'collection_job_id': '%s', 'error': '%s'}",
job_id,
Expand Down
7 changes: 6 additions & 1 deletion backend/app/services/collections/providers/registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from openai import OpenAI

from app.crud import get_provider_credential
from app.crud.rag.open_ai import OPENAI_MAX_RETRIES, OPENAI_TIMEOUT_SECONDS
from app.services.collections.providers.base import BaseProvider
from app.services.collections.providers.gemini import GeminiAIStudioProvider
from app.services.collections.providers.openai import OpenAIProvider
Expand Down Expand Up @@ -63,7 +64,11 @@ def get_llm_provider(
if provider == LLMProvider.OPENAI:
if "api_key" not in credentials:
raise ValueError("OpenAI credentials not configured for this project.")
client = OpenAI(api_key=credentials["api_key"])
client = OpenAI(
api_key=credentials["api_key"],
max_retries=OPENAI_MAX_RETRIES,
timeout=OPENAI_TIMEOUT_SECONDS,
)
elif provider == LLMProvider.GOOGLE_AISTUDIO:
if "api_key" not in credentials:
raise ValueError(
Expand Down
4 changes: 4 additions & 0 deletions backend/app/services/doctransform/registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ class TransformationError(Exception):
".markdown": "markdown",
".csv": "csv",
".json": "json",
".xlsx": "xlsx",
".xls": "xls",
}

# Map format names to file extensions
Expand All @@ -51,6 +53,8 @@ class TransformationError(Exception):
"markdown": ".md",
"csv": ".csv",
"json": ".json",
"xlsx": ".xlsx",
"xls": ".xls",
}


Expand Down
Loading
Loading