Feat/add file quota2 - #233
Conversation
UserWarning: Duplicate Operation ID list_models_v1_models_get for function list_models at /app/openrag/routers/openai.py warnings.warn(message, stacklevel=1)
Add relationship_id and parent_id fields to support document linking:
- relationship_id: Groups related documents (email threads, folders)
- parent_id: Hierarchical parent reference (parent email, parent folder)
Changes:
- Add SQLAlchemy columns and indexes to File model
- Add Alembic migration for new database columns
- Add PartitionFileManager query methods with recursive CTE for ancestors
- Add VectorDB wrapper methods and Milvus INVERTED indexes
- Add API endpoints: GET /{partition}/relationships/{id} and ancestors
- Add include_related/include_ancestors params to search endpoints
- Add RelationshipAwareRetriever for context-aware retrieval
- Add unit tests (14) and integration tests (11)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
requires database migration to add file_quota column to users table.
Features:
- Add file_quota field to users table (nullable integer)
- Add DEFAULT_FILE_QUOTA env var to set global default quota
- Add PATCH /users/{user_id}/quota endpoint to update user quotas
- Add quota enforcement on file upload (indexed files + pending tasks)
- Admins bypass quota checks; quota<=0 for a user means unlimited; quota > 0 to limit user's quota
API Changes:
- POST /users/ now accepts optional file_quota parameter
- GET /users/ and GET /users/{id} return file_quota in response
- New PATCH /users/{user_id}/quota endpoint for quota management
- GET /users/info now renders additional fields: indexed_files, pending_files, total_files and file_quota
Tests:
- test_update_user_quota: verify quota update from 10 to 12
- test_user_default_quota: verify None quota defaults to 10
- TestUserQuotaEnforcement class in test_indexer.py:
- test_unlimited_quota_user_can_exceed_default: user with quota=0 uploads 11 files
- test_quota_limit_blocks_excess_uploads: user with quota=5 blocked on 6th file
Perform migration: Refer to the doc
https://github.com/linagora/openrag/blob/0b5f84cd880e7c75db6f78d12a7681227721c2ad/docs/content/docs/documentation/sql_migration.mdx?plain=1#L47-L53
requires database migration to add file_quota column to users table.
Features:
- Add file_quota field to users table (nullable integer)
- Add DEFAULT_FILE_QUOTA env var to set global default quota
- Add PATCH /users/{user_id}/quota endpoint to update user quotas
- Add quota enforcement on file upload (indexed files + pending tasks)
- Admins bypass quota checks; quota<=0 for a user means unlimited; quota > 0 to limit user's quota
API Changes:
- POST /users/ now accepts optional file_quota parameter
- GET /users/ and GET /users/{id} return file_quota in response
- New PATCH /users/{user_id}/quota endpoint for quota management
- GET /users/info now renders additional fields: file_count, pending_files, total_files and file_quota
Tests:
- test_update_user_quota: verify quota update from 10 to 12
- test_user_default_quota: verify None quota defaults to 10
- TestUserQuotaEnforcement class in test_indexer.py:
- test_unlimited_quota_user_can_exceed_default: user with quota=0 uploads 11 files
- test_quota_limit_blocks_excess_uploads: user with quota=5 blocked on 6th file
Perform migration: Refer to the doc
https://github.com/linagora/openrag/blob/0b5f84cd880e7c75db6f78d12a7681227721c2ad/docs/content/docs/documentation/sql_migration.mdx?plain=1#L47-L53
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a per-user file quota system: DB schema changes (user.file_quota, user.file_count), quota-aware API checks on upload/copy endpoints, task-aware pending counts, vectordb/indexer call-site updates to propagate user context, and new tests and config vars to exercise quota behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Router as Indexer Router
participant Validator as Quota Validator
participant VectorDB as VectorDB
participant TaskMgr as Task Manager
participant DB as Database
Client->>Router: POST /add_file (user token)
Router->>Validator: check_user_file_quota(user)
Validator->>VectorDB: get_user_by_token(user_token)
VectorDB->>DB: Query User (file_quota, file_count)
DB-->>VectorDB: User data
Validator->>TaskMgr: get_user_pending_task_count(user_id)
TaskMgr-->>Validator: pending_count
Validator->>Validator: total = file_count + pending_count
alt total >= user_quota
Validator-->>Router: HTTP 403 Quota Exceeded
Router-->>Client: 403 Error
else within quota
Validator-->>Router: approved
Router->>VectorDB: add_file_to_partition(file_id, user_id)
VectorDB->>DB: Insert File, Increment User.file_count
DB-->>VectorDB: Success
VectorDB-->>Router: File added / task queued
Router-->>Client: Task queued response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
openrag/components/indexer/vectordb/vectordb.py (1)
883-890:⚠️ Potential issue | 🔴 Critical
delete_usercallsdelete_partitionwithout requireduser_idargument.
delete_partitioninPartitionFileManagernow requires auser_idparameter (as seen in the relevant snippet fromutils.pyline 269), but line 889 calls it without one. This will raise aTypeErrorat runtime when deleting a user who owns partitions.Proposed fix
async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: - self.partition_file_manager.delete_partition(partition) + self.partition_file_manager.delete_partition(partition, user_id=user_id) self.partition_file_manager.delete_user(user_id)openrag/components/indexer/vectordb/utils.py (2)
270-290:⚠️ Potential issue | 🟠 Major
delete_partitiondecrements the caller'sfile_countby the total files in the partition, ignoring multi-user uploads.If user A uploaded 3 files and user B uploaded 2 files to the same partition, and user A deletes the partition, user A's
file_countis decremented by 5 (all files), not 3 (only A's files). This corrupts the count for user A.To fix this correctly, you'd need to track which user uploaded each file (e.g., a
user_idcolumn on theFilemodel) and decrement each user's count by their respective file count in the partition. Alternatively, count files per user via the existingfile_countper-user-per-partition before deletion.🛠️ Sketch of a more correct approach
One option: add
uploaded_byto theFilemodel. Then indelete_partition:# Count files per uploader from sqlalchemy import func file_counts = ( session.query(File.uploaded_by, func.count()) .filter(File.partition_name == partition) .group_by(File.uploaded_by) .all() ) for uid, cnt in file_counts: user = session.query(User).filter(User.id == uid).first() if user: user.file_count = func.greatest(0, User.file_count - cnt)If adding a column is too invasive now, consider reconciling
file_countfrom the live query (get_user_file_count) after partition deletion instead.
245-268:⚠️ Potential issue | 🟠 Major
remove_file_from_partitiondecrements the caller's file_count, not the original uploader's.The
user_idparameter is presumably the user requesting the deletion. If a different user uploaded the file, this user'sfile_countis incorrectly decremented. This is the same ownership-tracking gap as noted indelete_partition.
🤖 Fix all issues with AI agents
In @.gitmodules:
- Line 7: The .gitmodules entry currently points the submodule branch "branch =
feat/handling_user_file_quota", which risks breakage when feature branches are
deleted; update the .gitmodules configuration for the affected submodule (the
"branch" field) to reference a stable branch (e.g., "dev" or "main") or remove
the branch field and pin the submodule to a specific commit hash instead (or
after merging the feature, set the branch to the merged stable branch) so that
git submodule update remains reproducible.
In `@docs/content/docs/documentation/env_vars.md`:
- Line 420: The table entry for DEFAULT_FILE_QUOTA is incorrect — change its
default value from `0` to `-1` and update the description to match the
config.yaml/data_model.md semantics: show `-1` as the default and state that
`<=0` (or specifically `-1`) disables quotas globally while `>0` sets the
default per-user quota when no explicit quota exists; reference the
DEFAULT_FILE_QUOTA setting and ensure the wording aligns with the value defined
in config.yaml and documented in data_model.md.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 231-236: The current read-modify-write pattern user.file_count +=
1 in the indexer functions causes a race; replace the in-memory
increment/decrement with SQL-level atomic updates using a column expression so
the DB executes SET file_count = file_count + 1 (and -1) atomically. Concretely,
in the function where you increment (the block using
session.query(User).filter(User.id == user_id).first()) and in
remove_file_from_partition and delete_partition, switch to a query.update() (or
session.execute with an update clause) targeting User.file_count =
User.file_count + 1 (and User.file_count = User.file_count - 1 for decrements)
filtered by User.id == user_id, then flush/commit as before to persist the
change.
In `@openrag/routers/users.py`:
- Around line 60-93: The response key for indexed files in get_current_user_info
is "file_count" but the docstring/documentation references "indexed_files";
update either the endpoint response to include "indexed_files" (e.g., add
"indexed_files": indexed_count alongside "file_count" or rename the key to
"indexed_files") or change the docstring to reference "file_count" so the docs
and the get_current_user_info response remain consistent; ensure any tests or
consumers referencing file_count/indexed_files are updated accordingly.
- Around line 260-275: The update_user_quota endpoint currently calls
vectordb.update_user_quota.remote and discards its return value, and if the
target user doesn't exist the underlying _check_user_exists exception bubbles up
as a 500; update the update_user_quota function to await
vectordb.update_user_quota.remote into a variable (e.g., updated = await
vectordb.update_user_quota.remote(...)), catch the specific exception raised by
_check_user_exists (or a user-not-found error from the vectordb client) and
return JSONResponse(status_code=404, content={"detail": "User not found"}), and
on success return the updated user dict in the response body (HTTP 200) while
keeping the existing debug log via logger.debug; ensure require_admin and
vectordb dependency usage remains unchanged.
In `@openrag/routers/utils.py`:
- Around line 194-204: Fix the typo "to disabled" → "to disable" in the
docstring and the inline comment near the quota-check logic that references
DEFAULT_FILE_QUOTA and the admin check (the block using user.get("is_admin") and
the conditional "if DEFAULT_FILE_QUOTA <= 0"). Update both the docstring line
describing quota semantics and the inline comment above the DEFAULT_FILE_QUOTA
check to read "to disable quota checking".
In
`@openrag/scripts/migrations/alembic/versions/bd24665e3451_add_file_quota_and_file_count_fields.py`:
- Around line 24-25: The migration adds a NOT NULL users.file_count column
without a DB default which will fail on existing rows; update the op.add_column
call that creates "file_count" (in this migration file where
op.add_column("users", sa.Column("file_count", ...))) to include a
server_default (e.g. server_default=sa.text("0")) so existing rows get 0 at
schema time, keeping nullable=False; after applying the migration you can
optionally remove the server_default in a follow-up migration if you don't want
it permanently.
In `@tests/api_tests/test_indexer.py`:
- Around line 619-620: The finally block in test_file_count_decrements_on_delete
only calls self._cleanup_partition(api_client, partition_name) and omits
cleaning up the test user; add a call to self._cleanup_user(api_client, user_id)
in that finally block (matching other tests) so the created user is
removed—ensure the test uses the same user_id variable created earlier in the
test and keep the partition cleanup call intact (functions referenced:
_cleanup_partition, _cleanup_user, and user_id).
- Around line 486-514: The test test_quota_limit_blocks_excess_uploads relies on
quota enforcement being enabled; ensure the test explicitly verifies or sets
that precondition by checking DEFAULT_FILE_QUOTA > 0 (import DEFAULT_FILE_QUOTA
from openrag.routers.utils) at the start of the test or by using the
partition/user PATCH quota endpoint after creating the user (methods involved:
_create_user_with_quota, _create_partition, _upload_file) to set a positive
file_quota so quota enforcement is guaranteed regardless of environment.
In `@tests/api_tests/test_users.py`:
- Around line 68-84: The test_user_default_quota test is asserting a hardcoded
10 which couples it to environment config; instead import or read the configured
DEFAULT_FILE_QUOTA used by the app (e.g., from the config/settings module or
environment variable) and use that value as the expected value in the assertion.
Update the test to obtain DEFAULT_FILE_QUOTA via the same symbol your app uses
(e.g., DEFAULT_FILE_QUOTA from settings/config) and replace the literal 10 in
the assertion so api_client-based creation and the get_response check compare
against the configured constant.
🧹 Nitpick comments (8)
openrag/routers/partition.py (2)
414-427: Unusedrequestparameter inget_related_files.The
requestparameter is declared but never used in the function body, unlikeget_file_ancestorswhere it's also unused. Consider removing it to keep the signature clean.Proposed fix
async def get_related_files( - request: Request, partition: str, relationship_id: str, vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ):
455-475: Unusedrequestparameter inget_file_ancestors.Same as
get_related_files—requestis injected but never referenced.docs/content/docs/documentation/linked_files.md (1)
65-65: Add language specifiers to fenced code blocks.Several code blocks (lines 65, 228, 233, 241, 282) lack language identifiers, flagged by markdownlint (MD040). Use
textorplaintextfor non-code examples.Also applies to: 228-246
openrag/routers/utils.py (1)
219-239: TOCTOU window between quota check and task registration.There's a race condition: two concurrent upload requests for the same user can both pass the quota check before either task is registered in
TaskStateManager, allowing the user to exceed their quota by one or more files. Since this is a soft limit and not a financial control, it's likely acceptable — but worth documenting if precision matters.tests/api_tests/test_users.py (1)
54-57: Brittle assertion on message string content.Asserting
str(user_id) in update_data["message"]and"12" in update_data["message"]couples the test to the exact message format. Consider asserting on structured response fields instead (e.g., checking the updated quota value directly, which you already do on line 63).tests/api_tests/test_indexer.py (1)
470-480: Redundant status check after assertion.Line 476 checks
if response.status_code in [200, 201, 202], but line 472 already asserts this — so the condition is alwaysTrueat that point. The same pattern appears in the other tests (lines 504-508). While harmless, it adds noise.✂️ Simplify by removing the redundant check
assert response.status_code in [200, 201, 202], ( f"File {i} upload failed with status {response.status_code}: {response.text}" ) # Wait for task if needed - if response.status_code in [200, 201, 202]: - data = response.json() - if "task_status_url" in data: - task_id = get_task_id(data) - wait_for_task(api_client, task_id, headers={"Authorization": f"Bearer {user_token}"}) + data = response.json() + if "task_status_url" in data: + task_id = get_task_id(data) + wait_for_task(api_client, task_id, headers={"Authorization": f"Bearer {user_token}"})openrag/components/indexer/vectordb/utils.py (2)
337-343: No-op assignment on line 341.
file_quota = file_quotadoes nothing. You can remove theelifbranch or add a comment if the intent is to document the pass-through case.✂️ Simplify
if self.file_quota_per_user > 0: # Quotas enabled globally if file_quota is None: file_quota = self.file_quota_per_user # default to default quota - elif file_quota > 0: - file_quota = file_quota # use specified quota else: - pass # unlimited + pass # file_quota > 0: use as-is; file_quota <= 0: unlimited
560-580:update_user_quotadoesn't guard againstuserbeingNone.If the caller doesn't go through the vectordb actor's
_check_user_existsgate (or if a future refactor removes that check), line 569 will raiseAttributeError. A defensive check would be prudent.🛡️ Proposed fix
with self.Session() as s: user = s.query(User).filter(User.id == user_id).first() + if not user: + raise ValueError(f"User {user_id} not found") user.file_quota = file_quota
| [submodule "extern/indexer-ui"] | ||
| path = extern/indexer-ui | ||
| url = https://github.com/linagora/openrag-admin-ui.git | ||
| branch = feat/handling_user_file_quota |
There was a problem hiding this comment.
Avoid pointing submodules to feature branches.
The submodule is now configured to track the feature branch feat/handling_user_file_quota. This creates a maintenance risk because feature branches are typically deleted after merging, which will break git submodule update for anyone cloning or updating the repository later.
Recommended approach:
- Merge the indexer-ui feature branch to a stable branch (e.g.,
devormain) first - Update this
.gitmodulesto point to that stable branch, or - Use a specific commit hash instead of a branch name for reproducibility
📌 Example: Using a specific commit instead of a branch
After merging the indexer-ui PR, pin to the merge commit:
[submodule "extern/indexer-ui"]
path = extern/indexer-ui
url = https://github.com/linagora/openrag-admin-ui.git
- branch = feat/handling_user_file_quota
+ branch = devOr remove the branch field entirely and use git submodule update --remote with a specific commit when needed.
🤖 Prompt for AI Agents
In @.gitmodules at line 7, The .gitmodules entry currently points the submodule
branch "branch = feat/handling_user_file_quota", which risks breakage when
feature branches are deleted; update the .gitmodules configuration for the
affected submodule (the "branch" field) to reference a stable branch (e.g.,
"dev" or "main") or remove the branch field and pin the submodule to a specific
commit hash instead (or after merging the feature, set the branch to the merged
stable branch) so that git submodule update remains reproducible.
| def test_user_default_quota(self, api_client): | ||
| """Test that user created with None quota gets default value (10).""" | ||
| # Create a user without specifying quota (None) | ||
| create_response = api_client.post( | ||
| "/users/", | ||
| data={"display_name": "default_quota_user"}, | ||
| ) | ||
| assert create_response.status_code == 201 | ||
| user_data = create_response.json() | ||
| user_id = user_data["id"] | ||
|
|
||
| try: | ||
| # Verify the user has the default quota (10) | ||
| get_response = api_client.get(f"/users/{user_id}") | ||
| assert get_response.status_code == 200 | ||
| get_data = get_response.json() | ||
| assert get_data["file_quota"] == 10, "User should have default quota of 10" |
There was a problem hiding this comment.
Hardcoded default quota value couples test to environment config.
The expected default quota of 10 (line 84) must match the DEFAULT_FILE_QUOTA configured in the test environment. If the config changes, this test silently becomes incorrect. Consider reading the expected default from the same config source or documenting the dependency on the test environment configuration.
🤖 Prompt for AI Agents
In `@tests/api_tests/test_users.py` around lines 68 - 84, The
test_user_default_quota test is asserting a hardcoded 10 which couples it to
environment config; instead import or read the configured DEFAULT_FILE_QUOTA
used by the app (e.g., from the config/settings module or environment variable)
and use that value as the expected value in the assertion. Update the test to
obtain DEFAULT_FILE_QUOTA via the same symbol your app uses (e.g.,
DEFAULT_FILE_QUOTA from settings/config) and replace the literal 10 in the
assertion so api_client-based creation and the get_response check compare
against the configured constant.
13515d2 to
8b24f46
Compare
There was a problem hiding this comment.
Actionable comments posted: 12
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
openrag/components/indexer/indexer.py (1)
296-312:⚠️ Potential issue | 🟡 MinorWrong type hint:
partition: intshould bepartition: str.In
set_details, thepartitionparameter is typed asintbut it holds a partition name string (e.g., used in metadata dict on line 309).Proposed fix
async def set_details( self, task_id: str, *, file_id: str, - partition: int, + partition: str, metadata: dict, user_id: int, ):openrag/components/indexer/vectordb/vectordb.py (2)
53-62: 🛠️ Refactor suggestion | 🟠 MajorAbstract base class signatures are out of sync with the implementation.
BaseVectorDB.delete_partition(line 53) andBaseVectorDB.delete_file(line 61) don't includeuser_id, butMilvusDBnow requires it. This breaks the Liskov Substitution Principle and will confuse any future implementers ofBaseVectorDB.♻️ Update ABC signatures
`@abstractmethod` - async def delete_partition(self, partition: str): + async def delete_partition(self, partition: str, user_id: int): pass ... `@abstractmethod` - async def delete_file(self, file_id: str, partition: str): + async def delete_file(self, file_id: str, partition: str, user_id: int): pass
898-905:⚠️ Potential issue | 🔴 CriticalBug:
delete_partitioncalled without requireduser_idargument.Line 904 calls
self.partition_file_manager.delete_partition(partition), but the updated signature inutils.pyline 300 now requiresuser_id: intas a second positional argument. This will raise aTypeErrorat runtime when deleting a user who owns partitions.🐛 Proposed fix
async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: - self.partition_file_manager.delete_partition(partition) + self.partition_file_manager.delete_partition(partition, user_id=user_id) self.partition_file_manager.delete_user(user_id)
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/linked_files.md`:
- Line 65: Several fenced code blocks in the markdown lack language specifiers
(triple-backtick blocks), causing markdownlint MD040 failures; edit each
offending fenced code block and add the appropriate language identifier—use
`text` for plain-text examples and `txt` for diagram-like blocks—specifically
update the triple-backtick blocks mentioned in the comment so they read ```text
or ```txt as appropriate.
- Line 32: Fix the minor grammar in the sentence that currently reads "The
`File` model includes now two relationship fields:" by changing it to "The
`File` model now includes two relationship fields:" so the placement of "now" is
correct; update the sentence where the `File` model is described in the docs
(the exact line containing "The `File` model includes now two relationship
fields:").
In `@openrag/components/indexer/indexer.py`:
- Around line 148-157: delete_file assumes user is a dict and calls
user.get("id"), but update_file_metadata can pass None; either make user
required in update_file_metadata or make delete_file accept None: change
delete_file signature to user: dict | None = None (or keep current signature)
and guard before accessing user (e.g., compute user_id = user.get("id") if user
else None) and pass that to vectordb.delete_file.remote; also update type hints
and any callers (notably update_file_metadata) to match the new optional user
handling so no AttributeError occurs.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 645-693: The recursive CTE in get_file_ancestors can loop on
cyclic parent_id chains; fix by adding a depth guard/limit to the CTE (e.g.,
introduce a MAX_DEPTH parameter and increment depth in the recursive step, then
constrain recursion with "WHERE a.depth < :max_depth" or equivalent) and surface
max_depth as a parameter to get_file_ancestors so callers can control it; ensure
the query binds the new max_depth parameter when executing the text SQL and
still returns results ordered from root to the file.
- Around line 575-595: update_user_quota currently assumes the user exists and
then accesses/returns fields after committing which can cause AttributeError or
DetachedInstanceError; update the method to first check if user is None after
the query (e.g., if user is None: raise ValueError or a suitable NotFound error)
to guard against missing users, and ensure the returned fields are safe to
access by either reading needed attributes before calling s.commit() or calling
s.refresh(user) after s.commit() (use the existing Session variable s and the
user object) before building and returning the dict.
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 330-342: The calls to index_params.add_index for dynamic document
fields use field_name but Milvus expects dynamic-field indexing to pass
json_path and a mandatory json_cast_type; update the two add_index calls (the
ones adding "relationship_id_idx" and "parent_id_idx") to replace
field_name="relationship_id"/"parent_id" with
json_path="$.relationship_id"/"$.parent_id" (or the correct JSON path used by
your stored documents) and add json_cast_type set to the proper Milvus type
(e.g., "VARCHAR" if these are strings or "DOUBLE"/"BOOL" as appropriate); keep
index_type="INVERTED" and index_name the same.
In `@openrag/routers/utils.py`:
- Around line 304-310: The code mixes safe .get() access
(llm_param.get("base_url"), llm_param.get("model"), llm_param.get("api_key"))
with direct indexing llm_param["base_url"] / llm_param["model"] when creating
the logger; change the logger.bind call to use the already-read variables
(base_url and model) or use llm_param.get(...) consistently (with sensible
defaults) so no KeyError can be raised—update the logger.bind invocation
(symbol: logger.bind and variables llm_param, base_url, model, log) to reference
base_url and model (or llm_param.get with defaults) instead of direct dict
indexing.
In
`@openrag/scripts/migrations/alembic/versions/344b49ce4f69_add_file_quota_field.py`:
- Around line 1-32: This migration adds the same file_quota column already added
in bd24665e3451_add_file_quota_and_file_count_fields.py, causing
duplicate/branching migrations; fix by consolidating to a single migration
chain: remove the op.add_column("users", sa.Column("file_quota", ...)) and
corresponding op.drop_column("users", "file_quota") from this file
(344b49ce4f69_add_file_quota_field.py) so only bd24665e3451... owns the column,
or alternatively merge the unique changes into one file and update the
revision/down_revision values accordingly; locate the op.add_column and
op.drop_column calls in this module to perform the change and ensure the
revision and down_revision identifiers remain consistent.
In `@openrag/tests/test_relationships_integration.py`:
- Around line 66-72: Replace the use of set(...) on a generator with a set
comprehension to satisfy Ruff C401: in
test_index_folder_files_share_relationship_id compute relationship_ids with a
comprehension like {doc.metadata.get("relationship_id") for doc in
folder_documents} instead of set(doc.metadata.get("relationship_id") for doc in
folder_documents); apply the same change to the other similar test that builds
relationship_ids in the file.
- Around line 12-14: Reorder and group the imports into standard-library,
third-party, and local groups with a blank line between groups; specifically
move "from unittest.mock import AsyncMock, MagicMock, patch" above "import
pytest" and keep "from langchain_core.documents.base import Document" in the
third-party group below a blank line, so the import order becomes:
standard-library (unittest.mock...), third-party (pytest, langchain_core...),
ensuring imports reference AsyncMock, MagicMock, patch, pytest, and Document as
shown.
In `@tests/api_tests/test_search.py`:
- Around line 90-93: Remove the Git conflict markers (<<<<<<<, =======, >>>>>>>)
in this test file and keep the intended new test declaration for
test_search_with_include_related; specifically locate the conflict around the
test_search_with_include_related function and delete the conflict marker lines
so the function signature and body remain valid and the module parses correctly.
- Around line 93-143: The test uses missing fixtures indexed_folder_partition,
exact_match_query, and folder_files which causes collection failures; add
fixture implementations in tests/api_tests/conftest.py that (1)
indexed_folder_partition: create and return a test partition id with an indexed
folder/thread relationship, (2) folder_files: create and index three files
(e.g., "file1.txt","file2.txt","file3.txt") under the same relationship_id and
return a mapping like {filename: (content, relationship_id)} so the test can
assert expected IDs, and (3) exact_match_query: return a query string that will
match a known chunk from file1.txt; ensure the fixtures handle indexing and
cleanup and remove any leftover merge conflict markers related to this test
addition.
🧹 Nitpick comments (7)
openrag/routers/users.py (2)
1-9: Module-level Ray actor initialization.
task_state_manager = get_task_state_manager()(line 9) runs at import time. This works becauseray.init()is called inapi.pybefore routers are imported, but it creates an implicit ordering dependency. If this module is ever imported in a context where Ray hasn't been initialized (e.g., tests, scripts), it will fail.
123-139:file_quotaparameter ordering may confuse API consumers.
file_quota(line 128) is declared aftervectordb(a DI dependency, line 127) in the function signature. While FastAPI handles this correctly, placing all user-facing form parameters together improves readability.♻️ Suggested reorder
async def create_user( display_name: str | None = Form(None), external_user_id: str | None = Form(None), is_admin: bool = Form(False), + file_quota: int | None = Form(None), vectordb=Depends(get_vectordb), - file_quota: int | None = Form(None), admin_user=Depends(require_admin), ):openrag/routers/partition.py (1)
455-475: Unusedrequestparameter.The
request: Requestparameter is injected but never referenced in the function body. Same applies toget_related_filesabove. Consider removing them to keep the signature clean, unless they're needed for future use.Proposed fix
async def get_file_ancestors( - request: Request, partition: str, file_id: str, vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ):openrag/components/retriever.py (1)
206-247: Consider parallelizing expansion fetches for lower latency.Related and ancestor chunks are fetched sequentially in two loops. When there are multiple unique
relationship_ids orfile_ids from the initial results, this could add noticeable latency. Usingasyncio.gather(withreturn_exceptions=True) would allow concurrent fetches.openrag/components/test_relationships.py (1)
41-49: Testto_dict()structure diverges from productionFile.to_dict().The test helper nests metadata under
"file_metadata"key, while the productionFile.to_dict()inopenrag/components/indexer/vectordb/utils.py(lines 69–78) spreads metadata into the top-level dict with**metadata. This means test assertions likefile_dict["file_metadata"]["filename"](line 457) wouldn't catch regressions where production flattens these fields differently.Not blocking since these tests focus on relationship fields, but be aware of this when extending coverage.
openrag/routers/search.py (1)
17-99: Sequential remote calls in loops could slow expansion significantly.Lines 57-75 and 78-97 each
awaita Ray remote call per unique relationship/file. With many distinct relationship IDs or file IDs in the result set, this becomes a waterfall of serial RPCs. Consider parallelizing withasyncio.gather.♻️ Proposed refactor — parallelize remote calls
# Fetch related chunks by relationship_id if include_related: + tasks = [] for partition, rel_id in relationship_ids: if partition and rel_id: - try: - related_chunks = await vectordb.get_related_chunks.remote( + tasks.append(vectordb.get_related_chunks.remote( partition=partition, relationship_id=rel_id, limit=related_limit, - ) - for chunk in related_chunks: - chunk_id = chunk.metadata.get("_id") - if chunk_id and chunk_id not in seen_ids: - seen_ids.add(chunk_id) - expanded_results.append(chunk) - except Exception as e: - logger.warning( - "Failed to fetch related chunks", - relationship_id=rel_id, - error=str(e), - ) + )) + related_results = await asyncio.gather(*tasks, return_exceptions=True) + for result in related_results: + if isinstance(result, Exception): + logger.warning("Failed to fetch related chunks", error=str(result)) + continue + for chunk in result: + chunk_id = chunk.metadata.get("_id") + if chunk_id and chunk_id not in seen_ids: + seen_ids.add(chunk_id) + expanded_results.append(chunk)Apply the same pattern to the ancestor chunk fetching loop.
openrag/components/indexer/vectordb/utils.py (1)
33-35: Module-levelload_config()call — will fail at import time if config is unavailable.
config = load_config()andDEFAULT_FILE_QUOTA = config.rdb.get(...)execute at module import time. If the config subsystem isn't initialized (e.g., in isolated unit tests), importing this module will raise an exception. This may explain why the relationship integration tests usepytest.skipas a workaround.Consider lazy-loading the config value or providing a fallback, similar to how other modules handle this pattern.
| import pytest | ||
| from unittest.mock import AsyncMock, MagicMock, patch | ||
| from langchain_core.documents.base import Document |
There was a problem hiding this comment.
Fix import ordering to unblock the pipeline.
The linting pipeline is failing due to unsorted imports (Ruff I001). pytest should come after unittest.mock in standard import ordering, and langchain_core is a third-party import that should be in a separate group.
🔧 Proposed fix
-import pytest
from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
from langchain_core.documents.base import Document📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import pytest | |
| from unittest.mock import AsyncMock, MagicMock, patch | |
| from langchain_core.documents.base import Document | |
| from unittest.mock import AsyncMock, MagicMock, patch | |
| from langchain_core.documents.base import Document | |
| import pytest |
🧰 Tools
🪛 GitHub Actions: Linting
[error] 12-12: Ruff check failed: Import block is un-sorted or un-formatted.
🪛 GitHub Check: lint (3.12)
[failure] 12-14: Ruff (I001)
openrag/tests/test_relationships_integration.py:12:1: I001 Import block is un-sorted or un-formatted
🤖 Prompt for AI Agents
In `@openrag/tests/test_relationships_integration.py` around lines 12 - 14,
Reorder and group the imports into standard-library, third-party, and local
groups with a blank line between groups; specifically move "from unittest.mock
import AsyncMock, MagicMock, patch" above "import pytest" and keep "from
langchain_core.documents.base import Document" in the third-party group below a
blank line, so the import order becomes: standard-library (unittest.mock...),
third-party (pytest, langchain_core...), ensuring imports reference AsyncMock,
MagicMock, patch, pytest, and Document as shown.
8b24f46 to
64c6801
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openrag/components/indexer/vectordb/vectordb.py (1)
898-905:⚠️ Potential issue | 🔴 Critical
delete_usercallsdelete_partitionwithout the requireduser_idargument — will raiseTypeErrorat runtime.
delete_partitioninPartitionFileManager(utils.py line 301) now requiresuser_id: intas a positional parameter (no default value). But line 904 calls it withoutuser_id, which will crash when an admin tries to delete a user who owns partitions.Proposed fix
async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: - self.partition_file_manager.delete_partition(partition) + self.partition_file_manager.delete_partition(partition, user_id=user_id) self.partition_file_manager.delete_user(user_id)
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/data_model.md`:
- Around line 103-111: The docs claim DEFAULT_FILE_QUOTA defaults to -1 but the
code sets DEFAULT_FILE_QUOTA via config.rdb.get("default_file_quota", 0) in
openrag/routers/utils.py and openrag/components/indexer/vectordb/utils.py;
update either the docs or the code so they match: change the fallback in both
modules to use -1 (i.e., set DEFAULT_FILE_QUOTA =
config.rdb.get("default_file_quota", -1)) or change the documentation text to
state the code-default is 0; ensure you only modify the DEFAULT_FILE_QUOTA
default value in the two referenced modules (DEFAULT_FILE_QUOTA) or the
documentation paragraph to keep behavior and docs consistent.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 285-290: The decrement of User.file_count is vulnerable to race
conditions because it reads the Python-cached value then writes a SQL-level
value; replace the current check-and-assign with an atomic SQL expression using
func.greatest to floor at zero (consistent with delete_partition). For example,
perform an update via session.query(User).filter(User.id ==
user_id).update({User.file_count: func.greatest(User.file_count - 1, 0)}) (and
ensure func is imported from sqlalchemy) so the decrement is done server-side
and cannot go negative under concurrency.
- Around line 682-694: Import the json module and parse row.file_metadata (which
can be a JSON string when using session.execute/text) before unpacking it with
**; inside the list comprehension that builds the returned dict (the code around
session.execute(...) and the dict with keys
"file_id","partition","parent_id","relationship_id","depth"), detect if
row.file_metadata is a str and call json.loads(row.file_metadata) (or use {}
when null/empty) and then merge that parsed dict into the returned dict to avoid
the TypeError when unpacking a raw JSON string.
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 917-919: The return type annotation for update_user_quota is
incorrect: update_user_quota currently declares -> bool but delegates to
partition_file_manager.update_user_quota which returns a dict of user fields;
update the signature of update_user_quota to return a dict (e.g., -> dict or ->
dict[str, Any] and add Any import if using typing) and adjust any callers or
tests expecting a bool accordingly; reference the update_user_quota function in
vectordb and the partition_file_manager.update_user_quota implementation when
making the change.
🧹 Nitpick comments (1)
openrag/routers/users.py (1)
60-93: Quota resolution logic is duplicated between here andcheck_user_file_quota.The quota precedence logic (admin → kill-switch → per-user → default fallback) in
get_current_user_info(lines 66–75) is a near-copy of the same logic incheck_user_file_quota(utils.py lines 202–216). Consider extracting a shared helper likeresolve_user_quota(user) -> int | floatto avoid the two implementations drifting apart.
|
|
||
| # Decrement file_count for the user | ||
| user = session.query(User).filter(User.id == user_id).first() | ||
| if user and user.file_count > 0: | ||
| user.file_count = User.file_count - 1 | ||
|
|
There was a problem hiding this comment.
file_count decrement is not fully guarded against going negative under concurrency.
Line 288 checks user.file_count > 0 using the Python-cached value, but the actual decrement on line 289 is SQL-level. Under concurrent deletes, two sessions could both read file_count == 1, both pass the check, and both decrement — resulting in file_count == -1. Use func.greatest for consistency with delete_partition (line 313).
Proposed fix
# Decrement file_count for the user
user = session.query(User).filter(User.id == user_id).first()
- if user and user.file_count > 0:
- user.file_count = User.file_count - 1
+ if user:
+ user.file_count = func.greatest(0, User.file_count - 1)🤖 Prompt for AI Agents
In `@openrag/components/indexer/vectordb/utils.py` around lines 285 - 290, The
decrement of User.file_count is vulnerable to race conditions because it reads
the Python-cached value then writes a SQL-level value; replace the current
check-and-assign with an atomic SQL expression using func.greatest to floor at
zero (consistent with delete_partition). For example, perform an update via
session.query(User).filter(User.id == user_id).update({User.file_count:
func.greatest(User.file_count - 1, 0)}) (and ensure func is imported from
sqlalchemy) so the decrement is done server-side and cannot go negative under
concurrency.
| result = session.execute(query, {"file_id": file_id, "partition": partition}) | ||
|
|
||
| return [ | ||
| { | ||
| "file_id": row.file_id, | ||
| "partition": row.partition_name, | ||
| "parent_id": row.parent_id, | ||
| "relationship_id": row.relationship_id, | ||
| "depth": row.depth, | ||
| **(row.file_metadata or {}), | ||
| } | ||
| for row in result | ||
| ] |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, check if json is imported in utils.py
grep -n "^import json" openrag/components/indexer/vectordb/utils.pyRepository: linagora/openrag
Length of output: 42
🏁 Script executed:
# Get the full context around lines 682-694 to see the actual code
sed -n '675,700p' openrag/components/indexer/vectordb/utils.pyRepository: linagora/openrag
Length of output: 1029
🏁 Script executed:
# Check the imports at the top of utils.py
head -50 openrag/components/indexer/vectordb/utils.pyRepository: linagora/openrag
Length of output: 1189
🏁 Script executed:
# Look at test_relationships.py line 131 to verify the claim
sed -n '125,140p' openrag/components/indexer/vectordb/tests/test_relationships.pyRepository: linagora/openrag
Length of output: 166
🏁 Script executed:
# Search for test_relationships.py
fd test_relationships.pyRepository: linagora/openrag
Length of output: 100
🏁 Script executed:
# Check if json is imported anywhere in utils.py (full file)
grep -n "json" openrag/components/indexer/vectordb/utils.pyRepository: linagora/openrag
Length of output: 42
🏁 Script executed:
# Check the query definition to see how file_metadata is constructed
sed -n '660,694p' openrag/components/indexer/vectordb/utils.pyRepository: linagora/openrag
Length of output: 1516
🏁 Script executed:
# Check test_relationships.py around line 131
sed -n '125,145p' openrag/components/test_relationships.pyRepository: linagora/openrag
Length of output: 860
🏁 Script executed:
# Also check the File model to see how file_metadata column is defined
grep -n "file_metadata" openrag/components/indexer/vectordb/utils.py | head -20Repository: linagora/openrag
Length of output: 546
Add json import and parse file_metadata before unpacking.
When using session.execute(text(...)), SQLAlchemy doesn't infer column types, so file_metadata (a JSON column) is returned as a raw string by the database driver rather than a parsed dict. Attempting to unpack it with ** on line 691 will raise TypeError. The test file (test_relationships.py, line 131) uses the same pattern with json.loads() to handle this correctly.
Proposed fix
+import json
+
...
return [
{
"file_id": row.file_id,
"partition": row.partition_name,
"parent_id": row.parent_id,
"relationship_id": row.relationship_id,
"depth": row.depth,
- **(row.file_metadata or {}),
+ **(json.loads(row.file_metadata) if row.file_metadata else {}),
}
for row in result
]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| result = session.execute(query, {"file_id": file_id, "partition": partition}) | |
| return [ | |
| { | |
| "file_id": row.file_id, | |
| "partition": row.partition_name, | |
| "parent_id": row.parent_id, | |
| "relationship_id": row.relationship_id, | |
| "depth": row.depth, | |
| **(row.file_metadata or {}), | |
| } | |
| for row in result | |
| ] | |
| import json | |
| result = session.execute(query, {"file_id": file_id, "partition": partition}) | |
| return [ | |
| { | |
| "file_id": row.file_id, | |
| "partition": row.partition_name, | |
| "parent_id": row.parent_id, | |
| "relationship_id": row.relationship_id, | |
| "depth": row.depth, | |
| **(json.loads(row.file_metadata) if row.file_metadata else {}), | |
| } | |
| for row in result | |
| ] |
🤖 Prompt for AI Agents
In `@openrag/components/indexer/vectordb/utils.py` around lines 682 - 694, Import
the json module and parse row.file_metadata (which can be a JSON string when
using session.execute/text) before unpacking it with **; inside the list
comprehension that builds the returned dict (the code around
session.execute(...) and the dict with keys
"file_id","partition","parent_id","relationship_id","depth"), detect if
row.file_metadata is a str and call json.loads(row.file_metadata) (or use {}
when null/empty) and then merge that parsed dict into the returned dict to avoid
the TypeError when unpacking a raw JSON string.
f32a983 to
e4a777e
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
openrag/components/indexer/indexer.py (1)
164-189:⚠️ Potential issue | 🟡 Minor
update_file_metadatastill declaresuser: dict | None = Nonebutdelete_filerequiresuser: dict.Line 170 allows
userto beNone, but Line 183 passes it directly toself.delete_file(...)which requires a non-optionaldict(Line 149). While current API callers always supply a user, the type mismatch could cause anAttributeErrorif this method is ever called without a user.Proposed fix — tighten the type hint
async def update_file_metadata( self, file_id: str, metadata: dict, partition: str, - user: dict | None = None, + user: dict, ):openrag/components/indexer/vectordb/utils.py (1)
243-250:⚠️ Potential issue | 🟠 Major
PartitionMembershipcreated with potentiallyNoneuser_id.
user_idis typed asint | None = None(line 219), but when called fromvectordb.py:387withuser.get("id"), it can beNoneif the "id" key is missing. If the partition doesn't exist yet, line 249 creates aPartitionMembership(user_id=None). Theuser_idcolumn onPartitionMembershiphas aNOT NULLconstraint (line 133), so this will raise anIntegrityErrorat commit time. Either validateuser_idearly or guard this block.Proposed guard
if not partition_obj: partition_obj = Partition(partition=partition) session.add(partition_obj) log.info("Created new partition") - membership = PartitionMembership(partition_name=partition, user_id=user_id, role="owner") - session.add(membership) + if user_id is not None: + membership = PartitionMembership(partition_name=partition, user_id=user_id, role="owner") + session.add(membership) + else: + log.warning("No user_id provided; partition created without owner membership")
🤖 Fix all issues with AI agents
In `@openrag/components/indexer/indexer.py`:
- Around line 148-157: The call to vectordb.delete_file in indexer.delete_file
passes a user_id keyword (user.get("id")) but vectordb.delete_file only accepts
(self, file_id: str, partition: str), causing a TypeError; fix by either
removing the user_id kwarg from the call in indexer.delete_file (the call site
in delete_file) or update the vectordb.delete_file method signature to accept
(file_id, partition, user_id=None) and handle or log the user_id there—pick the
approach consistent with access control expectations and update the
corresponding function (indexer.delete_file or Vectordb.delete_file)
accordingly.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 71-80: In to_dict, avoid metadata silently overwriting explicit
columns by either spreading metadata first or filtering out reserved keys:
change the dict construction in to_dict so metadata is applied before the
explicit keys (e.g., d = {**(self.file_metadata or {}), "partition":
self.partition_name, "file_id": self.file_id, "relationship_id":
self.relationship_id, "parent_id": self.parent_id}) or filter self.file_metadata
to remove reserved keys ("partition","file_id","relationship_id","parent_id")
before merging; update the to_dict method accordingly.
In `@openrag/routers/utils.py`:
- Around line 184-187: The check_user_file_quota function currently injects an
unused dependency vectordb via Depends(get_vectordb); remove the vectordb
parameter from the function signature (leave user=Depends(current_user) intact)
so the get_vectordb dependency is not resolved unnecessarily, and then remove
any now-unused import or reference to get_vectordb if it becomes unused
elsewhere.
🧹 Nitpick comments (1)
openrag/routers/users.py (1)
9-9: Consider makingtask_state_managera FastAPI dependency instead of a module-level singleton.
get_vectordbis injected viaDepends()(lazy), buttask_state_manageris eagerly resolved at import time. If Ray isn't initialized when this module loads (e.g., during unit tests or import-time errors), this will fail. Wrapping it in a dependency (or alru_cache-backed getter used viaDepends) would make it consistent withget_vectordband easier to mock in tests.
| def to_dict(self): | ||
| metadata = self.file_metadata or {} | ||
| d = {"partition": self.partition_name, "file_id": self.file_id, **metadata} | ||
| d = { | ||
| "partition": self.partition_name, | ||
| "file_id": self.file_id, | ||
| "relationship_id": self.relationship_id, | ||
| "parent_id": self.parent_id, | ||
| **metadata, | ||
| } | ||
| return d |
There was a problem hiding this comment.
Potential key collision: file_metadata could overwrite relationship_id or parent_id.
The **metadata spread on line 78 is applied after the explicit keys on lines 76-77. If file_metadata contains a "relationship_id" or "parent_id" key, it will silently overwrite the column values. Consider placing **metadata first or filtering out reserved keys.
Proposed fix — spread metadata first so column values take precedence
def to_dict(self):
metadata = self.file_metadata or {}
d = {
- "partition": self.partition_name,
- "file_id": self.file_id,
- "relationship_id": self.relationship_id,
- "parent_id": self.parent_id,
**metadata,
+ "partition": self.partition_name,
+ "file_id": self.file_id,
+ "relationship_id": self.relationship_id,
+ "parent_id": self.parent_id,
}
return d🤖 Prompt for AI Agents
In `@openrag/components/indexer/vectordb/utils.py` around lines 71 - 80, In
to_dict, avoid metadata silently overwriting explicit columns by either
spreading metadata first or filtering out reserved keys: change the dict
construction in to_dict so metadata is applied before the explicit keys (e.g., d
= {**(self.file_metadata or {}), "partition": self.partition_name, "file_id":
self.file_id, "relationship_id": self.relationship_id, "parent_id":
self.parent_id}) or filter self.file_metadata to remove reserved keys
("partition","file_id","relationship_id","parent_id") before merging; update the
to_dict method accordingly.
| async def check_user_file_quota( | ||
| user=Depends(current_user), | ||
| vectordb=Depends(get_vectordb), | ||
| ): |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cat -n openrag/routers/utils.py | sed -n '184,241p'Repository: linagora/openrag
Length of output: 2114
Remove unused vectordb parameter.
The vectordb parameter on line 186 is injected via Depends(get_vectordb) but never used in the function body. This unnecessarily forces the dependency to be resolved on every call.
Proposed fix
async def check_user_file_quota(
user=Depends(current_user),
- vectordb=Depends(get_vectordb),
):🤖 Prompt for AI Agents
In `@openrag/routers/utils.py` around lines 184 - 187, The check_user_file_quota
function currently injects an unused dependency vectordb via
Depends(get_vectordb); remove the vectordb parameter from the function signature
(leave user=Depends(current_user) intact) so the get_vectordb dependency is not
resolved unnecessarily, and then remove any now-unused import or reference to
get_vectordb if it becomes unused elsewhere.
e4a777e to
64457db
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
openrag/components/indexer/vectordb/vectordb.py (1)
878-885:⚠️ Potential issue | 🔴 Critical
delete_partitioncall missing requireduser_idargument — will raiseTypeErrorat runtime.
PartitionFileManager.delete_partition(partition, user_id)requiresuser_idas a positional argument (seeopenrag/components/indexer/vectordb/utils.pylines 281–301), but the call on line 884 omits it. When deleting a user who owns partitions, this will crash.🐛 Proposed fix
async def delete_user(self, user_id: int): self._check_user_exists(user_id) user_partitions = [ p["partition"] for p in self.partition_file_manager.list_user_partitions(user_id) if p["role"] == "owner" ] for partition in user_partitions: - self.partition_file_manager.delete_partition(partition) + self.partition_file_manager.delete_partition(partition, user_id=user_id) self.partition_file_manager.delete_user(user_id)openrag/components/indexer/indexer.py (1)
164-189:⚠️ Potential issue | 🟡 Minor
update_file_metadatapassesusertodelete_file, butusercan beNone.
useris typeddict | None = Noneat line 170, butdelete_fileat line 149 expectsuser: dict(non-optional). IfuserisNone,user.get("id")insidedelete_filewill raiseAttributeError.Consider either making
userrequired here or adding a guard.
🤖 Fix all issues with AI agents
In
`@openrag/scripts/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py`:
- Around line 24-25: Add a data backfill step in the migration that, after
adding the file_count column (added via op.add_column with server_default="0"),
runs an UPDATE to set users.file_count from actual file counts (joining files,
partitions, partition_memberships and filtering pm.role IN ('owner','editor') or
whatever ownership logic your app uses) so existing users reflect true usage;
then remove the server_default on file_count (alter the column to drop the
default) so future inserts behave normally. Also mention check_user_file_quota
and the usage pattern user.get("file_count", 0) so reviewers know this backfill
is to keep quota enforcement consistent.
In `@tests/api_tests/test_indexer.py`:
- Around line 458-518: The test test_quota_limit_blocks_excess_uploads can fail
if quota enforcement is globally disabled (DEFAULT_FILE_QUOTA <= 0); make the
test self-contained by asserting or ensuring quota enforcement is enabled before
exercising limits: either add a guard at the top of
test_quota_limit_blocks_excess_uploads that checks the server's effective
DEFAULT_FILE_QUOTA (or an endpoint/flag) and skips/fails the test if quotas are
disabled, or explicitly set the user's quota via the PATCH /users/{id}/quota (or
use the helper _create_user_with_quota and then call the PATCH) to a positive
value and verify the server acknowledged it before uploading files; reference
the test name test_quota_limit_blocks_excess_uploads and the PATCH
/users/{id}/quota (or _create_user_with_quota) to locate where to add the guard
or quota-enforcing call.
🧹 Nitpick comments (3)
openrag/components/indexer/vectordb/vectordb.py (1)
52-62: Abstract base class signatures are out of sync with theMilvusDBoverrides.
BaseVectorDB.delete_partition(line 53) andBaseVectorDB.delete_file(line 61) still lack theuser_idparameter that theMilvusDBimplementations now require. This breaks the Liskov substitution principle and will confuse anyone coding against the abstract interface.♻️ Proposed fix
`@abstractmethod` - async def delete_partition(self, partition: str): + async def delete_partition(self, partition: str, user_id: int): pass `@abstractmethod` def list_partition_files(self, partition: str, limit: int | None = None): pass `@abstractmethod` - async def delete_file(self, file_id: str, partition: str): + async def delete_file(self, file_id: str, partition: str, user_id: int): passdocs/content/docs/documentation/data_model.md (1)
22-23: New columns documented correctly; minor formatting nit.The
file_quotaandfile_countcolumn documentation matches the migration schema. However, line 23 is missing a space before the|separator after the column name (compare with other rows).-| `file_count`| Integer (default=0) | Number of uploaded files.| +| `file_count` | Integer (default=0) | Number of uploaded files. |openrag/routers/users.py (1)
60-93: Quota resolution logic is duplicated between here andcheck_user_file_quota.The quota determination logic (admin → kill-switch → per-user → default fallback) on lines 66–75 is nearly identical to lines 200–215 of
openrag/routers/utils.py. If the semantics evolve (e.g., new quota tiers), both must be updated in lockstep.Consider extracting a shared helper like
resolve_user_quota(user, default_file_quota) -> int | floatthat both sites call.
| op.add_column("users", sa.Column("file_quota", sa.Integer(), nullable=True)) | ||
| op.add_column("users", sa.Column("file_count", sa.Integer(), nullable=False, server_default="0")) |
There was a problem hiding this comment.
file_count will be 0 for all existing users after migration — no backfill for existing files.
Existing users who already have indexed files will have file_count = 0 after this migration. Since the quota enforcement relies on file_count to track usage (indexed_count = user.get("file_count", 0) in check_user_file_quota), these users will appear to have no files and could upload beyond their actual usage until the count is reconciled.
Consider adding a data migration step that backfills file_count from the actual file counts, e.g.:
UPDATE users SET file_count = (
SELECT COUNT(*) FROM files f
JOIN partitions p ON f.partition_name = p.partition
JOIN partition_memberships pm ON pm.partition_name = p.partition
WHERE pm.user_id = users.id AND pm.role IN ('owner', 'editor')
);The exact query depends on how file ownership is determined in your system. Without this, the quota system will be inaccurate for pre-existing users.
🤖 Prompt for AI Agents
In
`@openrag/scripts/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.py`
around lines 24 - 25, Add a data backfill step in the migration that, after
adding the file_count column (added via op.add_column with server_default="0"),
runs an UPDATE to set users.file_count from actual file counts (joining files,
partitions, partition_memberships and filtering pm.role IN ('owner','editor') or
whatever ownership logic your app uses) so existing users reflect true usage;
then remove the server_default on file_count (alter the column to drop the
default) so future inserts behave normally. Also mention check_user_file_quota
and the usage pattern user.get("file_count", 0) so reviewers know this backfill
is to keep quota enforcement consistent.
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
3 similar comments
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
62369d6 to
f124e97
Compare
requires database migration to add file_quota column to users table.
Features:
- Add file_quota field to users table (nullable integer)
- Add DEFAULT_FILE_QUOTA env var to set global default quota
- Add PATCH /users/{user_id}/quota endpoint to update user quotas
- Add quota enforcement on file upload (indexed files + pending tasks)
- Admins bypass quota checks; quota<=0 for a user means unlimited; quota > 0 to limit user's quota
API Changes:
- POST /users/ now accepts optional file_quota parameter
- GET /users/ and GET /users/{id} return file_quota in response
- New PATCH /users/{user_id}/quota endpoint for quota management
- GET /users/info now renders additional fields: file_count, pending_files, total_files and file_quota
Tests:
- test_update_user_quota: verify quota update from 10 to 12
- test_user_default_quota: verify None quota defaults to 10
- TestUserQuotaEnforcement class in test_indexer.py:
- test_unlimited_quota_user_can_exceed_default: user with quota=0 uploads 11 files
- test_quota_limit_blocks_excess_uploads: user with quota=5 blocked on 6th file
Perform migration: Refer to the doc
https://github.com/linagora/openrag/blob/0b5f84cd880e7c75db6f78d12a7681227721c2ad/docs/content/docs/documentation/sql_migration.mdx?plain=1#L47-L53
d75b5fe to
b46a7af
Compare
b46a7af to
546d73e
Compare
546d73e to
87ba4e4
Compare
BREAKING CHANGE: Introduces DEFAULT_FILE_QUOTA environment variable and requires database migration to add file_quota column to users table.
Features:
API Changes:
Tests:
Perform migration: Refer to the doc
openrag/docs/content/docs/documentation/sql_migration.mdx
Lines 47 to 53 in 0b5f84c
Frontend:
IndexerUI has been updated to show users their current quota usage and enforce quota limit. See this PR. One has to manually update the submodule in order for Indexer UI to point to the newest version.
Summary by CodeRabbit
New Features
Documentation
Database Migration
Tests