Add workspaces - #272
Merged
Merged
Conversation
|
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 workspace support: DB tables + migration, ORM models, FastAPI workspace router/endpoints, workspace-file associations and lifecycle in vector DB and indexer, workspace-scoped retrieval/filtering, docs and tests, plus a minor .gitignore tweak. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client as Client
participant SearchAPI as Search Endpoint
participant Retriever as Retriever
participant VectorDB as MilvusDB
participant Postgres as PostgreSQL
participant Milvus as Milvus
Client->>SearchAPI: GET /search?workspace=ws123
SearchAPI->>Retriever: retrieve(partitions, query, filter={workspace_id:"ws123"})
Retriever->>VectorDB: async_search(partitions, query, filter)
alt filter contains workspace_id
VectorDB->>Postgres: SELECT file_id FROM workspace_files WHERE workspace_id='ws123'
Postgres-->>VectorDB: [file_id1, file_id2, ...]
VectorDB->>Milvus: async_search(partitions, query, expr: file_id IN [...])
else no workspace filter
VectorDB->>Milvus: async_search(partitions, query)
end
Milvus-->>VectorDB: [matching vectors]
VectorDB-->>Retriever: [Documents]
Retriever-->>SearchAPI: [Documents]
SearchAPI-->>Client: 200 OK [documents]
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
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 |
coderabbitai
Bot
added
breaking-change
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 (1)
openrag/components/retriever.py (1)
176-187:⚠️ Potential issue | 🟠 MajorDuplicate
super().__init__call.
HyDeRetriever.__init__callssuper().__init__twice - once at line 176-185 with full parameters, and again at line 187 with onlytop_kandsimilarity_threshold. The second call overwrites instance attributes set by the first call.Proposed fix - remove duplicate call
super().__init__( top_k, similarity_threshold, with_surrounding_chunks, include_related, include_ancestors, related_limit, max_ancestor_depth, **kwargs, ) - super().__init__(top_k, similarity_threshold, **kwargs) if llm is None: raise ValueError("llm must be provided for HyDeRetriever")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/retriever.py` around lines 176 - 187, HyDeRetriever.__init__ currently calls super().__init__ twice which causes the second call to overwrite attributes set by the first; remove the duplicate invocation (the one that only passes top_k and similarity_threshold) and keep the initial super().__init__ call that forwards all parameters (top_k, similarity_threshold, with_surrounding_chunks, include_related, include_ancestors, related_limit, max_ancestor_depth, **kwargs) so the parent class is initialized once with the complete argument set.
🧹 Nitpick comments (9)
tests/api_tests/test_workspaces.py (2)
76-86: Tests add non-existent files to workspaces.These tests add files (
file-a,file-b) to workspaces without first creating/indexing actual files in the partition. While this validates the current implementation, it also demonstrates the lack of file existence validation discussed earlier.Consider adding tests that:
- Upload actual files and associate them with workspaces
- Verify workspace-scoped search returns those files
- Test error handling when associating non-existent files (if such validation is added)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/test_workspaces.py` around lines 76 - 86, Update the test_add_files_to_workspace test to create/upload real files before associating them with a workspace: call the API endpoint(s) that index/upload files (use api_client.post to the partition file upload/index endpoint) to create file-a and file-b, then call the existing workspace files POST (the same path used in test_add_files_to_workspace) and assert the response includes those file_ids; additionally add an assertion that a workspace-scoped search endpoint (call the partition/{workspace_partition}/workspaces/{workspace_id}/search or relevant search function) returns those files, and add a separate test that attempts to associate a clearly non-existent file id and asserts the expected error response to cover error handling.
10-20: Consider using context manager or try-finally for cleanup.The fixture cleanup in lines 17-20 silently swallows all exceptions. While acceptable for test cleanup, consider logging the exception for debugging purposes:
🔧 Optional: Log cleanup failures
yield name try: api_client.delete(f"/partition/{name}") - except Exception: - pass + except Exception as e: + import logging + logging.debug(f"Cleanup of partition {name} failed: {e}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/test_workspaces.py` around lines 10 - 20, The workspace_partition fixture currently swallows all exceptions during cleanup; change it to use a try/finally (or contextmanager) pattern so the delete is always attempted and if it fails log the exception instead of silently passing; update the fixture function workspace_partition to wrap the yield in a try/finally and call api_client.delete(f"/partition/{name}") in the finally block, catching exceptions only to log the error (including exception details) rather than using a bare except: pass.openrag/components/pipeline.py (1)
177-183: Workspace filtering not implemented for completions endpoint.The
_prepare_for_chat_completionmethod extractsworkspacefrom metadata and applies filtering (lines 138, 149-152), but_prepare_for_completionsdoes not. This means workspace-scoped search only works for/v1/chat/completions, not/v1/completions.If this is intentional, consider documenting the limitation. Otherwise, apply the same workspace filtering pattern:
♻️ Add workspace filtering to completions
async def _prepare_for_completions(self, partition: list[str], payload: dict): prompt = payload["prompt"] + metadata = payload.get("metadata", {}) + workspace = metadata.get("workspace") # 1. get the query query = await self.generate_query(messages=[{"role": "user", "content": prompt}]) # 2. get docs - docs = await self.retriever_pipeline.retrieve_docs(partition=partition, query=query) + filter_dict = {"workspace_id": workspace} if workspace else None + docs = await self.retriever_pipeline.retrieve_docs(partition=partition, query=query, filter=filter_dict)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/pipeline.py` around lines 177 - 183, _prepare_for_completions lacks the workspace metadata extraction and filtering applied in _prepare_for_chat_completion, causing workspace-scoped searches to be ignored for the /v1/completions path; fix it by mirroring the same pattern from _prepare_for_chat_completion: extract workspace (e.g. workspace = payload.get("metadata", {}).get("workspace") or similar) from the incoming payload, and when calling self.retriever_pipeline.retrieve_docs(partition=..., query=...), pass or apply the workspace filter (or include workspace in the retrieve call or filter the returned docs) so completions honor workspace scoping just like _prepare_for_chat_completion does.openrag/routers/search.py (1)
1-5: Use absolute imports from theopenrag/directory.Per coding guidelines, imports should use absolute paths from the
openrag/directory root.Proposed fix
-from components.retriever import _expand_with_related_chunks +from openrag.components.retriever import _expand_with_related_chunks from fastapi import APIRouter, Depends, Query, Request, status from fastapi.responses import JSONResponse -from utils.dependencies import get_indexer, get_vectordb -from utils.logger import get_logger +from openrag.utils.dependencies import get_indexer, get_vectordb +from openrag.utils.logger import get_loggerAs per coding guidelines: "Use absolute imports from the
openrag/directory (which is the Python path root)".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/search.py` around lines 1 - 5, Change the relative-style imports at the top of search.py to use absolute imports from the openrag package root: replace components.retriever import _expand_with_related_chunks with openrag.components.retriever import _expand_with_related_chunks, and replace utils.dependencies.get_indexer/get_vectordb and utils.logger.get_logger with openrag.utils.dependencies.get_indexer/get_vectordb and openrag.utils.logger.get_logger respectively; keep third-party imports (fastapi, fastapi.responses) unchanged. Ensure the module names (_expand_with_related_chunks, get_indexer, get_vectordb, get_logger) remain the same so callers in this file (e.g., any uses of _expand_with_related_chunks, get_indexer, get_vectordb, get_logger) continue to work.openrag/components/indexer/vectordb/utils.py (2)
767-773: Consider handling duplicate workspace_id gracefully.
create_workspacewill raise anIntegrityErrorif a workspace with the sameworkspace_idalready exists (due to the unique constraint). Consider catching this and raising a more descriptive custom exception, or checking for existence first.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 767 - 773, create_workspace currently commits a Workspace and will raise an IntegrityError on duplicate workspace_id; update it to handle duplicates gracefully by either (a) querying Session for existing Workspace(workspace_id=workspace_id, partition_name=partition) and returning or raising a custom DuplicateWorkspaceError, or (b) wrapping the session.commit() in a try/except that catches sqlalchemy.exc.IntegrityError, calls session.rollback(), and raises a clearer custom exception (e.g., DuplicateWorkspaceError) with a descriptive message; reference the create_workspace function, Workspace model, Session context, and the new DuplicateWorkspaceError when implementing the change.
813-817: Consider handling duplicate file additions gracefully.
add_files_to_workspacewill raise anIntegrityErrorif a file is already in the workspace (due toUniqueConstraint("workspace_id", "file_id")). Consider using an upsert pattern or checking existence first to handle idempotent requests.Potential approach using INSERT...ON CONFLICT DO NOTHING
def add_files_to_workspace(self, workspace_id: str, file_ids: list[str]): with self.Session() as session: for fid in file_ids: stmt = insert(WorkspaceFile).values( workspace_id=workspace_id, file_id=fid ).on_conflict_do_nothing( index_elements=['workspace_id', 'file_id'] ) session.execute(stmt) session.commit()Note: Requires importing
insertfromsqlalchemy.dialects.postgresqlfor PostgreSQL-specific ON CONFLICT syntax.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 813 - 817, The add_files_to_workspace method currently inserts WorkspaceFile rows and will raise IntegrityError on duplicate (UniqueConstraint workspace_id,file_id); change it to perform an idempotent insert (upsert) or ignore duplicates: inside add_files_to_workspace replace the plain session.add loop with a per-file INSERT...ON CONFLICT DO NOTHING statement (use insert from sqlalchemy.dialects.postgresql and .on_conflict_do_nothing(index_elements=['workspace_id','file_id'])) or alternatively check existence before insert or catch IntegrityError around the insert and continue; ensure you still commit at the end and reference WorkspaceFile and add_files_to_workspace when making the change.openrag/routers/workspaces.py (2)
14-20: AddConfigDict(extra='allow')to Pydantic request models.Per coding guidelines, Pydantic request models in API endpoints must use
ConfigDict(extra='allow')to accept vendor-specific fields.Proposed fix
+from pydantic import BaseModel, ConfigDict + + class CreateWorkspaceRequest(BaseModel): + model_config = ConfigDict(extra="allow") workspace_id: str display_name: str | None = None class AddFilesRequest(BaseModel): + model_config = ConfigDict(extra="allow") file_ids: list[str]As per coding guidelines: "Pydantic request models in API endpoints and tests must use
ConfigDict(extra='allow')to accept vendor-specific fields likeextra_body".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 14 - 20, The CreateWorkspaceRequest and AddFilesRequest Pydantic models must allow vendor-specific fields; add a model_config = ConfigDict(extra='allow') to both classes (CreateWorkspaceRequest and AddFilesRequest) so extra keys like extra_body are accepted, and ensure ConfigDict is imported from pydantic if not already present.
7-9: Use absolute imports from theopenrag/directory.Proposed fix
-from utils.dependencies import get_vectordb +from openrag.utils.dependencies import get_vectordbAs per coding guidelines: "Use absolute imports from the
openrag/directory (which is the Python path root)".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 7 - 9, The imports in workspaces.py use non-absolute paths; replace "from utils.dependencies import get_vectordb" with an absolute import from the package root (use openrag.utils.dependencies and the symbol get_vectordb) and replace the relative "from .utils import require_partition_editor, require_partition_owner, require_partition_viewer" with an absolute import from the routers package (use openrag.routers.utils and the symbols require_partition_editor, require_partition_owner, require_partition_viewer) so all imports are rooted at openrag.openrag/components/indexer/vectordb/vectordb.py (1)
464-472: Performance impact with large workspaces. Thelist_workspace_files()call at line 468 returns all file IDs unbounded, and the constructed Milvus expression at line 472 creates a potentially largefile_id in [...]clause. For workspaces with thousands of files, this may degrade query performance. The same pattern appears inget_related_documents()(line 1042) andget_document_ancestors()(line 1081).Consider paginating the file ID list or implementing a maximum file count with appropriate error handling for large workspaces.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` around lines 464 - 472, The current use of partition_file_manager.list_workspace_files(...) to build a large "file_id in [...]" Milvus expression (in the block that mutates filter and appends to expr_parts) can produce huge IN clauses and hurt performance; update the code in vectordb.py (and the similar spots in get_related_documents and get_document_ancestors) to page or cap the file ID set: call list_workspace_files with pagination or a max_count, and if the total exceeds a safe threshold return an error/empty result or switch to an alternative query strategy (e.g., query by workspace_id metadata or batch queries) instead of constructing a giant IN list; ensure you reference and modify the code paths around partition_file_manager.list_workspace_files, the filter handling that pops "workspace_id", and the place that builds expr_parts with f"file_id in [{id_list}]" so downstream callers handle the new error/limit response.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/content/docs/documentation/data_model.md`:
- Around line 178-179: The docs claim that `file_id` is a foreign key, but the
`WorkspaceFile` class in openrag/components/indexer/vectordb/utils.py currently
defines `file_id` as a plain String; either remove the FK claim from docs or add
the FK to the model—preferably add referential integrity by changing
`WorkspaceFile.file_id` to use a SQLAlchemy ForeignKey referencing
`files.file_id` with ondelete='CASCADE' (and update any imports/constraints as
needed), or if you choose docs instead, update
docs/content/docs/documentation/data_model.md to remove the FK arrow/“CASCADE”
for `file_id` so it matches the implementation.
In `@docs/content/docs/documentation/workspaces.md`:
- Line 53: The docs contain inconsistent base paths: the intro sentence uses
"/indexer/partition/{partition}/workspaces" while the API table and examples use
"/partition/{partition}/workspaces"; pick the correct base path used by the
server (either with or without the "/indexer" prefix), then update every
occurrence so they match (update the header sentence, the API table paths,
examples, and any curl examples) — search for the strings
"/indexer/partition/{partition}/workspaces" and
"/partition/{partition}/workspaces" and replace with the confirmed canonical
path and ensure any descriptive text references that canonical path
consistently.
In `@openrag/routers/indexer.py`:
- Around line 176-184: The current code queues the indexing task via
indexer.add_file and immediately calls vectordb.add_files_to_workspace which can
create workspace_files pointing at a file that may never be created; change the
flow so workspace association is performed only after successful indexing:
modify the Indexer.add_file task (the add_file implementation) to accept
workspace_ids (or pass through parsed_workspace_ids) and perform the
vectordb.add_files_to_workspace logic inside Indexer.add_file after the file
record is persisted, or alternatively implement cleanup in the task failure path
by using task_state_manager to detect failure and removing any workspace_files
created; update callers that currently invoke vectordb.add_files_to_workspace
(the list comprehension using vectordb.add_files_to_workspace.remote) to stop
creating associations immediately and instead let Indexer.add_file handle it
post-insert.
---
Outside diff comments:
In `@openrag/components/retriever.py`:
- Around line 176-187: HyDeRetriever.__init__ currently calls super().__init__
twice which causes the second call to overwrite attributes set by the first;
remove the duplicate invocation (the one that only passes top_k and
similarity_threshold) and keep the initial super().__init__ call that forwards
all parameters (top_k, similarity_threshold, with_surrounding_chunks,
include_related, include_ancestors, related_limit, max_ancestor_depth, **kwargs)
so the parent class is initialized once with the complete argument set.
---
Nitpick comments:
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 767-773: create_workspace currently commits a Workspace and will
raise an IntegrityError on duplicate workspace_id; update it to handle
duplicates gracefully by either (a) querying Session for existing
Workspace(workspace_id=workspace_id, partition_name=partition) and returning or
raising a custom DuplicateWorkspaceError, or (b) wrapping the session.commit()
in a try/except that catches sqlalchemy.exc.IntegrityError, calls
session.rollback(), and raises a clearer custom exception (e.g.,
DuplicateWorkspaceError) with a descriptive message; reference the
create_workspace function, Workspace model, Session context, and the new
DuplicateWorkspaceError when implementing the change.
- Around line 813-817: The add_files_to_workspace method currently inserts
WorkspaceFile rows and will raise IntegrityError on duplicate (UniqueConstraint
workspace_id,file_id); change it to perform an idempotent insert (upsert) or
ignore duplicates: inside add_files_to_workspace replace the plain session.add
loop with a per-file INSERT...ON CONFLICT DO NOTHING statement (use insert from
sqlalchemy.dialects.postgresql and
.on_conflict_do_nothing(index_elements=['workspace_id','file_id'])) or
alternatively check existence before insert or catch IntegrityError around the
insert and continue; ensure you still commit at the end and reference
WorkspaceFile and add_files_to_workspace when making the change.
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 464-472: The current use of
partition_file_manager.list_workspace_files(...) to build a large "file_id in
[...]" Milvus expression (in the block that mutates filter and appends to
expr_parts) can produce huge IN clauses and hurt performance; update the code in
vectordb.py (and the similar spots in get_related_documents and
get_document_ancestors) to page or cap the file ID set: call
list_workspace_files with pagination or a max_count, and if the total exceeds a
safe threshold return an error/empty result or switch to an alternative query
strategy (e.g., query by workspace_id metadata or batch queries) instead of
constructing a giant IN list; ensure you reference and modify the code paths
around partition_file_manager.list_workspace_files, the filter handling that
pops "workspace_id", and the place that builds expr_parts with f"file_id in
[{id_list}]" so downstream callers handle the new error/limit response.
In `@openrag/components/pipeline.py`:
- Around line 177-183: _prepare_for_completions lacks the workspace metadata
extraction and filtering applied in _prepare_for_chat_completion, causing
workspace-scoped searches to be ignored for the /v1/completions path; fix it by
mirroring the same pattern from _prepare_for_chat_completion: extract workspace
(e.g. workspace = payload.get("metadata", {}).get("workspace") or similar) from
the incoming payload, and when calling
self.retriever_pipeline.retrieve_docs(partition=..., query=...), pass or apply
the workspace filter (or include workspace in the retrieve call or filter the
returned docs) so completions honor workspace scoping just like
_prepare_for_chat_completion does.
In `@openrag/routers/search.py`:
- Around line 1-5: Change the relative-style imports at the top of search.py to
use absolute imports from the openrag package root: replace components.retriever
import _expand_with_related_chunks with openrag.components.retriever import
_expand_with_related_chunks, and replace
utils.dependencies.get_indexer/get_vectordb and utils.logger.get_logger with
openrag.utils.dependencies.get_indexer/get_vectordb and
openrag.utils.logger.get_logger respectively; keep third-party imports (fastapi,
fastapi.responses) unchanged. Ensure the module names
(_expand_with_related_chunks, get_indexer, get_vectordb, get_logger) remain the
same so callers in this file (e.g., any uses of _expand_with_related_chunks,
get_indexer, get_vectordb, get_logger) continue to work.
In `@openrag/routers/workspaces.py`:
- Around line 14-20: The CreateWorkspaceRequest and AddFilesRequest Pydantic
models must allow vendor-specific fields; add a model_config =
ConfigDict(extra='allow') to both classes (CreateWorkspaceRequest and
AddFilesRequest) so extra keys like extra_body are accepted, and ensure
ConfigDict is imported from pydantic if not already present.
- Around line 7-9: The imports in workspaces.py use non-absolute paths; replace
"from utils.dependencies import get_vectordb" with an absolute import from the
package root (use openrag.utils.dependencies and the symbol get_vectordb) and
replace the relative "from .utils import require_partition_editor,
require_partition_owner, require_partition_viewer" with an absolute import from
the routers package (use openrag.routers.utils and the symbols
require_partition_editor, require_partition_owner, require_partition_viewer) so
all imports are rooted at openrag.
In `@tests/api_tests/test_workspaces.py`:
- Around line 76-86: Update the test_add_files_to_workspace test to
create/upload real files before associating them with a workspace: call the API
endpoint(s) that index/upload files (use api_client.post to the partition file
upload/index endpoint) to create file-a and file-b, then call the existing
workspace files POST (the same path used in test_add_files_to_workspace) and
assert the response includes those file_ids; additionally add an assertion that
a workspace-scoped search endpoint (call the
partition/{workspace_partition}/workspaces/{workspace_id}/search or relevant
search function) returns those files, and add a separate test that attempts to
associate a clearly non-existent file id and asserts the expected error response
to cover error handling.
- Around line 10-20: The workspace_partition fixture currently swallows all
exceptions during cleanup; change it to use a try/finally (or contextmanager)
pattern so the delete is always attempted and if it fails log the exception
instead of silently passing; update the fixture function workspace_partition to
wrap the yield in a try/finally and call api_client.delete(f"/partition/{name}")
in the finally block, catching exceptions only to log the error (including
exception details) rather than using a bare except: pass.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: df97cbab-6d41-4594-addf-ba8d73d5ecca
📒 Files selected for processing (15)
.gitignoreCLAUDE.mddocs/content/docs/documentation/data_model.mddocs/content/docs/documentation/sql_migration.mdxdocs/content/docs/documentation/workspaces.mdopenrag/api.pyopenrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/components/pipeline.pyopenrag/components/retriever.pyopenrag/routers/indexer.pyopenrag/routers/search.pyopenrag/routers/workspaces.pyopenrag/scripts/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.pytests/api_tests/test_workspaces.py
Closed
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
EnjoyBacon7
reviewed
Mar 11, 2026
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (3)
openrag/components/indexer/vectordb/utils.py (3)
820-824:⚠️ Potential issue | 🟠 MajorMake bulk workspace-file adds idempotent.
Re-adding an existing
(workspace_id, file_id)pair hitsuix_workspace_fileand aborts the whole request. For a bulk add endpoint this turns harmless retries into 500s; use conflict-safe inserts or filter existing memberships first.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 820 - 824, The bulk add in add_files_to_workspace currently inserts all (workspace_id, file_id) pairs and will fail on the unique constraint uix_workspace_file; update add_files_to_workspace to first query existing WorkspaceFile.file_id rows for the given workspace_id (using self.Session()/session and filter(WorkspaceFile.workspace_id==workspace_id, WorkspaceFile.file_id.in_(file_ids))) compute the set difference to find only new file_ids, then insert only those (either via session.add_all([WorkspaceFile(...) for fid in new_ids]) or session.bulk_save_objects) and commit; this makes the operation idempotent and avoids unique constraint violations.
767-774:⚠️ Potential issue | 🟠 MajorTranslate duplicate workspace IDs into a conflict instead of a raw DB error.
workspace_idis unique, but this method always inserts and commits. Re-creating an existing workspace will currently bubble anIntegrityErrorthrough the actor and show up as a 500 instead of a 409-style conflict.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 767 - 774, The create_workspace method currently always inserts and commits and so a duplicate Workspace.workspace_id will raise a raw DB IntegrityError; modify create_workspace (involving self.Session and Workspace) to catch sqlalchemy.exc.IntegrityError around session.commit(), roll back the session, and convert it into a conflict response (e.g., raise a dedicated ConflictError or HTTP 409-style exception used in the project) with a clear message about duplicate workspace_id; ensure other exceptions are re-raised after rollback so the session is clean.
176-188:⚠️ Potential issue | 🔴 CriticalWorkspace-file membership needs partition scope.
WorkspaceFileonly storesfile_id, butFileis unique on(file_id, partition_name). That makes these new workspace operations ambiguous across partitions:delete_workspace()can misclassify orphaned files, andremove_file_from_all_workspaces()can remove memberships for a different partition that happens to reuse the samefile_id. This needs eitherpartition_nameonworkspace_fileswith a partition-scoped key/FK, or every lookup/delete must join through the workspace’s partition before touching rows.Also applies to: 803-818, 843-847
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 176 - 188, WorkspaceFile lacks partition scoping so operations like delete_workspace() and remove_file_from_all_workspaces() can hit/clear rows for the same file_id in other partitions; fix by adding a partition_name column to WorkspaceFile and making the uniqueness key partition-scoped. Concretely, add Column(String, nullable=False, index=True, name="partition_name") to the WorkspaceFile model, change the UniqueConstraint to UniqueConstraint("workspace_id","file_id","partition_name", name=...), and ensure any FK/lookup logic (or add a composite FK if supported) references both file_id and partition_name (or always join to File on file_id+partition_name) so delete_workspace() and remove_file_from_all_workspaces() operate only within the correct partition.
🧹 Nitpick comments (1)
openrag/routers/workspaces.py (1)
7-10: Useopenrag.*imports in the new router module.This file mixes bare and relative imports (
utils...,.utils). Please switch them toopenrag.utils.../openrag.routers.utils...to match the repository import convention.As per coding guidelines,
Use absolute imports from the openrag/ directory (which is the Python path root).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 7 - 10, The imports in this module use bare/relative paths; update them to absolute openrag package imports: replace get_vectordb import from utils.dependencies with openrag.utils.dependencies.get_vectordb, replace get_logger from utils.logger with openrag.utils.logger.get_logger, and replace the relative .utils imports for require_partition_editor, require_partition_owner, and require_partition_viewer with openrag.routers.utils.require_partition_editor, openrag.routers.utils.require_partition_owner, and openrag.routers.utils.require_partition_viewer respectively so all imports use the openrag.* convention.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 465-472: The workspace_id translation currently only constrains
file_id but must also constrain partition_name: after popping workspace_id (in
the same block where filter is converted and file_ids are computed via
partition_file_manager.list_workspace_files), fetch the workspace metadata
(e.g., via the workspace lookup used elsewhere) to obtain its partition_name,
then 1) if the incoming filter/args include an explicit partition set that does
not contain that partition_name return [] immediately, otherwise append a
partition constraint to expr_parts like `partition_name == "<partition>"` (in
addition to the existing file_id IN clause) so Milvus queries are pinned to the
workspace partition. Ensure you use the same identifiers (filter, workspace_id,
partition_file_manager.list_workspace_files, expr_parts, partition_name) so the
change integrates with the existing expression building.
In `@openrag/routers/workspaces.py`:
- Around line 25-29: The workspace validation currently calls the Ray actor
method directly (require_workspace_in_partition uses
vectordb.get_workspace.remote(workspace_id)); replace direct .remote() awaits
with the centralized call_ray_actor_with_timeout() helper from
components.ray_utils so all Ray actor calls use the shared timeout/cancellation
logic; update require_workspace_in_partition and the other Ray calls in this
module (the other places invoking vectordb.get_workspace.remote or similar
.remote() calls) to call call_ray_actor_with_timeout(vectordb.get_workspace,
workspace_id) (or the equivalent method and args) and return/use its result
exactly where ws is used.
- Around line 16-22: Update the Pydantic request models CreateWorkspaceRequest
and AddFilesRequest to accept vendor-specific extras by adding a ConfigDict with
extra='allow' (e.g., define a Config attribute or use model_config =
ConfigDict(extra='allow') depending on Pydantic version) so the models no longer
reject fields like extra_body; locate the classes CreateWorkspaceRequest and
AddFilesRequest and add the model configuration referencing ConfigDict and
extra='allow' accordingly.
---
Duplicate comments:
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 820-824: The bulk add in add_files_to_workspace currently inserts
all (workspace_id, file_id) pairs and will fail on the unique constraint
uix_workspace_file; update add_files_to_workspace to first query existing
WorkspaceFile.file_id rows for the given workspace_id (using
self.Session()/session and filter(WorkspaceFile.workspace_id==workspace_id,
WorkspaceFile.file_id.in_(file_ids))) compute the set difference to find only
new file_ids, then insert only those (either via
session.add_all([WorkspaceFile(...) for fid in new_ids]) or
session.bulk_save_objects) and commit; this makes the operation idempotent and
avoids unique constraint violations.
- Around line 767-774: The create_workspace method currently always inserts and
commits and so a duplicate Workspace.workspace_id will raise a raw DB
IntegrityError; modify create_workspace (involving self.Session and Workspace)
to catch sqlalchemy.exc.IntegrityError around session.commit(), roll back the
session, and convert it into a conflict response (e.g., raise a dedicated
ConflictError or HTTP 409-style exception used in the project) with a clear
message about duplicate workspace_id; ensure other exceptions are re-raised
after rollback so the session is clean.
- Around line 176-188: WorkspaceFile lacks partition scoping so operations like
delete_workspace() and remove_file_from_all_workspaces() can hit/clear rows for
the same file_id in other partitions; fix by adding a partition_name column to
WorkspaceFile and making the uniqueness key partition-scoped. Concretely, add
Column(String, nullable=False, index=True, name="partition_name") to the
WorkspaceFile model, change the UniqueConstraint to
UniqueConstraint("workspace_id","file_id","partition_name", name=...), and
ensure any FK/lookup logic (or add a composite FK if supported) references both
file_id and partition_name (or always join to File on file_id+partition_name) so
delete_workspace() and remove_file_from_all_workspaces() operate only within the
correct partition.
---
Nitpick comments:
In `@openrag/routers/workspaces.py`:
- Around line 7-10: The imports in this module use bare/relative paths; update
them to absolute openrag package imports: replace get_vectordb import from
utils.dependencies with openrag.utils.dependencies.get_vectordb, replace
get_logger from utils.logger with openrag.utils.logger.get_logger, and replace
the relative .utils imports for require_partition_editor,
require_partition_owner, and require_partition_viewer with
openrag.routers.utils.require_partition_editor,
openrag.routers.utils.require_partition_owner, and
openrag.routers.utils.require_partition_viewer respectively so all imports use
the openrag.* convention.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ab7f5acf-b29e-48c4-81aa-3536c7e6fd76
📒 Files selected for processing (4)
openrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/routers/workspaces.pytests/api_tests/test_workspaces.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/api_tests/test_workspaces.py
There was a problem hiding this comment.
♻️ Duplicate comments (3)
openrag/components/indexer/vectordb/vectordb.py (2)
465-472:⚠️ Potential issue | 🟠 MajorPin workspace searches to the workspace’s partition, not just its file IDs.
A workspace is partition-owned, but this branch only adds
file_id in [...]. Sincefile_idis only unique per partition, apartition=["all"]search can still pull chunks from another accessible partition that reuses the same ID. Fetch the workspace metadata here, reject mismatched explicit partition filters, and always append a Milvuspartition == "<workspace partition>"clause alongside thefile_idfilter.Suggested fix
if filter: filter = dict(filter) # don't mutate caller's dict if "workspace_id" in filter: workspace_id = filter.pop("workspace_id") + workspace = self.partition_file_manager.get_workspace(workspace_id) + if not workspace: + return [] + workspace_partition = workspace["partition_name"] + if partition != ["all"] and workspace_partition not in partition: + return [] file_ids = self.partition_file_manager.list_workspace_files(workspace_id) if not file_ids: return [] # Empty workspace → no results + expr_parts.append(f'partition == "{workspace_partition}"') id_list = ", ".join(f'"{fid}"' for fid in file_ids) expr_parts.append(f"file_id in [{id_list}]")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` around lines 465 - 472, In the branch inside vectordb.py that handles "workspace_id" on the incoming filter (the block that calls self.partition_file_manager.list_workspace_files and appends to expr_parts), fetch the workspace metadata to get its partition id, then: (1) if the caller supplied an explicit partition filter ensure it matches the workspace partition and raise/reject on mismatch, (2) always append a Milvus partition clause like partition == "<workspace_partition>" to expr_parts in addition to the existing file_id in [...] clause, and (3) keep the existing defensive copy of filter and use the workspace partition value rather than assuming file_id uniqueness across partitions.
606-606:⚠️ Potential issue | 🔴 CriticalScope workspace cleanup by partition before deleting associations.
remove_file_from_all_workspaces()currently deletes byfile_idalone, so deletingfile_id="X"from partition A will also strip workspace memberships for partition B if that partition reused the same ID. This is a cross-partition data corruption path.Suggested call-site change
- self.partition_file_manager.remove_file_from_all_workspaces(file_id) + self.partition_file_manager.remove_file_from_all_workspaces(file_id, partition)The helper in
openrag/components/indexer/vectordb/utils.pyalso needs to filter throughWorkspace.partition_name == partitionbefore deleting. Based on learnings,files.file_idis not globally unique; workspace cleanup must stay partition-scoped.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` at line 606, The call to partition_file_manager.remove_file_from_all_workspaces(file_id) must be made partition-scoped to avoid cross-partition deletion: change the call site to pass the current partition (e.g., remove_file_from_all_workspaces(file_id, partition) or a new partition-scoped method) and update the helper in openrag/components/indexer/vectordb/utils.py to filter workspace deletions by Workspace.partition_name == partition before removing associations; ensure the helper and partition_file_manager method signatures are updated together and all call sites use the partition argument.docs/content/docs/documentation/data_model.md (1)
66-70:⚠️ Potential issue | 🟡 MinorRemove the
FKmarker fromworkspace_files.file_idin the ER diagram.Line 69 still documents
file_idas a foreign key, but the migration creates it as a plainStringwith no database FK. The prose below was already corrected, so the Mermaid schema is now the only remaining mismatch.Suggested doc fix
workspace_files { int id PK varchar workspace_id FK - varchar file_id FK + varchar file_id }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/content/docs/documentation/data_model.md` around lines 66 - 70, The ER diagram incorrectly marks workspace_files.file_id as a foreign key; update the Mermaid schema for the workspace_files table to remove the "FK" marker from file_id so it matches the migration and prose—edit the workspace_files block in the data_model document and change "varchar file_id FK" to "varchar file_id" (leaving workspace_files.id and workspace_id as-is).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@docs/content/docs/documentation/data_model.md`:
- Around line 66-70: The ER diagram incorrectly marks workspace_files.file_id as
a foreign key; update the Mermaid schema for the workspace_files table to remove
the "FK" marker from file_id so it matches the migration and prose—edit the
workspace_files block in the data_model document and change "varchar file_id FK"
to "varchar file_id" (leaving workspace_files.id and workspace_id as-is).
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 465-472: In the branch inside vectordb.py that handles
"workspace_id" on the incoming filter (the block that calls
self.partition_file_manager.list_workspace_files and appends to expr_parts),
fetch the workspace metadata to get its partition id, then: (1) if the caller
supplied an explicit partition filter ensure it matches the workspace partition
and raise/reject on mismatch, (2) always append a Milvus partition clause like
partition == "<workspace_partition>" to expr_parts in addition to the existing
file_id in [...] clause, and (3) keep the existing defensive copy of filter and
use the workspace partition value rather than assuming file_id uniqueness across
partitions.
- Line 606: The call to
partition_file_manager.remove_file_from_all_workspaces(file_id) must be made
partition-scoped to avoid cross-partition deletion: change the call site to pass
the current partition (e.g., remove_file_from_all_workspaces(file_id, partition)
or a new partition-scoped method) and update the helper in
openrag/components/indexer/vectordb/utils.py to filter workspace deletions by
Workspace.partition_name == partition before removing associations; ensure the
helper and partition_file_manager method signatures are updated together and all
call sites use the partition argument.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f60a8a1f-fa0d-4ca1-99b0-9ef781671d79
📒 Files selected for processing (3)
docs/content/docs/documentation/data_model.mddocs/content/docs/documentation/workspaces.mdopenrag/components/indexer/vectordb/vectordb.py
✅ Files skipped from review due to trivial changes (1)
- docs/content/docs/documentation/workspaces.md
There was a problem hiding this comment.
Actionable comments posted: 1
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)
598-607:⚠️ Potential issue | 🟡 MinorNon-atomic delete: RDB cleanup failures leave orphaned workspace_files rows.
The Milvus delete (lines 601-604) commits before
remove_file_from_all_workspacesandremove_file_from_partitionare called. If either RDB operation fails, the file is deleted from Milvus but workspace_files references remain, causing workspace file listings to show non-existent files.Consider wrapping the RDB operations in a try-except to log failures, or performing RDB cleanup first (before Milvus delete) so rollback is possible.
🛡️ Suggested improvement
try: res = await self._async_client.delete( collection_name=self.collection_name, filter=f"partition == '{partition}' and file_id == '{file_id}'", ) - self.partition_file_manager.remove_file_from_all_workspaces(file_id, partition) - self.partition_file_manager.remove_file_from_partition(file_id=file_id, partition=partition) + try: + self.partition_file_manager.remove_file_from_all_workspaces(file_id, partition) + self.partition_file_manager.remove_file_from_partition(file_id=file_id, partition=partition) + except Exception as rdb_err: + log.error("RDB cleanup failed after Milvus delete — manual cleanup may be needed", error=str(rdb_err)) + raise log.info("Deleted file chunks from partition.", count=res.get("delete_count", 0))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` around lines 598 - 607, The Milvus deletion in delete_file currently commits before the relational DB cleanup, risking orphaned workspace_files if remove_file_from_all_workspaces or remove_file_from_partition fails; change the flow in delete_file so you first perform the RDB cleanup via partition_file_manager.remove_file_from_all_workspaces(file_id, partition) and partition_file_manager.remove_file_from_partition(file_id=file_id, partition=partition) inside a try/except that ensures both succeed (log and raise on failure), and only after successful RDB operations call self._async_client.delete(collection_name=self.collection_name, filter=...) (or alternatively, if you must delete Milvus first, wrap the RDB calls in try/except and on failure attempt to restore the Milvus entry or surface a clear error); ensure failures are logged with context (file_id, partition) and do not leave Milvus deleted while RDB still references the file.
♻️ Duplicate comments (2)
openrag/routers/workspaces.py (1)
16-22:⚠️ Potential issue | 🟡 MinorPydantic request models missing
ConfigDict(extra='allow').Per coding guidelines, request models must accept vendor-specific fields like
extra_body. BothCreateWorkspaceRequestandAddFilesRequestwill currently reject unexpected fields.🔧 Proposed fix
-from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict class CreateWorkspaceRequest(BaseModel): + model_config = ConfigDict(extra="allow") workspace_id: str display_name: str | None = None class AddFilesRequest(BaseModel): + model_config = ConfigDict(extra="allow") file_ids: list[str]As per coding guidelines,
Pydantic request models in API endpoints and tests must use ConfigDict(extra='allow') to accept vendor-specific fields like extra_body.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 16 - 22, CreateWorkspaceRequest and AddFilesRequest are strict Pydantic models that will reject vendor-specific fields; update each to allow extra fields by adding a ConfigDict with extra='allow' (e.g., add a nested Config/ModelConfig using ConfigDict(extra='allow') on CreateWorkspaceRequest and AddFilesRequest) so endpoints and tests accept fields like extra_body while preserving existing typed fields.openrag/components/indexer/vectordb/vectordb.py (1)
465-472:⚠️ Potential issue | 🟠 MajorWorkspace search filter missing partition constraint — potential cross-partition data leak.
The workspace file_ids lookup extracts files but doesn't add the workspace's
partition_nameto the Milvus expression. Sincefile_idis only unique per partition, a search withpartition=["all"]could return chunks from another partition that happens to reuse the samefile_id.The workspace's
partition_nameshould be used to constrain the search, or validate that the requested partitions include the workspace's partition.🔒 Proposed fix
if filter: filter = dict(filter) # don't mutate caller's dict if "workspace_id" in filter: workspace_id = filter.pop("workspace_id") + ws = self.partition_file_manager.get_workspace(workspace_id) + if ws: + ws_partition = ws["partition_name"] + # Constrain search to workspace's partition + if partition != ["all"] and ws_partition not in partition: + return [] # Workspace not in requested partitions + expr_parts.append(f'partition == "{ws_partition}"') file_ids = self.partition_file_manager.list_workspace_files(workspace_id) if not file_ids: return [] # Empty workspace → no results id_list = ", ".join(f'"{fid}"' for fid in file_ids) expr_parts.append(f"file_id in [{id_list}]")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` around lines 465 - 472, The workspace filtering logic in the method that builds Milvus expressions currently uses file_ids from partition_file_manager.list_workspace_files(workspace_id) but does not constrain by the workspace's partition, risking cross-partition matches; update the block where filter/workspace_id is handled to fetch the workspace's partition_name (via partition_file_manager or the same API that returns file_ids) and either add a partition constraint to expr_parts (e.g., partition_name == "…") or validate that any requested partitions include the workspace partition before returning/continuing; ensure you still copy filter (filter = dict(filter)), pop workspace_id, build the file_id IN expression into expr_parts, and then add the partition_name constraint or raise/adjust when partitions aren't compatible.
🧹 Nitpick comments (5)
openrag/components/pipeline.py (1)
151-152: Consider usingcall_ray_actor_with_timeout()for consistency.The direct
.remote()call bypasses the centralized timeout and cancellation handling. While this may be acceptable for a quick lookup, it deviates from the established pattern.Based on learnings,
Use the centralized call_ray_actor_with_timeout() utility from components.ray_utils for all Ray actor method calls to ensure proper timeout and cancellation handling.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/pipeline.py` around lines 151 - 152, The direct Ray actor call to vectordb.get_workspace (vectordb.get_workspace.remote(workspace)) bypasses centralized timeout/cancellation handling; replace this direct .remote() invocation with the call_ray_actor_with_timeout utility from components.ray_utils (call_ray_actor_with_timeout(vectordb, "get_workspace", workspace)) so the Vectordb actor method uses the standard timeout and cancellation behavior and follows existing patterns.openrag/routers/workspaces.py (1)
25-30: Direct.remote()calls bypass centralized timeout handling.All workspace endpoints call Ray actor methods directly (e.g.,
vectordb.get_workspace.remote()) instead of usingcall_ray_actor_with_timeout(). This bypasses the shared timeout and cancellation handling the rest of the service uses.💡 Example fix for one call site
+from components.ray_utils import call_ray_actor_with_timeout ... async def require_workspace_in_partition(partition: str, workspace_id: str, vectordb=Depends(get_vectordb)) -> dict: """Validate that a workspace exists and belongs to the given partition.""" - ws = await vectordb.get_workspace.remote(workspace_id) + ws = await call_ray_actor_with_timeout(vectordb.get_workspace.remote(workspace_id)) if not ws or ws["partition_name"] != partition: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Workspace not found") return wsApply the same pattern to all
.remote()calls in this file.Based on learnings,
Use the centralized call_ray_actor_with_timeout() utility from components.ray_utils for all Ray actor method calls to ensure proper timeout and cancellation handling.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 25 - 30, Replace the direct Ray actor call in require_workspace_in_partition so it uses the centralized timeout wrapper: import call_ray_actor_with_timeout from components.ray_utils, call await call_ray_actor_with_timeout(vectordb.get_workspace, workspace_id) instead of vectordb.get_workspace.remote(workspace_id), then perform the same existence and partition check on the returned ws and raise HTTPException if invalid; apply the same pattern to any other .remote() usages in this module to ensure shared timeout/cancellation behavior.openrag/routers/indexer.py (1)
284-285:put_filedoesn't supportworkspace_ids— consider feature parity.The
add_fileendpoint now acceptsworkspace_idsto associate the file with workspaces during indexing, butput_file(which replaces a file) doesn't offer this option. When a file is replaced:
- The old file is deleted (which removes workspace associations via cascade)
- The new file is indexed without workspace associations
Users would need to manually re-add the file to workspaces after a PUT. Consider adding
workspace_idsparameter for consistency, or document this as intentional behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/indexer.py` around lines 284 - 285, The put_file handler currently calls indexer.add_file.remote(path=file_path, metadata=metadata, partition=partition, user=user) but lacks the workspace_ids parameter, causing replaced files to lose workspace associations; update the put_file endpoint (function put_file) to accept an optional workspace_ids argument (e.g., from request body/params), validate/normalize it as done in the add_file flow, and pass workspace_ids through to indexer.add_file.remote (include workspace_ids=workspace_ids in the call) so replaced files are indexed with the same workspace associations.openrag/components/indexer/vectordb/utils.py (2)
165-169: Consider addingindex=Truetopartition_namefor query performance.The migration at
e7f8a9b0c1d2_add_workspaces.pydoesn't explicitly indexpartition_namebeyond the foreign key, butlist_workspacesqueries bypartition_name. While SQLAlchemy FK constraints don't auto-create indexes in PostgreSQL, adding an explicit index would improve lookup performance when listing workspaces by partition.💡 Optional improvement
partition_name = Column( String, ForeignKey("partitions.partition", ondelete="CASCADE"), nullable=False, + index=True, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 165 - 169, Add an explicit index on the partition_name column to speed up queries: update the Column definition named partition_name in openrag/components/indexer/vectordb/utils.py to include index=True, and add an Alembic migration that creates an index on that column for existing databases (use a clear name like ix_<table>_partition_name and include a downgrade to drop it); this ensures list_workspaces queries by partition_name benefit from the index without relying on FK behavior.
768-774: Missing error handling for duplicateworkspace_idconstraint violation.If
create_workspaceis called with a duplicateworkspace_id, SQLAlchemy will raise anIntegrityErrorthat propagates as an unhandled exception. The router handles this with a pre-check, but adding defensive handling here would improve robustness.🛡️ Proposed defensive handling
+from sqlalchemy.exc import IntegrityError +from utils.exceptions.vectordb import VDBInsertError def create_workspace(self, workspace_id: str, partition: str, user_id: int | None, display_name: str | None = None): with self.Session() as session: - ws = Workspace( - workspace_id=workspace_id, partition_name=partition, created_by=user_id, display_name=display_name - ) - session.add(ws) - session.commit() + try: + ws = Workspace( + workspace_id=workspace_id, partition_name=partition, created_by=user_id, display_name=display_name + ) + session.add(ws) + session.commit() + except IntegrityError: + session.rollback() + raise VDBInsertError( + f"Workspace '{workspace_id}' already exists", + status_code=409, + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/utils.py` around lines 768 - 774, create_workspace currently commits a new Workspace without handling IntegrityError from duplicate workspace_id; wrap the commit in a try/except that catches sqlalchemy.exc.IntegrityError, call session.rollback() on error, and then either raise a clear domain/HTTP error (or re-raise a wrapped exception) so callers get a controlled failure instead of an unhandled DB exception; reference the create_workspace method, the Workspace model and self.Session for where to apply this change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/indexer/indexer.py`:
- Around line 120-129: Task is being marked COMPLETED via
task_state_manager.set_state.remote(task_id, "COMPLETED") before calling
vectordb.add_files_to_workspace.remote which can fail and leave workspace
association inconsistent; change flow so workspace association is attempted
before marking COMPLETED or wrap the association in a try/except: call
asyncio.gather on [vectordb.add_files_to_workspace.remote(ws_id, [file_id]) for
ws_id in workspace_ids] first, on success then call
task_state_manager.set_state.remote(task_id, "COMPLETED"); if you choose to keep
marking first, catch Exception around the gather, log the error (include
task_id/file_id/workspace_ids) and set the task state to a failure state (e.g.,
"FAILED" or a retry state) or enqueue a retry so failures aren’t silently
ignored.
---
Outside diff comments:
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 598-607: The Milvus deletion in delete_file currently commits
before the relational DB cleanup, risking orphaned workspace_files if
remove_file_from_all_workspaces or remove_file_from_partition fails; change the
flow in delete_file so you first perform the RDB cleanup via
partition_file_manager.remove_file_from_all_workspaces(file_id, partition) and
partition_file_manager.remove_file_from_partition(file_id=file_id,
partition=partition) inside a try/except that ensures both succeed (log and
raise on failure), and only after successful RDB operations call
self._async_client.delete(collection_name=self.collection_name, filter=...) (or
alternatively, if you must delete Milvus first, wrap the RDB calls in try/except
and on failure attempt to restore the Milvus entry or surface a clear error);
ensure failures are logged with context (file_id, partition) and do not leave
Milvus deleted while RDB still references the file.
---
Duplicate comments:
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 465-472: The workspace filtering logic in the method that builds
Milvus expressions currently uses file_ids from
partition_file_manager.list_workspace_files(workspace_id) but does not constrain
by the workspace's partition, risking cross-partition matches; update the block
where filter/workspace_id is handled to fetch the workspace's partition_name
(via partition_file_manager or the same API that returns file_ids) and either
add a partition constraint to expr_parts (e.g., partition_name == "…") or
validate that any requested partitions include the workspace partition before
returning/continuing; ensure you still copy filter (filter = dict(filter)), pop
workspace_id, build the file_id IN expression into expr_parts, and then add the
partition_name constraint or raise/adjust when partitions aren't compatible.
In `@openrag/routers/workspaces.py`:
- Around line 16-22: CreateWorkspaceRequest and AddFilesRequest are strict
Pydantic models that will reject vendor-specific fields; update each to allow
extra fields by adding a ConfigDict with extra='allow' (e.g., add a nested
Config/ModelConfig using ConfigDict(extra='allow') on CreateWorkspaceRequest and
AddFilesRequest) so endpoints and tests accept fields like extra_body while
preserving existing typed fields.
---
Nitpick comments:
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 165-169: Add an explicit index on the partition_name column to
speed up queries: update the Column definition named partition_name in
openrag/components/indexer/vectordb/utils.py to include index=True, and add an
Alembic migration that creates an index on that column for existing databases
(use a clear name like ix_<table>_partition_name and include a downgrade to drop
it); this ensures list_workspaces queries by partition_name benefit from the
index without relying on FK behavior.
- Around line 768-774: create_workspace currently commits a new Workspace
without handling IntegrityError from duplicate workspace_id; wrap the commit in
a try/except that catches sqlalchemy.exc.IntegrityError, call session.rollback()
on error, and then either raise a clear domain/HTTP error (or re-raise a wrapped
exception) so callers get a controlled failure instead of an unhandled DB
exception; reference the create_workspace method, the Workspace model and
self.Session for where to apply this change.
In `@openrag/components/pipeline.py`:
- Around line 151-152: The direct Ray actor call to vectordb.get_workspace
(vectordb.get_workspace.remote(workspace)) bypasses centralized
timeout/cancellation handling; replace this direct .remote() invocation with the
call_ray_actor_with_timeout utility from components.ray_utils
(call_ray_actor_with_timeout(vectordb, "get_workspace", workspace)) so the
Vectordb actor method uses the standard timeout and cancellation behavior and
follows existing patterns.
In `@openrag/routers/indexer.py`:
- Around line 284-285: The put_file handler currently calls
indexer.add_file.remote(path=file_path, metadata=metadata, partition=partition,
user=user) but lacks the workspace_ids parameter, causing replaced files to lose
workspace associations; update the put_file endpoint (function put_file) to
accept an optional workspace_ids argument (e.g., from request body/params),
validate/normalize it as done in the add_file flow, and pass workspace_ids
through to indexer.add_file.remote (include workspace_ids=workspace_ids in the
call) so replaced files are indexed with the same workspace associations.
In `@openrag/routers/workspaces.py`:
- Around line 25-30: Replace the direct Ray actor call in
require_workspace_in_partition so it uses the centralized timeout wrapper:
import call_ray_actor_with_timeout from components.ray_utils, call await
call_ray_actor_with_timeout(vectordb.get_workspace, workspace_id) instead of
vectordb.get_workspace.remote(workspace_id), then perform the same existence and
partition check on the returned ws and raise HTTPException if invalid; apply the
same pattern to any other .remote() usages in this module to ensure shared
timeout/cancellation behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: e5fe4c4d-8551-47ab-85e6-b67d0210354b
📒 Files selected for processing (7)
openrag/components/indexer/indexer.pyopenrag/components/indexer/vectordb/utils.pyopenrag/components/indexer/vectordb/vectordb.pyopenrag/components/pipeline.pyopenrag/routers/indexer.pyopenrag/routers/search.pyopenrag/routers/workspaces.py
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
openrag/components/indexer/vectordb/vectordb.py (1)
466-481:⚠️ Potential issue | 🟠 MajorDon't widen
async_search()beyond the caller's partition list.If a caller passes
partition=["foo"]and a workspace frombar, Lines 480-481 remove the original scope and searchbaranyway.async_search()should only pin to the workspace partition after confirming that partition is allowed by the caller; otherwise it should return[].Suggested fix
ws_partition = ws["partition_name"] + if partition != ["all"] and ws_partition not in partition: + return [] # Replace any outer partition filter with the workspace's partition expr_parts = [p for p in expr_parts if not p.startswith("partition in ")] expr_parts.append(f'partition == "{ws_partition}"')🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/indexer/vectordb/vectordb.py` around lines 466 - 481, The async_search() workspace pin currently unconditionally replaces caller partition filters with the workspace's partition (ws_partition), widening scope; change the logic in the workspace_id block so that after computing ws_partition you first check whether the caller explicitly constrained partitions (e.g. a local variable like partition or any existing "partition in ..." expr_parts): if the caller provided allowed partitions and ws_partition is not among them, return [] immediately; otherwise (no caller partition constraint or ws_partition is allowed) proceed to remove any outer partition filters and append the pinned f'partition == "{ws_partition}"' to expr_parts as before. Ensure you reference the workspace_id handling, ws = self.partition_file_manager.get_workspace(...), ws_partition, expr_parts and the early returns so the change preserves current behavior when the workspace partition is permitted.
🧹 Nitpick comments (2)
openrag/components/pipeline.py (1)
10-10: Use the repo'sopenrag.*import path for the new Ray helper.This new import keeps the legacy
components.*style alive in another touched line. Please root the added utility import underopenragwhen updating this section.As per coding guidelines,
**/*.py: Use absolute imports from theopenrag/directory (which is the Python path root).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/pipeline.py` at line 10, Replace the relative-style import for the Ray helper with the repo-root absolute import: change the line importing call_ray_actor_with_timeout (currently "from components.ray_utils import call_ray_actor_with_timeout") to use the openrag package root, e.g. "from openrag.components.ray_utils import call_ray_actor_with_timeout", so the symbol call_ray_actor_with_timeout is resolved via the openrag.* import path.openrag/routers/workspaces.py (1)
5-12: Useopenrag.*imports throughout this new router module.The new file introduces
components.*,utils.*,config, and relative.utilsimports, which adds another exception to the repo's import rule. Please root these underopenragbefore this module becomes part of the public API surface.As per coding guidelines,
**/*.py: Use absolute imports from theopenrag/directory (which is the Python path root).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/routers/workspaces.py` around lines 5 - 12, This file uses non-rooted imports; change them to absolute imports under the openrag package: replace components.ray_utils -> openrag.components.ray_utils (for call_ray_actor_with_timeout), config -> openrag.config (for load_config), utils.dependencies -> openrag.utils.dependencies (for get_vectordb), utils.logger -> openrag.utils.logger (for get_logger), and the relative .utils -> openrag.routers.utils (for require_partition_editor, require_partition_owner, require_partition_viewer); keep the imported names/APIRouter/Depends/HTTPException/status the same.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/pipeline.py`:
- Around line 152-166: The workspace partition check drops workspace scoping
when partition == ["all"]; update the conditional in the block that calls
vectordb.get_workspace (used via ray.get_actor("Vectordb") and
call_ray_actor_with_timeout) so it only treats a workspace as missing if ws is
falsy OR ( "all" not in partition AND ws["partition_name"] not in partition ).
Keep the rest of the logic intact and ensure filter_dict still becomes
{"workspace_id": workspace} when workspace remains valid.
In `@openrag/routers/workspaces.py`:
- Around line 112-127: The response currently reports orphaned_files_deleted as
len(orphaned) even though some deletions may have failed; modify the
post-deletion logic in the block using call_ray_actor_with_timeout and
vectordb.delete_file.remote to compute successful deletions by iterating over
results (the results from asyncio.gather bound to variable results) and counting
entries that are not Exception, optionally collecting failed file_ids (use
logger.warning with file_id and error via str(result) as already done), then
return {"status":"deleted","orphaned_files_deleted": <count_of_successes>,
"orphaned_files_failed": <count_of_failures>} (or at minimum set
orphaned_files_deleted to the count of non-exception results) so the API
accurately reflects what actually succeeded; references: orphaned, results,
call_ray_actor_with_timeout, vectordb.delete_file.remote, logger.warning.
---
Duplicate comments:
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 466-481: The async_search() workspace pin currently
unconditionally replaces caller partition filters with the workspace's partition
(ws_partition), widening scope; change the logic in the workspace_id block so
that after computing ws_partition you first check whether the caller explicitly
constrained partitions (e.g. a local variable like partition or any existing
"partition in ..." expr_parts): if the caller provided allowed partitions and
ws_partition is not among them, return [] immediately; otherwise (no caller
partition constraint or ws_partition is allowed) proceed to remove any outer
partition filters and append the pinned f'partition == "{ws_partition}"' to
expr_parts as before. Ensure you reference the workspace_id handling, ws =
self.partition_file_manager.get_workspace(...), ws_partition, expr_parts and the
early returns so the change preserves current behavior when the workspace
partition is permitted.
---
Nitpick comments:
In `@openrag/components/pipeline.py`:
- Line 10: Replace the relative-style import for the Ray helper with the
repo-root absolute import: change the line importing call_ray_actor_with_timeout
(currently "from components.ray_utils import call_ray_actor_with_timeout") to
use the openrag package root, e.g. "from openrag.components.ray_utils import
call_ray_actor_with_timeout", so the symbol call_ray_actor_with_timeout is
resolved via the openrag.* import path.
In `@openrag/routers/workspaces.py`:
- Around line 5-12: This file uses non-rooted imports; change them to absolute
imports under the openrag package: replace components.ray_utils ->
openrag.components.ray_utils (for call_ray_actor_with_timeout), config ->
openrag.config (for load_config), utils.dependencies ->
openrag.utils.dependencies (for get_vectordb), utils.logger ->
openrag.utils.logger (for get_logger), and the relative .utils ->
openrag.routers.utils (for require_partition_editor, require_partition_owner,
require_partition_viewer); keep the imported
names/APIRouter/Depends/HTTPException/status the same.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ba2f8f25-63a4-482e-a7fa-305ca02e90ac
📒 Files selected for processing (5)
.hydra_config/config.yamlopenrag/components/indexer/vectordb/vectordb.pyopenrag/components/pipeline.pyopenrag/routers/search.pyopenrag/routers/workspaces.py
🚧 Files skipped from review as they are similar to previous changes (1)
- openrag/routers/search.py
EnjoyBacon7
force-pushed
the
feat/workspace
branch
from
March 12, 2026 10:40
a78023c to
1a0dc1a
Compare
We improve the disk file saving by doing the following: - Do not load the whole file in memory, but rather stream the http buffer and write in chunks - Do not make blocking I/O , to allow parallel writes - Add random prefix to saved files, to avoid name collisions
Add a new `v1/tools` endpoint , to allow custom tools execution. This is useful to execute specific openRAG features, such as text extraction, without the semantic search functionnalities
BREAKING CHANGE: adds `workspaces` and `workspace_files` tables to PostgreSQL. Apply the Alembic migration before restarting the service: https://linagora.github.io/openrag/documentation/sql_migration/#step-2-apply-the-migration
asyncio.gather without return_exceptions=True would raise on the first failed delete_file, leaving the workspace deleted from PostgreSQL but orphaned chunks remaining in Milvus. Use return_exceptions=True and log individual failures so the endpoint always returns successfully after the DB commit.
…aint in implementation
…y with Ray actor pattern
…filtering search results
…artition to prevent cross-partition collisions
Move workspace_files insertion into Indexer.add_file() so it only runs after the file record is committed (COMPLETED state). Previously the router associated workspaces immediately after queuing, leaving orphaned workspace_files rows if indexing failed.
…actor_with_timeout Replace bare await actor.method.remote() calls in workspaces.py, search.py, and pipeline.py with call_ray_actor_with_timeout() so workspace endpoints participate in the shared timeout and cancellation handling used by the rest of the service. Add vectordb_timeout config key (default 30s, env VECTORDB_TIMEOUT) under ray.indexer.
Add ConfigDict(extra="allow") to CreateWorkspaceRequest and AddFilesRequest so the API tolerates vendor-specific fields like extra_body, consistent with other request models in the service.
…c_search file_id is only unique per (file_id, partition_name). The previous code built a 'file_id in [...]' expression without also scoping to the workspace's partition, so a partition=["all"] search could return chunks from a different partition that reuses the same file_id. Fix: fetch the workspace metadata to get its partition_name, strip any outer partition filter, and replace it with an exact equality expression for the workspace's own partition before appending the file_id filter.
Wrap the asyncio.gather call that links the file to workspaces in a try/except so that a transient DB error after set_state(COMPLETED) logs a warning instead of propagating to the outer except handler, which would incorrectly overwrite COMPLETED→FAILED via set_failed_if_not_cancelled.
The guard 'ws["partition_name"] not in partition' always evaluated True when partition==["all"] because "all" is never a real partition name, silently dropping the workspace filter and falling back to an unscoped search. Add an '"all" not in partition' short-circuit so the filter is only discarded when the workspace genuinely doesn't belong to the explicitly named partition set.
orphaned_files_deleted previously returned len(orphaned) even when some delete_file calls had failed. Now only successful deletions are counted and a new orphaned_files_failed list is included in the response so callers can distinguish partial failures.
EnjoyBacon7
force-pushed
the
feat/workspace
branch
from
March 12, 2026 10:51
1a0dc1a to
66bbf7a
Compare
This was referenced Mar 12, 2026
This was referenced Apr 19, 2026
Merged
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
3 participants
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.
Workspace is a new concept that can be seen as sub-partition spaces for files.
Inside a partition, we can now organize files into workplaces and search for files for a given workspace.
The main reason to use workspace rather than dedicated partitions is to avoid file duplication: a single file can belong to multiple workplace, while a same file can exist in multiple partitions.
This is typically useful to clients with 1 partition per user, each user having the possibility to have workplaces with specific files into it, to create specialized context assistants.
Summary by CodeRabbit
New Features
Documentation
Tests
Chores