Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
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
155 changes: 127 additions & 28 deletions backend/app/crud/rag/open_ai.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,47 @@
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 tenacity import (
RetryCallState,
retry,
retry_if_exception_type,
stop_after_attempt,
wait_exponential,
)

from app.models import Document, ProviderType

logger = logging.getLogger(__name__)

OPENAI_PROVIDER = ProviderType.openai.value

# Under the Celery soft time limit so a hung call can't eat the whole task window.
# SDK-level retries are off (registry.py: max_retries=0); tenacity is the sole
# retry layer, wrapping batch create+index.
OPENAI_TIMEOUT_SECONDS = 30

BATCH_POLL_INTERVAL_SECONDS = 2

# Retry batch create+index on any OpenAI/indexing failure, exponential backoff
# (~2s, 4s, 8s), all inside one Celery soft-time-limit window.
BATCH_INDEX_MAX_ATTEMPTS = 4
BATCH_RETRY_BACKOFF_BASE_SECONDS = 2


def _log_batch_retry(retry_state: RetryCallState) -> None:
logger.warning(
f"[OpenAIVectorStoreCrud._create_and_index_batch] Batch attempt failed, retrying | "
f"attempt={retry_state.attempt_number}, "
f"error={retry_state.outcome.exception() if retry_state.outcome else None}"
)


def vs_ls(client: OpenAI, vector_store_id: str):
kwargs = {}
Expand Down Expand Up @@ -85,7 +115,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 +131,100 @@ 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:
"""Poll until indexing finishes; the Celery soft time limit is the deadline."""
while True:
batch = self._retrieve_file_batch(batch_id, vector_store_id)
if batch.status != "in_progress":

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate and inspect the target file and relevant literals/usages without running repo code.
if [ -f backend/app/crud/rag/open_ai.py ]; then
  echo "== target file size =="
  wc -l backend/app/crud/rag/open_ai.py
  echo
  echo "== relevant file sections =="
  sed -n '120,145p;270,290p' backend/app/crud/rag/open_ai.py
  echo
  echo "== batch.status/status value literals in target =="
  python3 - <<'PY'
import re
from pathlib import Path
p=Path('backend/app/crud/rag/open_ai.py')
text=p.read_text()
for i,line in enumerate(text.splitlines(),1):
    if '.status' in line or re.search(r'"(?:in_progress|completed|failed|pending|success|cancelled)"', line):
        print(f"{i}: {line}")
PY
else
  echo "backend/app/crud/rag/open_ai.py not found"
  git ls-files | rg 'open_ai\.py$|rag' || true
fi

echo
echo "== tests/files mentioning open_ai batch status literals =="
git ls-files | rg 'test|spec|py$' | xargs -r grep -E 'batch\.status|in_progress|completed|success' | head -n 200 || true

Repository: ProjectTech4DevAI/kaapi-backend

Length of output: 22815


Centralize OpenAI file-batch status values.

"in_progress" and "completed" are hard-coded directly in the CRUD path. Extract batch-status constants or an enum here and use the same values in any tests/callers that depend on this protocol state.

