fix: replace OpenAI SDK with raw HTTP to avoid Cloudflare WAF blocking, broaden exception handling#880
Open
hungcuong9125 wants to merge 4 commits into
Conversation
…g, broaden exception handling, align vector dimensions at startup Problem - The OpenAI Python SDK (AsyncOpenAI) injects x-stainless-* telemetry headers into every request. Cloudflare-proxied OpenAI-compatible endpoints block these, causing embedding generation to fail with PermissionDeniedError. These exceptions are not ValueError subclasses, so they bypassed the except ValueError handler and became unhandled HTTP 500 errors. - pgvector column dimension mismatch between migration default (1536) and runtime config caused unhandled DataError on insert. - Missing broad exception catch around embedding calls in document CRUD. Changes - src/embedding_client.py: Replace AsyncOpenAI with raw httpx.AsyncClient POST requests, sending only Authorization and Content-Type headers. - src/crud/document.py: Add DataError alongside IntegrityError; add broad except Exception catch around embedding call → ValidationException (422). - docker/entrypoint.sh: Add scripts/configure_embeddings.py --yes step after migration to auto-align pgvector dimensions to current config. - tests/: Update all hardcoded 1536-dimension test vectors to use settings.EMBEDDING.VECTOR_DIMENSIONS across 6 test files (47 occurrences). Testing - Manual: create conclusion with session_id, without session_id, self-conclusion all return HTTP 201 instead of 500. - Unit: 27/27 conclusion route tests pass. - CI: ruff check, basedpyright, pytest pass.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe startup flow now runs embedding configuration after database provisioning. OpenAI embeddings use direct HTTP requests. Document CRUD broadens embedding and database error handling. Tests now use configured embedding vector dimensions instead of hardcoded 1536 values. ChangesConfigurable embedding dimensions and OpenAI HTTP client
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant EntryPoint as docker/entrypoint.sh
participant ProvisionDB as scripts/provision_db.py
participant ConfigureEmbeddings as scripts/configure_embeddings.py
EntryPoint->>ProvisionDB: start provisioning
ProvisionDB->>ProvisionDB: run database initialization
ProvisionDB->>ConfigureEmbeddings: execute --yes
ConfigureEmbeddings-->>ProvisionDB: stdout, stderr, exit code
ProvisionDB-->>EntryPoint: continue startup
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/crud/document.py (1)
896-899: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winBroad
except Exceptionconflates client errors with transient/infra failures.Wrapping any embedding-call exception (timeouts, connection resets, provider 5xx/rate-limit errors) as a
ValidationException(422) tells the client their input was invalid, when the real cause may be a transient dependency failure the client should retry rather than "fix." This trades one wrong status code (500) for another kind of misleading one (422).Consider catching narrower transport/HTTP exception types (e.g., from
httpx) and mapping them to a more appropriate status (e.g., 503) while reservingValidationExceptionfor genuinely malformed input.🤖 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 `@src/crud/document.py` around lines 896 - 899, The broad exception handling around the embedding generation block in document CRUD is mapping all failures to ValidationException, which incorrectly treats transport/provider outages as bad input. Narrow the except path around the embedding call to catch only genuine validation/malformed-input cases for ValidationException, and handle HTTP/transport failures separately in the relevant embedding helper or method with a more appropriate retryable/server-side exception (for example, using the document embedding flow in src/crud/document.py and the embedding-call wrapper that raises the current error).src/embedding_client.py (1)
263-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate request-building/error-handling logic between single and batch embed paths.
embed()and the OpenAI branch of_process_batch()both construct the same headers dict, POST to the same endpoint, checkresp.is_success, and parseresp.json()["data"]almost identically. Consider extracting a small private helper (e.g._openai_post_embeddings(http_client, payload)) to avoid maintaining two copies of this logic (e.g., if auth header format or error message ever needs to change).♻️ Proposed extraction
+ async def _openai_post_embeddings( + self, http_client: httpx.AsyncClient, payload: dict[str, Any] + ) -> dict[str, Any]: + resp = await http_client.post( + f"{self._openai_base_url}/embeddings", + json=payload, + headers={ + "Authorization": f"Bearer {self._openai_api_key}", + "Content-Type": "application/json", + }, + ) + if not resp.is_success: + raise RuntimeError(f"Embedding API error: {resp.status_code} {resp.text}") + return resp.json()Then call
data = await self._openai_post_embeddings(http_client, payload)from bothembed()and_process_batch().Also applies to: 478-500
🤖 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 `@src/embedding_client.py` around lines 263 - 282, The OpenAI embedding request logic is duplicated in both embed() and the OpenAI branch of _process_batch(), including headers, POST request, success check, and JSON parsing. Extract this into a small private helper such as _openai_post_embeddings(http_client, payload) inside EmbeddingClient, then have both embed() and _process_batch() call it so the auth header, endpoint, and error handling stay in one place.
🤖 Prompt for all review comments with 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.
Inline comments:
In `@docker/entrypoint.sh`:
- Around line 7-10: The startup script currently runs the DDL-capable
configure_embeddings.py --yes step unconditionally in entrypoint execution,
which can race across multiple replicas. Move this schema-altering embedding
configuration out of the container entrypoint and into a one-off migration/init
job or equivalent singleton startup path, using the configure_embeddings.py flow
only where it will run once per deployment rather than per app replica.
In `@src/crud/document.py`:
- Around line 1046-1049: The shared except block in the document creation flow
is too broad because it also catches DataError/IntegrityError from later
sync-state commits, causing a misleading “Failed to create observations” message
after the documents have already been committed. Narrow the handler around the
initial insert in the document creation path, or update the ValidationException
message in the create_observations logic to explicitly cover both document
creation and sync-state updates so the error text matches where the failure can
occur.
In `@src/embedding_client.py`:
- Line 9: The OpenAI test monkeypatches are targeting
src.embedding_client.AsyncOpenAI, but that symbol no longer exists and causes
AttributeError during test setup. Update the helpers in
tests/llm/test_embedding_client.py, especially _build_openai_client and related
setup, to patch httpx.AsyncClient behavior or use an injected fake client
instead of monkeypatching AsyncOpenAI. Keep the test coverage focused on the
embedding client’s request/response handling through the actual
httpx.AsyncClient path.
---
Nitpick comments:
In `@src/crud/document.py`:
- Around line 896-899: The broad exception handling around the embedding
generation block in document CRUD is mapping all failures to
ValidationException, which incorrectly treats transport/provider outages as bad
input. Narrow the except path around the embedding call to catch only genuine
validation/malformed-input cases for ValidationException, and handle
HTTP/transport failures separately in the relevant embedding helper or method
with a more appropriate retryable/server-side exception (for example, using the
document embedding flow in src/crud/document.py and the embedding-call wrapper
that raises the current error).
In `@src/embedding_client.py`:
- Around line 263-282: The OpenAI embedding request logic is duplicated in both
embed() and the OpenAI branch of _process_batch(), including headers, POST
request, success check, and JSON parsing. Extract this into a small private
helper such as _openai_post_embeddings(http_client, payload) inside
EmbeddingClient, then have both embed() and _process_batch() call it so the auth
header, endpoint, and error handling stay in one place.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e025854b-a406-47db-8d18-eafeb2634433
📒 Files selected for processing (9)
docker/entrypoint.shsrc/crud/document.pysrc/embedding_client.pytests/alembic/revisions/test_b8183c5ffb48_codify_document_level_and_times_derived.pytests/crud/test_document.pytests/deriver/test_vector_reconciliation.pytests/llm/test_embedding_client.pytests/routes/test_conclusions.pytests/utils/test_agent_tools.py
…ty, helper extraction, test fix Changes per review: 1. docker/entrypoint.sh — Move DDL configure_embeddings.py --yes into scripts/provision_db.py (DB init runs once) via subprocess.run non-fatal, avoiding race condition on ALTER COLUMN across replicas. 2. scripts/provision_db.py — Add non-fatal subprocess.run call to configure_embeddings.py --yes after alembic migration completes. 3. src/crud/document.py — Split try/except: document insert wrapped with (DataError, IntegrityError) → ValidationException; sync-state updates moved outside the insert try/except block. Broaden embedding-error handling: ValueError → 422, httpx.HTTPStatusError → 503, httpx.ConnectError/TimeoutException/RemoteProtocolError → 503, generic Exception → 500. 4. src/embedding_client.py — Extract _openai_post_embeddings(http_client, payload) helper to eliminate duplicate POST logic between embed() and _process_batch(). 5. tests/llm/test_embedding_client.py — Replace monkeypatched AsyncOpenAI with monkeypatched httpx.AsyncClient.post using httpx.Response fixtures, fixing AttributeError from removed class.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/llm/test_embedding_client.py (1)
315-326: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
fake_postmocking logic across tests.Three tests (
test_simple_batch_embed_respects_token_budget_per_request,test_simple_batch_embed_rejects_oversized_input,test_prepare_chunks_returns_ordered_chunks) reimplement inlinefake_post/client-construction logic that closely mirrors_build_openai_client. Consider extending_build_openai_client(e.g. optionalmax_input_tokens/max_tokens_per_requestparams) to cover these cases and reduce duplicated mock setup.Also applies to: 353-356, 379-382
🤖 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 `@tests/llm/test_embedding_client.py` around lines 315 - 326, The three embedding-client tests are duplicating the same `fake_post` and client setup logic instead of reusing the shared helper. Update `_build_openai_client` in the embedding client test utilities to accept the needed optional parameters such as `max_input_tokens` and/or `max_tokens_per_request`, then switch `test_simple_batch_embed_respects_token_budget_per_request`, `test_simple_batch_embed_rejects_oversized_input`, and `test_prepare_chunks_returns_ordered_chunks` to use that helper so the mocked `httpx.AsyncClient.post` setup is centralized.src/crud/document.py (1)
859-862: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring
Raisessection is now incomplete.The new embedding-failure paths (Lines 898-912) raise
HonchoExceptionfor provider HTTP errors and transient connectivity issues, but the docstring only documentsResourceNotFoundExceptionandValidationException.📝 Proposed docstring update
Raises: ResourceNotFoundException: If any session or peer is not found - ValidationException: If embedding generation fails or integrity constraint is violated + ValidationException: If embedding validation fails or a data/integrity constraint is violated + HonchoException: If the embedding provider returns an error or is unreachableAs per coding guidelines, "In Python code, use Google-style docstrings."
🤖 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 `@src/crud/document.py` around lines 859 - 862, Update the Google-style docstring for the method in document.py that handles embedding generation so the Raises section includes HonchoException alongside ResourceNotFoundException and ValidationException. Add a brief note that HonchoException is raised for provider HTTP errors and transient connectivity failures in the embedding-generation paths, and keep the documentation aligned with the new error behavior in the embedding workflow.Source: Coding guidelines
🤖 Prompt for all review comments with 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.
Inline comments:
In `@src/crud/document.py`:
- Around line 977-988: The pgvector-only sync-state update in the document flow
is missing the same error handling used in the external-vector-store branch, so
a transient failure in db.execute or db.commit can still bubble up as a 500.
Wrap the no-external-vector-store path in the same kind of try/except handling
used around upsert_many/db.execute/commit for the external branch in
src/crud/document.py, and make sure the exception path logs the failure and
returns a controlled response instead of propagating it.
---
Nitpick comments:
In `@src/crud/document.py`:
- Around line 859-862: Update the Google-style docstring for the method in
document.py that handles embedding generation so the Raises section includes
HonchoException alongside ResourceNotFoundException and ValidationException. Add
a brief note that HonchoException is raised for provider HTTP errors and
transient connectivity failures in the embedding-generation paths, and keep the
documentation aligned with the new error behavior in the embedding workflow.
In `@tests/llm/test_embedding_client.py`:
- Around line 315-326: The three embedding-client tests are duplicating the same
`fake_post` and client setup logic instead of reusing the shared helper. Update
`_build_openai_client` in the embedding client test utilities to accept the
needed optional parameters such as `max_input_tokens` and/or
`max_tokens_per_request`, then switch
`test_simple_batch_embed_respects_token_budget_per_request`,
`test_simple_batch_embed_rejects_oversized_input`, and
`test_prepare_chunks_returns_ordered_chunks` to use that helper so the mocked
`httpx.AsyncClient.post` setup is centralized.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8571742b-27d4-427d-9119-840a5da7b79c
📒 Files selected for processing (5)
docker/entrypoint.shscripts/provision_db.pysrc/crud/document.pysrc/embedding_client.pytests/llm/test_embedding_client.py
✅ Files skipped from review due to trivial changes (1)
- docker/entrypoint.sh
🚧 Files skipped from review as they are similar to previous changes (1)
- src/embedding_client.py
…cstring, wrap pgvector sync in try/except Changes: - tests/llm/test_embedding_client.py: Extend _build_openai_client helper with optional max_input_tokens/max_tokens_per_request params and refactor 3 tests (test_simple_batch_embed_respects_token_budget_per_request, test_simple_batch_embed_rejects_oversized_input, test_prepare_chunks_returns_ordered_chunks) to use it, eliminating duplicate fake_post mocking logic. - src/crud/document.py: Update Google-style docstring Raises section to include HonchoException for provider HTTP errors and transient connectivity failures. Wrap the pgvector-only sync-state update path (no external vector store) in try/except with logger.exception to prevent transient DB failures from bubbling up as HTTP 500.
…r path - Add await db.rollback() inside except block so the failed transaction is properly rolled back on transient DB errors. - Include workspace_name in the log message for better debugging context.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fix
POST /v3/workspaces/{workspace_id}/conclusionsreturning HTTP 500 due to embedding API failures and unhandled database exceptions.Root Cause
The OpenAI Python SDK (
AsyncOpenAI) automatically injectsx-stainless-*telemetry headers into every request. When using a Cloudflare-proxied OpenAI-compatible endpoint (e.g.,api.hungcuong.dev), Cloudflare's WAF blocks these non-standard headers, causing embedding generation to fail withPermissionDeniedError. Since the SDK's exception types (APIConnectionError,APIStatusError, etc.) are not subclasses ofValueError, they bypassed the existingexcept ValueErrorhandler and propagated to the global exception handler as HTTP 500.Secondary issue: a pgvector column dimension mismatch between the deployment config and migration default also caused an unhandled
DataErroron insert.Changes
src/embedding_client.py— Replace OpenAI SDK with raw HTTPProblem:
AsyncOpenAI(api_key=..., base_url=...)sendsx-stainless-arch,x-stainless-lang,x-stainless-package-version,x-stainless-runtime,x-stainless-runtime-versionheaders. Cloudflare WAF at common OpenAI-compatible reverse proxies rejects these.Fix: Replace
AsyncOpenAIwithhttpx.AsyncClientand make raw POST requests to{base_url}/embeddingswith onlyAuthorizationandContent-Typeheaders. This applies to bothembed()(single query) and_process_batch()(batch) code paths.src/crud/document.py— Broaden exception handlingDataErroralongsideIntegrityErrorin the DB commit catch block — covers pgvector dimension mismatches.except Exceptionaround the embedding call — any embedding failure (network, API, auth) now becomes aValidationException(422) instead of 500.docker/entrypoint.sh— Auto-align pgvector dimensionsMigration
b8183c5ffb48hardcodesVector(1536). If the deployment config uses a different dimension, inserts fail withDataError. Addedscripts/configure_embeddings.py --yesstep at startup to ALTER the vector column to match the current config.Test fixes
Replaced all hardcoded
[0.x] * 1536test vectors with[0.x] * settings.EMBEDDING.VECTOR_DIMENSIONSacross 6 test files (47 occurrences) to match the runtime configuration.Testing
session_id✅session_id✅observer_id == observed_id) ✅Checklist
ruff check,basedpyright)pytest tests/routes/test_conclusions.py)docker compose build)Summary by CodeRabbit
New Features
Bug Fixes
Tests
1536) and refactored OpenAI client tests to mockhttpx.