Feat/kb rag retrieval#924
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughAdds optional per-call document scoping by Changes
Sequence DiagramsequenceDiagram
participant Client
participant Endpoint as RAG Endpoint
participant Resolve as Name Resolver
participant KBService as KnowledgeService
participant Retrieval as Retrieval Engine
participant Meta as KB Meta Builder
Note over Client,Endpoint: Scoped retrieval request
Client->>Endpoint: POST internal_retrieve(document_names)
Endpoint->>Resolve: _resolve_document_names(names, kb_ids)
Resolve->>KBService: resolve_document_ids_by_names(db, kb_ids, names)
KBService-->>Resolve: [resolved_ids] / []
Resolve-->>Endpoint: resolved_document_ids
alt resolved_ids non-empty
Endpoint->>Retrieval: query(resolved_document_ids)
Retrieval-->>Endpoint: chunks/results
Endpoint-->>Client: 200 OK with results
else no ids found
Endpoint-->>Client: 200 OK with mode="rag_retrieval", records=[], message="Document names not found..."
end
Note over Client,Meta: KB metadata enrichment
Client->>Meta: _build_kb_meta_prompt(kb_ids)
Meta->>KBService: get_document_prompt_stats(db, kb_ids)
KBService-->>Meta: {kb_id: {total_count, searchable_count, spreadsheet_count}}
Meta-->>Client: kb_meta_list (includes stats)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 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]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
backend/app/api/endpoints/internal/rag.py (1)
216-230: Consider returningInternalRetrieveResponseinstance for type consistency.The early return at lines 224-230 uses a plain dict, while the normal path at lines 353-368 returns an
InternalRetrieveResponseinstance. While FastAPI handles dict-to-model conversion, using the model consistently improves type safety and makes the code more explicit.♻️ Optional: Return model instance for consistency
if not resolved_document_ids: - return { - "mode": "rag_retrieval", - "records": [], - "total": 0, - "total_estimated_tokens": 0, - "message": "Document names not found in the selected knowledge bases. Use kb_ls to inspect available documents first.", - } + return InternalRetrieveResponse( + mode="rag_retrieval", + records=[], + total=0, + total_estimated_tokens=0, + message="Document names not found in the selected knowledge bases. Use kb_ls to inspect available documents first.", + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/api/endpoints/internal/rag.py` around lines 216 - 230, The early return after resolving document names currently returns a plain dict; change it to return an InternalRetrieveResponse instance for consistency and type safety. In the branch that runs when resolved_document_ids is empty (after calling _resolve_document_names with db, knowledge_base_ids, request.document_names), construct and return InternalRetrieveResponse(mode="rag_retrieval", records=[], total=0, total_estimated_tokens=0, message="Document names not found in the selected knowledge bases. Use kb_ls to inspect available documents first.") instead of the dict so the function consistently returns the same model type as the normal path.backend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.py (1)
116-144: Test could verify resolution arguments for stronger assertions.The test verifies that
_resolve_document_namesis called, but doesn't assert the arguments passed. Consider verifying the call arguments to ensure the correct KB IDs and document names are forwarded.♻️ Optional improvement for stronger test assertion
assert response.status_code == 200 - mock_resolve.assert_called_once() + mock_resolve.assert_called_once_with( + db=ANY, + knowledge_base_ids=[12], + document_names=["release.md"], + ) mock_query.assert_awaited_once()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.py` around lines 116 - 144, Update the test_internal_retrieve_resolves_document_names_before_query test to assert the arguments passed to the mocked _resolve_document_names call; after posting to "/api/internal/rag/retrieve" assert mock_resolve.assert_called_once_with(knowledge_base_ids=[12], document_names=["release.md"]) (or equivalent tuple/kwargs matching how _resolve_document_names is invoked) so the test verifies the correct KB IDs and document names are forwarded; keep the existing assertions for response.status_code and mock_query.assert_awaited_once().
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/app/services/knowledge/knowledge_service.py`:
- Around line 892-916: The resolve_document_ids_by_names method currently
returns document IDs without checking document active state; update the DB query
inside resolve_document_ids_by_names to add a filter on
KnowledgeDocument.is_active == True so only active documents are resolved
(mirror the active filtering used in other helpers like
get_active_document_text_length_stats); locate the query using
db.query(KnowledgeDocument.id) and augment its .filter(...) call to include the
is_active condition.
In `@chat_shell/chat_shell/tools/builtin/knowledge_base.py`:
- Around line 481-496: The _apply_call_scoped_filters context manager mutates
instance state (self.document_ids / self.document_names), which is unsafe for
async/concurrent executions; change callers to accept and forward scoped filter
parameters instead of relying on mutated attributes (stop using
_apply_call_scoped_filters), or replace the mutable approach with an async-aware
solution using contextvars to store per-call document_ids/document_names; update
methods that currently read self.document_ids/self.document_names to accept
optional parameters (e.g., document_ids/document_names) and thread the scoped
values through the call chain, or create contextvar-backed getters/setters and
use those in place of direct attribute access.
---
Nitpick comments:
In `@backend/app/api/endpoints/internal/rag.py`:
- Around line 216-230: The early return after resolving document names currently
returns a plain dict; change it to return an InternalRetrieveResponse instance
for consistency and type safety. In the branch that runs when
resolved_document_ids is empty (after calling _resolve_document_names with db,
knowledge_base_ids, request.document_names), construct and return
InternalRetrieveResponse(mode="rag_retrieval", records=[], total=0,
total_estimated_tokens=0, message="Document names not found in the selected
knowledge bases. Use kb_ls to inspect available documents first.") instead of
the dict so the function consistently returns the same model type as the normal
path.
In `@backend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.py`:
- Around line 116-144: Update the
test_internal_retrieve_resolves_document_names_before_query test to assert the
arguments passed to the mocked _resolve_document_names call; after posting to
"/api/internal/rag/retrieve" assert
mock_resolve.assert_called_once_with(knowledge_base_ids=[12],
document_names=["release.md"]) (or equivalent tuple/kwargs matching how
_resolve_document_names is invoked) so the test verifies the correct KB IDs and
document names are forwarded; keep the existing assertions for
response.status_code and mock_query.assert_awaited_once().
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 710ef808-9475-4dce-ab88-237ab1db3688
📒 Files selected for processing (12)
backend/app/api/endpoints/internal/rag.pybackend/app/services/chat/preprocessing/contexts.pybackend/app/services/chat/preprocessing/kb_meta.pybackend/app/services/knowledge/knowledge_service.pybackend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.pybackend/tests/services/rag/test_retrieval_service.pybackend/tests/services/test_kb_meta_prompt.pychat_shell/chat_shell/tools/builtin/knowledge_base.pychat_shell/tests/test_knowledge_base_scoped_search.pydocs/superpowers/specs/2026-04-04-kb-rag-retrieval-design.mdshared/prompts/knowledge_base.pyshared/tests/test_prompts.py
Contributor
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)
chat_shell/chat_shell/tools/builtin/knowledge_base.py (1)
940-944:⚠️ Potential issue | 🟡 MinorImportError fallback doesn't forward document filters.
When the
ImportErrorexception is caught (lines 939-944), the fallback call to_retrieve_with_strategy_via_httpdoesn't passdocument_idsordocument_names. This means scoped document filtering would silently fail in this code path.🔧 Proposed fix to forward filters in fallback
except ImportError: result = await self._retrieve_with_strategy_via_http( query=query, max_results=max_results, route_mode=route_mode, + document_ids=resolved_document_ids, + document_names=document_names, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@chat_shell/chat_shell/tools/builtin/knowledge_base.py` around lines 940 - 944, The ImportError except block currently calls _retrieve_with_strategy_via_http without forwarding document_ids and document_names, which drops scoped filters; update the fallback call inside the ImportError handler to pass through the same document_ids and document_names parameters (alongside query, max_results, route_mode) so the fallback path honors document filtering in _retrieve_with_strategy_via_http.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/app/services/knowledge/knowledge_service.py`:
- Around line 866-876: The spreadsheet_count aggregation currently counts all
KnowledgeDocument.file_extension matches regardless of is_active; update the
CASE predicate used for spreadsheet_count in the method (the prompt-oriented
document stats block in knowledge_service.py) to require the document be active
as well (e.g., include KnowledgeDocument.is_active == True alongside the
file_extension IN spreadsheet_exts) so it mirrors the filtering behavior of
searchable_document_count and only counts active spreadsheet documents.
---
Outside diff comments:
In `@chat_shell/chat_shell/tools/builtin/knowledge_base.py`:
- Around line 940-944: The ImportError except block currently calls
_retrieve_with_strategy_via_http without forwarding document_ids and
document_names, which drops scoped filters; update the fallback call inside the
ImportError handler to pass through the same document_ids and document_names
parameters (alongside query, max_results, route_mode) so the fallback path
honors document filtering in _retrieve_with_strategy_via_http.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 32490fe3-ce6e-48c4-89e6-6ca8885e929e
📒 Files selected for processing (5)
backend/app/api/endpoints/internal/rag.pybackend/app/services/knowledge/knowledge_service.pybackend/tests/services/knowledge/test_knowledge_service.pychat_shell/chat_shell/tools/builtin/knowledge_base.pychat_shell/tests/test_knowledge_base_scoped_search.py
✅ Files skipped from review due to trivial changes (1)
- backend/tests/services/knowledge/test_knowledge_service.py
🚧 Files skipped from review as they are similar to previous changes (1)
- backend/app/api/endpoints/internal/rag.py
kissghosts
force-pushed
the
feat/kb-rag-retrieval
branch
from
96a039a to
9ac8342
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
backend/app/services/knowledge/knowledge_service.py (1)
857-857: Consider extracting spreadsheet extensions to a class constant.The
spreadsheet_extslist is embedded within the method. Extracting it to a class-level constant would improve maintainability and make it reusable if needed elsewhere.♻️ Proposed refactor
Add near line 961 (alongside
NOTEBOOK_MAX_DOCUMENTS):SPREADSHEET_EXTENSIONS = ["csv", "xls", "xlsx"]Then update line 857:
- spreadsheet_exts = ["csv", "xls", "xlsx"] + spreadsheet_exts = KnowledgeService.SPREADSHEET_EXTENSIONSAs per coding guidelines: "Python: Extract magic numbers to constants".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/knowledge_service.py` at line 857, Extract the inline list spreadsheet_exts = ["csv", "xls", "xlsx"] into a class-level constant named SPREADSHEET_EXTENSIONS (placed alongside NOTEBOOK_MAX_DOCUMENTS in the same class or module) and then replace the inline spreadsheet_exts usage in the method (where spreadsheet_exts is defined) with SPREADSHEET_EXTENSIONS to remove the magic literal and make it reusable; ensure imports/refs remain unchanged and run tests to verify behavior in KnowledgeService (knowledge_service.py).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@backend/app/services/knowledge/knowledge_service.py`:
- Line 857: Extract the inline list spreadsheet_exts = ["csv", "xls", "xlsx"]
into a class-level constant named SPREADSHEET_EXTENSIONS (placed alongside
NOTEBOOK_MAX_DOCUMENTS in the same class or module) and then replace the inline
spreadsheet_exts usage in the method (where spreadsheet_exts is defined) with
SPREADSHEET_EXTENSIONS to remove the magic literal and make it reusable; ensure
imports/refs remain unchanged and run tests to verify behavior in
KnowledgeService (knowledge_service.py).
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: dfbdd1f1-ddb5-4aa2-96c6-8b45268c83ad
📒 Files selected for processing (4)
backend/app/services/knowledge/knowledge_service.pybackend/tests/services/knowledge/test_knowledge_service.pyshared/prompts/knowledge_base.pyshared/tests/test_prompts.py
✅ Files skipped from review due to trivial changes (1)
- backend/tests/services/knowledge/test_knowledge_service.py
🚧 Files skipped from review as they are similar to previous changes (2)
- shared/tests/test_prompts.py
- shared/prompts/knowledge_base.py
qdaxb
approved these changes
Apr 7, 2026
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
2 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.
Summary by CodeRabbit
New Features
Bug Fixes / UX
Documentation
Tests