🤖 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` at line 136, Centralize the OpenAI
file-batch status values used by the CRUD flow instead of hard-coding
"in_progress" and "completed" in the status handling around batch.status. Define
reusable constants or an enum in the appropriate OpenAI module, update this path
and dependent tests/callers to reference them, and preserve the existing
protocol values.

Source: Coding guidelines

return batch
Comment thread
coderabbitai[bot] marked this conversation as resolved.
time.sleep(BATCH_POLL_INTERVAL_SECONDS)

def _raise_if_batch_incomplete(
self,
batch: VectorStoreFileBatch,
batch_id: str,
vector_store_id: str,
docs: list[Document],
) -> None:
"""Raise on any indexing failure so the batch attempt is retried."""
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,
filter="failed",
)
doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs}
parts = []
for f in failed_files:
d = doc_by_file_id.get(f.id)
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._raise_if_batch_incomplete] 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._raise_if_batch_incomplete] Could not fetch per-file errors | "
f"{{'batch_id': '{batch_id}', 'error': '{str(err)}'}}"
)
raise

# Only 'completed' is success; a 'cancelled'/'failed' batch with no per-file
# failures slips past the failed-count check above.
if batch.status != "completed":
error_message = (
f"[OPENAI] Vector store indexing did not complete "
f"(status: {batch.status}). Retry the collection."
)
logger.error(
f"[OpenAIVectorStoreCrud._raise_if_batch_incomplete] {error_message} | "
f"vector_store_id={vector_store_id}, batch_id={batch_id}, "
f"status={batch.status}"
)
raise RuntimeError(error_message)

@retry(
reraise=True,
stop=stop_after_attempt(BATCH_INDEX_MAX_ATTEMPTS),
wait=wait_exponential(multiplier=BATCH_RETRY_BACKOFF_BASE_SECONDS),
retry=retry_if_exception_type((OpenAIError, RuntimeError)),
before_sleep=_log_batch_retry,
)
def _create_and_index_batch(
self, vector_store_id: str, docs: list[Document]
) -> tuple[VectorStoreFileBatch, str]:
"""Create the file batch, wait for indexing, verify it completed. Retried as
a unit on any OpenAI/indexing failure; SoftTimeLimitExceeded is not retried
(not an OpenAIError/RuntimeError) so it aborts the task inside the window."""
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)
self._raise_if_batch_incomplete(batch, batch_id, vector_store_id, docs)
return batch, batch_id

def update(
self,
vector_store_id: str,
Expand All @@ -115,11 +239,7 @@ 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, batch_id = self._create_and_index_batch(vector_store_id, docs)
except openai.RateLimitError as e:
error_message = (
f"[OPENAI] Rate limit exceeded (code: {e.status_code}): "
Expand Down Expand Up @@ -215,30 +335,9 @@ 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,
filter="failed",
)
doc_by_file_id = {d.file_id[OPENAI_PROVIDER]: d for d in docs}
parts = []
for f in failed_files:
d = doc_by_file_id.get(f.id)
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}")
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)}'}}"
)
raise

def delete(self, vector_store_id: str, retries: int = 3):
if retries < 1:
Expand Down
13 changes: 9 additions & 4 deletions backend/app/services/collections/create_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -403,9 +403,11 @@ 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
]
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 @@ -507,7 +509,10 @@ def execute_batch_job(
webhook_secret=webhook_secret,
)

except (Timeout, SoftTimeLimitExceeded) as err:
except (Timeout, SoftTimeLimitExceeded):
# Batch-level retries happen in-task (tenacity in OpenAIVectorStoreCrud),
# so hitting the soft time limit means the window is spent — fail, don't
# re-queue.
timeout_err = TimeoutError("Task exceeded soft time limit")
logger.warning(
"[create_collection.execute_batch_job] Collection Creation Timed Out | {'collection_job_id': '%s', 'error': '%s'}",
Expand Down
7 changes: 6 additions & 1 deletion backend/app/services/collections/providers/registry.py
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_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=0,
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
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
38 changes: 38 additions & 0 deletions backend/app/services/documents/helpers.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import logging
from typing import Optional, Tuple, Iterable, Union
from uuid import UUID

Expand All @@ -11,6 +12,10 @@
)
from app.crud import DocTransformationJobCrud, DocumentCrud
from app.services.doctransform import job as transformation_job
from app.services.documents.validator import (
DocumentValidationError,
validate_document_content,
)
from app.models import (
DocTransformJobCreate,
TransformationStatus,
Expand All @@ -23,6 +28,39 @@
)


logger = logging.getLogger(__name__)


def validate_upload(
*,
src: UploadFile,
target_format: str | None,
transformer: str | None,
) -> Tuple[str, str | None]:
"""
Full pre-storage gate: extension and transformer validation plus a content
sanity check. Returns (source_format, actual_transformer_or_none).

Raises: HTTPException(400) on client errors.
"""
source_format, actual_transformer = pre_transform_validation(
src_filename=src.filename,
target_format=target_format,
transformer=transformer,
)

try:
validate_document_content(file=src, source_format=source_format)
except DocumentValidationError as e:
logger.warning(
f"[validate_upload] Document failed sanity check | "
f"filename: {e.filename} | format: {source_format} | reason: {e.reason}"
)
raise HTTPException(status_code=400, detail=e.client_message)

return source_format, actual_transformer


def calculate_file_size(file: UploadFile) -> float:
"""
Calculate the size of an uploaded file in kilobytes.
Expand Down
Loading
Loading