Feat/relationship3 - #228
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 document relationship support: nullable relationship_id and parent_id persisted during ingestion, DB migration and indexes, APIs to fetch related files and ancestors, retriever/search expansion for related/ancestor chunks, router endpoints, docs, and unit/integration tests. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant SearchRouter
participant Retriever
participant VectorDB
participant Database
Client->>SearchRouter: search(text, include_related=true, include_ancestors=true)
SearchRouter->>Retriever: retrieve(text, top_k)
Retriever->>VectorDB: similarity_search(text)
VectorDB->>Database: lookup file metadata (file_ids)
Database-->>VectorDB: base file metadata
VectorDB-->>Retriever: base Documents
SearchRouter->>Retriever: expand_search_results(results, include_related, include_ancestors)
alt include_related
Retriever->>VectorDB: get_related_chunks(partition, relationship_id, limit)
VectorDB->>Database: get_files_by_relationship(partition, relationship_id)
Database-->>VectorDB: related file_ids
VectorDB-->>Retriever: related Documents
end
alt include_ancestors
Retriever->>VectorDB: get_ancestor_chunks(partition, file_id, max_ancestor_depth)
VectorDB->>Database: get_file_ancestors(partition, file_id)
Database-->>VectorDB: ordered ancestor file_ids
VectorDB-->>Retriever: ancestor Documents
end
Retriever->>Retriever: deduplicate documents
Retriever-->>Client: expanded results
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 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)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/linked_files.md`:
- Around line 32-37: The markdown table under the `File` model has misaligned
pipes causing markdownlint MD060 failures; fix it by aligning the column
separators so each row has the same number of pipes and columns and the dashes
in the separator row match column widths for the two fields `relationship_id`
and `parent_id` (and apply the same alignment to the other occurrence around
lines 217-223); update the table rows so the header, separator, and content rows
are consistently formatted to satisfy the linter.
- Around line 45-58: Several fenced code blocks in the linked_files.md
documentation (for example the block starting with "Email A (root)" and the
block showing "Documents/2024/Q1/") lack language identifiers; update each
fenced triple-backtick to include an appropriate language tag (e.g., text for
plain diagrams/listings, http for HTTP snippets, json for JSON blobs) so
markdownlint MD040 is satisfied and readability improves—apply this change to
every fenced block in the document, especially those that contain directory
listings, thread diagrams, or HTTP/JSON examples.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 581-629: The recursive CTE in get_file_ancestors returns JSON
columns as raw strings, so unpacking row.file_metadata with **(row.file_metadata
or {}) will raise a TypeError; update the result mapping in get_file_ancestors
to deserialize JSON (e.g., use json.loads(row.file_metadata) if
row.file_metadata else {}) before merging into the dict so file_metadata becomes
a dict; import json at top if needed and keep references to get_file_ancestors,
session.execute, and row.file_metadata when applying the change.
- Around line 68-74: The to_dict() method currently merges user-supplied
metadata after the canonical fields so file_metadata can override canonical
keys; change the merge order in to_dict() to unpack metadata first and then set
the canonical keys (use **metadata before "partition": self.partition_name,
"file_id": self.file_id, "relationship_id": self.relationship_id, "parent_id":
self.parent_id) so that partition_name, file_id, relationship_id, and parent_id
in the returned dict always take precedence and cannot be shadowed by user
metadata.
In `@openrag/components/indexer/vectordb/vectordb.py`:
- Around line 1036-1045: The current filter_expr construction concatenates raw
file_ids which can break on quotes/backslashes; update the file_id_list building
in vectordb.py to json-encode each file id (use json.dumps(fid) for each entry
of file_ids) so the resulting filter_expr (the variable filter_expr) is
correctly escaped for the Milvus query call (self._async_client.query with
collection_name and partition); also add/import json at the top of the module if
not present.
- Around line 1066-1088: The returned Documents must be ordered by ancestry:
after calling partition_file_manager.get_ancestor_file_ids(...) and fetching
results via self._async_client.query(...), build a lookup from result["file_id"]
to the result row and then iterate ancestor_file_ids in order to produce
Documents so the output list follows the ancestor sequence; use the same
Document construction (page_content=res["text"], metadata={...}) for each
ancestor_file_id that exists in the lookup and skip or handle missing IDs
accordingly to preserve root-to-target ordering.
In `@openrag/routers/search.py`:
- Around line 138-180: Add a guard that validates include_related and
include_ancestors are mutually exclusive at the start of the request handling
for search_multiple_partitions (and the other search endpoint for a single
partition), returning an HTTP 400 error if both are true; specifically check the
boolean flags include_related and include_ancestors before calling
indexer.asearch.remote or _expand_with_related_chunks and raise/return a 400
response with a clear message when both flags are set, so the contract in the
docs is enforced.
In `@openrag/tests/test_relationships_integration.py`:
- Around line 32-41: The folder fixture documents set parent_id to the folder's
own ID which conflicts with the documented semantics that parent_id is not used
for folders; update the Document construction for folder-related test fixtures
(the Document instances created with metadata keys "relationship_id" and
"parent_id") to omit the "parent_id" key (or set it to None/absent) so only
"relationship_id" (folder_id) is used to mark folder membership, and leave any
ancestor logic unaffected.
🧹 Nitpick comments (8)
tests/api_tests/test_search.py (1)
169-173: Fragile assertion usingpop()on a set.Using
pop()on a set returns an arbitrary element since sets are unordered. Since the previous assertion (lines 165-167) already ensures noNonevalues exist, and the test expects all documents to share the samerelationship_id, this check could be clearer.💡 Suggested improvement
# verify that the relationship_id matches the expected one expected_relationship_id = folder_files["file1.txt"][1] # relationship_id used during indexing - assert relationship_ids.pop() == expected_relationship_id, ( - f"Documents should have relationship_id {expected_relationship_id}" - ) + assert relationship_ids == {expected_relationship_id}, ( + f"All documents should have relationship_id {expected_relationship_id}, got {relationship_ids}" + )docs/content/docs/documentation/linked_files.md (1)
276-276: Consider more specific wording than “very”.
Minor style nit: “large” or a concrete threshold reads crisper.openrag/components/retriever.py (1)
194-204: Deduplicate file lookups to avoid repeated ancestor fetches.
If multiple chunks belong to the same file, you’ll issue redundant ancestor calls. A set of tuples avoids that.♻️ Suggested refactor
- file_infos = [] # List of (partition, file_id) tuples + file_infos: set[tuple[str, str]] = set() ... - if self.include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) + if self.include_ancestors: + file_infos.add((metadata.get("partition"), metadata.get("file_id")))Also applies to: 227-242
openrag/routers/search.py (2)
44-54: Deduplicate ancestor fetches by file_id.
Multiple chunks from the same file will cause repeated ancestor requests; using a set avoids redundant calls.♻️ Suggested refactor
- file_infos = [] # List of (partition, file_id) tuples + file_infos: set[tuple[str, str]] = set() ... - if include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) + if include_ancestors: + file_infos.add((metadata.get("partition"), metadata.get("file_id")))Also applies to: 77-92
291-324: Avoid expansion call/logging when flags are always false.
search_filecurrently calls the helper and logs “Expanded results” even though no expansion can occur. Consider removing or guarding it.🧹 Suggested cleanup
- # Expand with related/ancestor chunks if requested - results = await _expand_with_related_chunks( - results=results, - vectordb=vectordb, - include_related=False, - include_ancestors=False, - ) - - log.info( - "Expanded results with related/ancestor chunks.", - results=len(results), - ) + # No relationship expansion for file-scoped searchopenrag/components/indexer/vectordb/utils.py (2)
45-51: Consider dropping redundant single-column indexes if composite indexes are sufficient.
relationship_id/parent_idare indexed both individually and as part of composite indexes withpartition_name. If you don’t have standalone lookups on those fields, the extra indexes add write overhead for little gain.♻️ Suggested adjustment (if standalone lookups aren’t needed)
- relationship_id = Column( - String, nullable=True, index=True - ) # Groups related documents (e.g., email thread ID, folder path) + relationship_id = Column( + String, nullable=True + ) # Groups related documents (e.g., email thread ID, folder path) @@ - parent_id = Column( - String, nullable=True, index=True - ) # Hierarchical parent reference (e.g., parent email, parent folder) + parent_id = Column( + String, nullable=True + ) # Hierarchical parent reference (e.g., parent email, parent folder)Also applies to: 61-63
213-225: Optional: validate parent/relationship consistency on insert.If
parent_idpoints to a file in another partition or an unrelated relationship, ancestor traversal can silently truncate or cross groups. If that’s not intended, consider validating the parent’s existence (and relationship_id match) or adding a constraint.Also applies to: 248-251
openrag/components/test_relationships.py (1)
1-49: Keep the test helper’sto_dictshape aligned with production.
FileModel.to_dict()nestsfile_metadata, while productionFile.to_dict()(openrag/components/indexer/vectordb/utils.py) flattens metadata into top-level keys. If these tests are intended to mirror production behavior, consider matching the output shape (or add a focused assertion to cover the production format) to avoid drift.
0376c43 to
dcb1a86
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@openrag/components/retriever.py`:
- Around line 191-192: The set comprehension that builds seen_ids uses
doc.metadata.get("_id") and will include None for docs missing an _id, causing
all such docs to be treated as duplicates; update the logic around seen_ids (the
variable created from chunks and the comprehension using
doc.metadata.get("_id")) to exclude None values (only add ids that are not None)
and ensure downstream deduplication logic treats docs without an _id as distinct
(e.g., by not inserting None into seen_ids or by using a fallback unique key for
those docs).
🧹 Nitpick comments (4)
docs/content/docs/documentation/linked_files.md (1)
65-73: Add language identifiers to fenced code blocks.Several code blocks are missing language identifiers (markdownlint MD040). Use
textfor plain text diagrams/listings:📝 Suggested fix
-``` +```text Documents/2024/Q1/-``` +```text Search: "budget report" → Returns 5 chunks from 5 different files-``` +```text Query → Hybrid Search → Reranking → Data Expansion → LLMAlso applies to: 228-246, 282-284
openrag/routers/partition.py (1)
413-426: Unusedrequestparameter inget_related_files.The
requestparameter is injected but never used in this endpoint. Consider removing it for cleaner code, or document if it's intended for future use.♻️ Suggested fix
async def get_related_files( - request: Request, partition: str, relationship_id: str, vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ):tests/api_tests/conftest.py (1)
151-156: Module-levelqueryvariable should be a constant or fixture.The
querystring is defined at module level rather than as a fixture or constant. Consider either:
- Renaming to
QUERY(constant convention)- Moving into a fixture for consistency with
exact_match_queryopenrag/routers/search.py (1)
296-329: Remove dead code in search_file endpoint.The expansion call with hardcoded
Falsevalues is effectively a no-op (the function returns early when both flags are false). This creates misleading logs ("Expanded results") and unnecessary code paths.Proposed cleanup
results = await indexer.asearch.remote(query=text, top_k=top_k, partition=partition, filter={"file_id": file_id}) log.info("Semantic search on specific file completed.", result_count=len(results)) - # Expand with related/ancestor chunks if requested - results = await _expand_with_related_chunks( - results=results, - vectordb=vectordb, - include_related=False, - include_ancestors=False, - ) - - log.info( - "Expanded results with related/ancestor chunks.", - results=len(results), - ) - documents = [Also remove the unused
vectordbdependency from this endpoint if expansion is not supported.
paultranvan
left a comment
paultranvan
left a comment
There was a problem hiding this comment.
I think this is a good foundation, but I have some concerns about API and database
| @@ -0,0 +1,306 @@ | |||
| --- | |||
| title: 🔗 Document Relationships & Linked Files | |||
There was a problem hiding this comment.
In this documentation, we use the folder and email as use-cases, which is good, but I think we should use them as example and not as API feature like it is currently presented. Indeed, we could have different datatype other than folder/email where this feature makes sense: conversations chat, agenda, etc. So in this sense, relationship_id is absolutely not restricted to folder ; just like ancestors that could be used in conversation threads, for instance.
There was a problem hiding this comment.
See the answer in the next comment
| AND f.partition_name = a.partition_name | ||
| ) | ||
| SELECT * FROM ancestors ORDER BY depth DESC | ||
| """) |
There was a problem hiding this comment.
That seems pretty dangerous: if a circular parent_id reference happens, we will end up with an infinite loop. We probably need to enforce a depth limit
There was a problem hiding this comment.
Also, this database command is quite complex. I think we need to have some perf measurements, to evaluate the impact on such query on the retrieval.
And maybe eventually we could envision to have some metrics on the endpoints, to be warned if it appears we kill performances by > X % @andyne13
There was a problem hiding this comment.
a depth limit parameter has been added with tests that go with it.
There was a problem hiding this comment.
Thanks for the depth limit, that helps :)
I'm still worried about query performance, though @andyne13
| @@ -0,0 +1,53 @@ | |||
| { | |||
There was a problem hiding this comment.
Nice to have some dataset for tests 👍
It's probably worth adding a "fixtures" folder, as we will certainly add more dummy data in the future
751d362 to
770dd23
Compare
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/API.mdx`:
- Line 129: The markdown contains a broken link "[context aware retrieval]()"
with an empty href; update the API.mdx sentence by either supplying the correct
target URL for "context aware retrieval" (replace the empty parentheses with the
proper path/anchor) or remove the link markup and leave plain text "context
aware retrieval" if there is no destination; target the exact link text
"[context aware retrieval]()" in the line "For context-aware search, see [search
endpoints](`#-semantic-search`) and [context aware retrieval]()."
In `@docs/content/docs/documentation/linked_files.md`:
- Line 119: The sentence about ancestors is incorrect/confusing: clarify that
parent_id points to the immediate parent and that ancestor traversal follows
parent links (so fileA.2's ancestors are the chain via parent_id, e.g., fileA.2
→ folderA (and then folderA's parent, etc.), not sibling files like fileA or
fileB). Update the text around "fileA.2 => fileA, fileB but not fileA.1" to a
correct example using the parent_id chain (mentioning parent_id and showing the
actual traversal like fileA.2 → folderA → ...), and explicitly state that
siblings are not considered ancestors.
- Around line 81-92: The markdown examples use fenced code blocks without
language tags (the blocks showing directory trees like the one starting with
"Documents/2024/Q1/"); add a language identifier such as ```text (or ```txt) to
each of those fenced code blocks to satisfy markdownlint MD040—update the block
that contains "Documents/2024/Q1/" and make the same change to the other
directory-structure examples in the file.
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 596-618: The recursive CTE "ancestors" (used in the Session block)
can run unbounded when max_ancestor_depth is None; change the query to guard
against cycles/runaway recursion by (1) adding a visited-path check in the CTE
(e.g., include a path/array of visited file_ids and add a WHERE clause like "AND
f.file_id NOT IN (path)" when recursing) and (2) enforce a default hard cap when
max_ancestor_depth is None (introduce a constant like MAX_SAFE_ANCESTOR_DEPTH
and apply it via depth_condition). Update the recursion logic that builds
"ancestors" to carry the path and check membership (instead of relying on UNION
ALL alone), and ensure any use of depth_condition/max_ancestor_depth ties into
that default cap so the query cannot loop indefinitely.
In `@openrag/components/retriever.py`:
- Around line 167-178: HyDeRetriever.__init__ currently invokes super().__init__
twice which causes the second call to overwrite relationship-related parameters;
remove the duplicate minimal call (the second super().__init__(top_k,
similarity_threshold, **kwargs)) so only the full super().__init__(top_k,
similarity_threshold, with_surrounding_chunks, include_related,
include_ancestors, related_limit, max_ancestor_depth, **kwargs) remains,
ensuring include_related/include_ancestors/related_limit/max_ancestor_depth are
preserved.
In `@openrag/components/test_relationships.py`:
- Around line 104-140: In get_file_ancestors, ensure the SQL parameter
:max_ancestor_depth is bound when depth_condition contains it: keep building
depth_condition as you do, and change the session.execute call so it
conditionally adds "max_ancestor_depth": max_ancestor_depth to the params dict
when max_ancestor_depth is not None; reference the depth_condition variable and
the session.execute(query, params) call so the params include file_id, partition
and optionally max_ancestor_depth.
In `@openrag/routers/search.py`:
- Around line 92-99: The call sites in openrag/routers/search.py are passing
vectordb= to _expand_with_related_chunks but the helper (defined in
components/retriever.py) expects the parameter name db; change all calls to pass
db=vectordb (not vectordb=) — specifically update the three invocations of
_expand_with_related_chunks (the blocks around results = await
_expand_with_related_chunks(...), including the occurrences that pass
include_related, include_ancestors, related_limit, and max_ancestor_depth) so
the keyword matches the helper signature (db=) to avoid the TypeError thrown by
_expand_with_related_chunks.
🧹 Nitpick comments (4)
.hydra_config/retriever/base.yaml (2)
5-6: Verify default behavior for existing deployments.Both
include_relatedandinclude_ancestorsdefault totrue. For existing deployments upgrading to this version, this could unexpectedly increase context size and token usage, potentially affecting LLM costs and response quality.Consider defaulting to
falsefor backward compatibility, allowing users to opt-in explicitly.💡 Suggested safer defaults
-include_related: ${oc.decode:${oc.env:INCLUDE_RELATED, true}} -include_ancestors: ${oc.decode:${oc.env:INCLUDE_ANCESTORS, true}} +include_related: ${oc.decode:${oc.env:INCLUDE_RELATED, false}} +include_ancestors: ${oc.decode:${oc.env:INCLUDE_ANCESTORS, false}}
8-8: Consider a more specific environment variable name.
MAX_DEPTHis generic and could conflict with other configurations. For consistency with other retriever env vars likeRETRIEVER_TOP_K, consider usingRETRIEVER_MAX_ANCESTOR_DEPTHorMAX_ANCESTOR_DEPTH.💡 Suggested fix
-max_ancestor_depth: ${oc.decode:${oc.env:MAX_DEPTH, 10}} # Maximum depth for ancestor retrieval (null = unlimited) +max_ancestor_depth: ${oc.decode:${oc.env:MAX_ANCESTOR_DEPTH, 10}} # Maximum depth for ancestor retrieval (null = unlimited)docs/content/docs/documentation/API.mdx (1)
84-84: Minor: Add period after "etc".Per American English style, abbreviations like "etc." require a period.
📝 Suggested fix
-...email threads ,etc (see [Document Relationships documentation]... +...email threads, etc. (see [Document Relationships documentation]...openrag/components/retriever.py (1)
236-256: Consider parallel fetching for related chunks.The current implementation fetches related chunks sequentially in a loop. For multiple relationship IDs, this could be slow. Consider using
asyncio.gatherfor parallel fetching.♻️ Suggested parallel implementation
# Fetch related chunks by relationship_id if include_related: - for partition, rel_id in relationship_ids: - if partition and rel_id: - try: - related_chunks = await db.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), - ) + tasks = [ + db.get_related_chunks.remote(partition=p, relationship_id=r, limit=related_limit) + for p, r in relationship_ids if p and r + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + for result in 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)
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/API.mdx`:
- Line 84: The phrase "folder-based relationships, email threads ,etc" in the
API documentation uses the abbreviation "etc" without a period; update that text
to "etc." (and remove the stray space before the comma if present) so the
abbreviation follows American English style — locate the string "folder-based
relationships, email threads ,etc" in the API.mdx content and replace it with
"folder-based relationships, email threads, etc.".
In `@openrag/components/indexer/vectordb/utils.py`:
- Around line 626-635: The ancestor result dict currently spreads
row.file_metadata after the canonical fields which allows metadata to override
canonical keys; change the merge to filter out those canonical keys from
row.file_metadata (file_id, partition, parent_id, relationship_id, depth) before
spreading so user metadata cannot collide — e.g., compute a safe_metadata from
row.file_metadata by excluding those keys and then include **safe_metadata in
the dict generation inside the list comprehension that builds the ancestor
result.
🧹 Nitpick comments (2)
docs/content/docs/documentation/linked_files.md (1)
144-146: Add language identifier to fenced code block.Add
textidentifier for this pipeline diagram.-``` +```text Query → Hybrid Search → Reranking → Data Expansion → LLMopenrag/routers/search.py (1)
244-255: Remove unnecessary expansion call.The
search_fileendpoint always passesinclude_related=Falseandinclude_ancestors=False, making the expansion call a no-op that adds overhead. Either remove the call or add expansion parameters to this endpoint if the feature is intended.♻️ Suggested simplification
results = await indexer.asearch.remote(query=text, top_k=top_k, partition=partition, filter={"file_id": file_id}) log.info("Semantic search on specific file completed.", result_count=len(results)) - # Expand with related/ancestor chunks if requested - results = await _expand_with_related_chunks( - results=results, - vectordb=vectordb, - include_related=False, - include_ancestors=False, - ) - - log.info( - "Expanded results with related/ancestor chunks.", - results=len(results), - ) - documents = [Also note:
vectordb=should bedb=if keeping this code.
8e2ed47 to
25105c5
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In @.hydra_config/retriever/base.yaml:
- Around line 4-8: Update the YAML defaults so they match the API/code defaults:
change include_related and include_ancestors to false and change related_limit
to 20 in the retriever config; this ensures the values applied by
RetrieverFactory.create_retriever() align with the parameters documented/used in
routers/search.py (referenced symbols: include_related, include_ancestors,
related_limit, RetrieverFactory.create_retriever, routers/search.py).
In `@tests/api_tests/test_search.py`:
- Around line 126-135: The relationship_id check is weak because using
relationship_ids.pop() only verifies one element; replace it with a strict check
that every document's relationship_id equals the expected one by comparing the
entire set (relationship_ids) to a singleton set containing
expected_relationship_id (or assert len(relationship_ids) == 1 and
expected_relationship_id in relationship_ids) to guarantee all docs share the
same relationship_id; update the assertions around relationship_ids and
expected_relationship_id accordingly in the test_search.py block that builds
relationship_ids from data_with["documents"].
🧹 Nitpick comments (2)
openrag/routers/search.py (1)
223-255: Skip expansion call when both flags are false.
search_filealways invokes_expand_with_related_chunkseven though both flags are hardcodedFalse. Consider skipping the call to avoid unnecessary work/log noise.Suggested refactor
- results = await _expand_with_related_chunks( - results=results, - db=vectordb, - include_related=False, - include_ancestors=False, - ) - - log.info( - "Expanded results with related/ancestor chunks.", - results=len(results), - ) + # No expansion for file search since flags are always falseopenrag/components/retriever.py (1)
225-235: Deduplicate ancestor fetch inputs to avoid redundant RPCs.If multiple results share the same (partition, file_id), you’ll issue repeated ancestor queries.
♻️ Suggested change
- file_infos = [] # List of (partition, file_id) tuples + file_infos = set() # Unique (partition, file_id) tuples @@ - if include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) + if include_ancestors and metadata.get("partition") and metadata.get("file_id"): + file_infos.add((metadata.get("partition"), metadata.get("file_id"))) @@ - for partition, file_id in file_infos: + for partition, file_id in file_infos:Also applies to: 258-273
cc52269 to
62ff653
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/API.mdx`:
- Around line 254-327: Update the section heading "Partitions & files
Management" to include the emoji prefix used elsewhere (add "📦" before the
heading text) so it matches the document's established style; ensure the new
heading reads "📦 Partitions & files Management" while leaving the endpoint docs
for GET /{partition}/relationships/{relationship_id} and GET
/{partition}/file/{file_id}/ancestors unchanged.
In `@openrag/routers/search.py`:
- Around line 244-255: Remove the no-op expansion call in search_file: the call
to _expand_with_related_chunks(results=results, db=vectordb,
include_related=False, include_ancestors=False) always early-returns and the
subsequent log.info("Expanded results with related/ancestor chunks.",
results=len(results)) is misleading; either delete both the
_expand_with_related_chunks invocation and that log message, or modify the call
to use the same expansion flags/parameters used by other endpoints (e.g., read
expansion params from the request and pass include_related/include_ancestors
accordingly) so expansion actually occurs when requested.
🧹 Nitpick comments (2)
tests/api_tests/test_search.py (1)
183-225: Consider adding an assertion that indexing succeeded for each email.The fixture silently continues (
time.sleep(3); continue) when no task info is returned (line 211-213), and skips on failure. But if the POST response itself fails (non-2xx status), there's no assertion. Addingassert response.status_code == 201after each POST would catch upload failures early instead of producing confusing assertion errors in downstream tests.Proposed improvement
with open(file_path, "rb") as f: response = api_client.post( f"/indexer/partition/{created_partition}/file/{email_id}", files={"file": (email_info["filename"], f, "text/plain")}, data={"metadata": json.dumps(metadata)}, ) + assert response.status_code == 201, f"Failed to upload {email_id}: {response.text}" data = response.json()openrag/components/indexer/vectordb/utils.py (1)
214-216: Minor: Inconsistent type annotation style.
None | stris used here, while the rest of the codebase consistently usesstr | None(e.g.,int | Noneon the same line 213,dict | Noneon line 212). Consider aligning for consistency.- relationship_id: None | str = None, - parent_id: None | str = None, + relationship_id: str | None = None, + parent_id: str | None = None,
62ff653 to
dff2005
Compare
coderabbitai
Bot
added
the
chore
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@docs/content/docs/documentation/linked_files.md`:
- Line 132: The document skips section 3: locate the heading "## **4. API
Endpoints**" and either insert the missing section 3 content (placeholder or
actual subsection) between "2. Data Model" and this heading, or renumber the
existing headings by changing "## **4. API Endpoints**" to "## **3. API
Endpoints**" (and similarly change any following "5." heading to "4.") so the
section numbers are sequential; update any internal references to these section
numbers accordingly.
🧹 Nitpick comments (9)
openrag/components/retriever.py (1)
229-234: Redundant remote calls for ancestor chunks from the same file.
file_infoscollects(partition, file_id)for every result document. When multiple chunks come from the same file, this triggers duplicateget_ancestor_chunkscalls for the samefile_id. While deduplication viaseen_idsprevents duplicate output, each redundant call is still an awaited remote RPC.Consider deduplicating
file_infosbefore the fetch loop:♻️ Suggested improvement
- file_infos = [] # List of (partition, file_id) tuples for doc in results: metadata = doc.metadata if include_related and metadata.get("relationship_id"): relationship_ids.add((metadata.get("partition"), metadata.get("relationship_id"))) if include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) + partition = metadata.get("partition") + file_id = metadata.get("file_id") + if partition and file_id: + file_infos.add((partition, file_id))And change line 227 to
file_infos = set().tests/api_tests/test_search.py (1)
183-225: Consider extracting the indexing-wait loop into a shared helper.The wait-for-indexing logic (poll task status until SUCCESS/FAILED, with 30 retries and 2s sleep) is duplicated between
indexed_partition(lines 27-46) andindexed_email_thread(lines 205-223). A shared fixture helper would reduce this duplication.openrag/routers/partition.py (2)
413-426: Unusedrequest: Requestparameter.The
requestparameter is injected but never used inget_related_files. Other endpoints in this file use it for URL generation (request.url_for). If HATEOAS links aren't planned for this response, consider removing it.♻️ Suggested change
async def get_related_files( - request: Request, partition: str, relationship_id: str, vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ):
455-478: Unusedrequest: Requestparameter inget_file_ancestors.Same as
get_related_files—requestis injected but unused.♻️ Suggested change
async def get_file_ancestors( - request: Request, partition: str, file_id: str, max_ancestor_depth: int | None = None, vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ):openrag/routers/search.py (1)
230-240: Unusedvectordbdependency and misleading log context insearch_file.
vectordb=Depends(get_vectordb)(line 230) is injected but never used — it was left behind after the no-op expansion removal. This creates an unnecessary Ray actor lookup per request. Additionally, the log bindsinclude_related=Falseandinclude_ancestors=False(lines 238-239) even though those parameters don't exist on this endpoint, which is confusing in log output.Suggested cleanup
async def search_file( request: Request, partition: str, file_id: str, text: str = Query(..., description="Text to search semantically"), top_k: int = Query(5, description="Number of top results to return"), indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ): log = logger.bind( partition=partition, file_id=file_id, query=text, top_k=top_k, - include_related=False, - include_ancestors=False, )openrag/components/test_relationships.py (2)
58-65: Mutable default argumentfile_metadata: dict = None.While benign here (the value is immediately passed to
json.dumps), using a mutable default is a common Python anti-pattern. Considerfile_metadata: dict | None = Nonefor consistency with the productionadd_file_to_partitionsignature.def add_file_to_partition( self, partition: str, file_id: str, - file_metadata: dict = None, + file_metadata: dict | None = None, relationship_id: str = None, parent_id: str = None, ):
22-49: Test helpers duplicate production model and manager logic.
FileModelandPartitionFileManagerHelperreplicate theFilemodel andPartitionFileManagermethods fromutils.py. If the production code evolves (e.g., new fields, changed CTE logic), these test doubles can silently drift. Consider importing and testing the production classes directly against an in-memory SQLite (theSessionis already injectable), or at minimum add a comment noting which production version these mirror.Also applies to: 52-151
openrag/components/indexer/vectordb/utils.py (2)
46-64: Redundant single-column indexes onrelationship_idandparent_id.Both columns have
index=True(lines 48, 51) creating standalone indexes, and are the leading columns in composite indexesix_relationship_partitionandix_parent_partition(lines 63-64). Most databases can satisfy single-column lookups using the composite index when the target column is the leading key. The extra standalone indexes consume storage and slow down writes with no query benefit.Suggested fix — remove the redundant `index=True`
- relationship_id = Column( - String, nullable=True, index=True - ) + relationship_id = Column( + String, nullable=True + ) - parent_id = Column( - String, nullable=True, index=True - ) + parent_id = Column( + String, nullable=True + )
214-215: Nit: type hint ordering convention.
None | stris unusual — PEP 604 and most codebases preferstr | None(non-None type first). The existingfile_metadataparam on line 212 already usesdict | None.- relationship_id: None | str = None, - parent_id: None | str = None, + relationship_id: str | None = None, + parent_id: str | None = None,
dff2005 to
db50a93
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 (1)
openrag/components/utils.py (1)
80-102:⚠️ Potential issue | 🟠 MajorReturn type annotation is incorrect — function now returns
tuple[str, int].The signature declares
-> strbut the function returns a two-tuple(formatted_context, n_docs)on both paths (lines 83 and 102). This will mislead callers and type checkers.-def format_context(docs: list[Document], max_context_tokens: int = 4096) -> str: +def format_context(docs: list[Document], max_context_tokens: int = 4096) -> tuple[str, int]:
🤖 Fix all issues with AI agents
In `@astro.config.mjs`:
- Around line 12-15: The closing brace/paren for the mermaid configuration is
misindented and the theme setting conflicts with autoTheme; adjust indentation
so the closing "})" aligns with the "mermaid({" line and decide which theme
behavior you want—either remove the explicit theme: 'forest' or set autoTheme:
false to enforce the forest theme—by updating the mermaid({ ... }) block (refer
to the mermaid(...) invocation and its theme and autoTheme properties).
In `@openrag/components/reranker.py`:
- Line 20: Replace the falsy check on top_k with an explicit None check: change
the expression that sets top_k ("top_k = min(top_k, len(documents)) if top_k
else len(documents)") to use "if top_k is not None" so a value of 0 is treated
correctly; update the line referencing top_k and documents in reranker.py (the
top_k assignment) to use this explicit None check.
🧹 Nitpick comments (7)
.hydra_config/retriever/base.yaml (1)
8-8: Environment variable nameMAX_DEPTHis too generic.The env var
MAX_DEPTHcould easily collide with other features. Consider renaming toMAX_ANCESTOR_DEPTHto match the config key and avoid ambiguity.-max_ancestor_depth: ${oc.decode:${oc.env:MAX_DEPTH, 10}} # Maximum depth for ancestor retrieval (null = unlimited) +max_ancestor_depth: ${oc.decode:${oc.env:MAX_ANCESTOR_DEPTH, 10}} # Maximum depth for ancestor retrieval (null = unlimited)tests/api_tests/test_search.py (1)
183-225: Consider extracting the wait-for-indexing loop into a shared helper.The indexing + polling pattern here (lines 196–223) is nearly identical to
indexed_folder_partitioninconftest.py. A shared helper likeindex_file_and_wait(api_client, partition, file_id, filename, file_handle, metadata)would reduce duplication and simplify future fixtures.openrag/components/retriever.py (2)
232-240: Deduplicatefile_infosto avoid redundant ancestor fetches.
file_infosis alist, so multiple chunks from the same file will trigger duplicateget_ancestor_chunkscalls. Sincerelationship_idsis aset(already deduplicated), apply the same treatment here.♻️ Proposed fix
# Collect unique relationship_ids and file_ids from results relationship_ids = set() - file_infos = [] # List of (partition, file_id) tuples + file_infos = set() # Set of (partition, file_id) tuples for doc in results: metadata = doc.metadata if include_related and metadata.get("relationship_id"): relationship_ids.add((metadata.get("partition"), metadata.get("relationship_id"))) if include_ancestors: - file_infos.append((metadata.get("partition"), metadata.get("file_id"))) + file_infos.add((metadata.get("partition"), metadata.get("file_id")))
39-40:expand_search_resultsreturnsNoneimplicitly in the base abstract class.
ABCRetriever.expand_search_resultsusespass, which returnsNone. If a subclass forgets to override it, callers will getNoneinstead of the original results. Consider returningresultsas a safe default, or mark it@abstractmethod.♻️ Proposed fix
async def expand_search_results(self, results: list[Document]) -> list[Document]: - pass + return resultsopenrag/routers/search.py (1)
230-240:vectordbdependency is injected but unused insearch_file.
get_vectordbis resolved on every request to this endpoint but never referenced. This adds unnecessary overhead (Ray actor lookup). Either remove it or, if expansion support is planned, add a TODO.♻️ Proposed fix
async def search_file( request: Request, partition: str, file_id: str, text: str = Query(..., description="Text to search semantically"), top_k: int = Query(5, description="Number of top results to return"), indexer=Depends(get_indexer), - vectordb=Depends(get_vectordb), partition_viewer=Depends(require_partition_viewer), ): log = logger.bind( partition=partition, file_id=file_id, query=text, top_k=top_k, - include_related=False, - include_ancestors=False, )openrag/components/test_relationships.py (1)
22-49: Test model diverges from productionFilemodel.
FileModelusesTextforfile_metadata(with manualjson.loads/json.dumps) while the productionFilemodel uses SQLAlchemyJSON. This is a reasonable concession for SQLite, but be aware that behavior differences (e.g., automatic deserialization in production vs. manual in tests) could mask issues. Consider adding a comment noting this intentional divergence.openrag/components/pipeline.py (1)
30-38: Type annotation mismatch:RetrieverFactory.create_retrieverreturnsABCRetriever.Line 32 annotates the retriever as
BaseRetriever, butRetrieverFactory.create_retrieverdeclares its return type asABCRetriever. While this works at runtime (all concrete retrievers inherit fromBaseRetriever), the annotation is stricter than the factory's contract. This could cause type-checker warnings.
db50a93 to
ff5bb7b
Compare
ff5bb7b to
98be85f
Compare
Ahmath-Gadji
added
breaking-change
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>
98be85f to
e4284f9
Compare
2 participants
Relationship-Aware Document Retrieval
This PR adds support for document relationships in OpenRAG, enabling context-aware retrieval across linked files such as email threads and folder-related document groups.
Key Changes
Filedatabase model withrelationship_idandparent_id, plus supporting indexes.BREAKING CHANGE
Perform migration: Refer to the doc
openrag/docs/content/docs/documentation/sql_migration.mdx
Lines 47 to 53 in 0b5f84c
include_relatedinclude_ancestorsOutcome
OpenRAG can now retrieve not only relevant chunks, but also related and contextual documents, improving performance on threaded and grouped datasets.
Summary by CodeRabbit
New Features
Documentation
Tests
Chores