Repository navigation
refactor(rag): build modular data plane foundation - #903
Merged
qdaxb merged 9 commits intoApr 3, 2026
Merged
Conversation
Contributor
|
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:
📝 WalkthroughWalkthroughThe PR splits RAG control vs data planes: adds runtime spec models, a runtime resolver, RagGateway protocol and LocalRagGateway, local data‑plane executors (index/query/delete), protected knowledge mediation for restricted mode, and refactors services/endpoints/Chat Shell to route indexing/retrieval through the gateway with persistence/mediation branching. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Chat Shell\n(KnowledgeBaseTool)
participant API as Backend\n(Internal RAG API)
participant Resolver as RagRuntimeResolver
participant Gateway as LocalRagGateway
participant DataPlane as Local Data Plane\n(Retrieval/Indexing)
participant Mediator as Protected\nKnowledge Mediator
participant Persist as Retrieval\nPersistence
Client->>API: POST /internal/rag/retrieve\n(query, kb_ids, restricted_mode, mediation_context)
API->>Resolver: build_query_runtime_spec(...)
Resolver-->>API: QueryRuntimeSpec
API->>Gateway: query(spec, db=db)
Gateway->>DataPlane: query_local(spec, db=db)
DataPlane->>DataPlane: RetrievalService.retrieve_for_chat_shell(...)
DataPlane-->>Gateway: {records, mode, total, total_estimated_tokens}
Gateway-->>API: retrieval result
API->>Persist: persist_retrieval_result(user_subtask_id, records, ...)
Persist-->>API: done
alt restricted_mode == True
API->>Mediator: transform(db, query, records, mediation_context, kb_ids)
Mediator->>Mediator: resolve model & summarize (LLM)
Mediator-->>API: ProtectedKnowledgeMediationResponse
API-->>Client: restricted_safe_summary response
else
API-->>Client: InternalRetrieveResponse (normal)
end
sequenceDiagram
participant Task as Index Task
participant Service as Indexing Service
participant Resolver as RagRuntimeResolver
participant KB as Knowledge DB
participant Gateway as LocalRagGateway
participant DataPlane as Local Data Plane\n(DocumentService)
Task->>Service: run_document_indexing(kb_id, document_id,...)
Service->>Resolver: build_index_runtime_spec(...)
Resolver->>KB: get_kb_index_info(...)
KB-->>Resolver: KnowledgeBaseIndexInfo
Resolver-->>Service: IndexRuntimeSpec
Service->>Gateway: index_document(spec, db=db)
Gateway->>DataPlane: index_document_local(spec, db=db)
DataPlane->>DataPlane: load retriever, create storage backend
DataPlane->>DataPlane: DocumentService.index_document(...)
DataPlane-->>Gateway: {status, ...}
Gateway-->>Service: result
Service-->>Task: {status: success}
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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: 7
🧹 Nitpick comments (3)
backend/app/services/rag/local_data_plane/indexing.py (1)
13-17: Trace this new async executor.
index_document_local()is the new local data-plane entrypoint, but it is not wrapped with the repo-standard async tracing decorator. That will leave this hop out of OpenTelemetry traces at the new boundary.As per coding guidelines, "Use
@trace_async(span_name, tracer_name, extract_attributes)decorator to trace entire async functions in OpenTelemetry".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/rag/local_data_plane/indexing.py` around lines 13 - 17, The async entrypoint index_document_local is missing the repo-standard OpenTelemetry decorator; import and apply the trace_async decorator to that coroutine (e.g., `@trace_async`("local.index_document", "local_data_plane", extract_attributes)) so the entire function is traced; implement or reuse an extract_attributes callable that reads useful fields from IndexRuntimeSpec (the spec param) and pass it as the third argument to trace_async, and ensure trace_async is imported where index_document_local is defined.backend/app/api/endpoints/internal/rag.py (1)
218-236: Avoid depending onRetrievalServiceprivate helpers from the endpoint.This path now executes retrieval through
RagGateway, but Lines 218-236 still reach intoRetrievalService._calculate_*for budget math. Please lift that logic into a small public helper shared by the endpoint and retrieval layer instead of binding the endpoint to private service internals.🤖 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 218 - 236, The endpoint is calling private helpers RetrievalService._calculate_ratio_based_direct_injection_budget and RetrievalService._calculate_available_injection_tokens directly; extract that budget/math logic into a small public helper (e.g., a new function like RetrievalBudget.calculate_available_injection_tokens and RetrievalBudget.calculate_ratio_based_direct_injection_budget or a single public function like RagBudget.calculate_injection_budget) and use this new public helper from both the endpoint (currently calling the underscored methods) and the RetrievalService implementation so the endpoint no longer depends on private internals; update calls in this file to invoke the new public symbol(s) and replace usages inside RetrievalService to delegate to the same helper.backend/app/services/knowledge/index_runtime.py (1)
20-25: Local import to avoid circular dependency is acceptable.The deferred import of
_is_organization_namespaceinside the function body is a valid pattern to prevent circular imports betweenknowledge_serviceandindex_runtime. Consider adding a brief comment explaining this if it's not obvious to future maintainers.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/index_runtime.py` around lines 20 - 25, Add a brief inline comment above the local import in is_organization_namespace explaining that the import of _is_organization_namespace from knowledge_service is deferred to avoid a circular import between index_runtime and knowledge_service; keep the import inside the function and the existing return _is_organization_namespace(db, namespace) call unchanged, just document the reason for the local import referencing both is_organization_namespace and _is_organization_namespace so future maintainers understand the pattern.
🤖 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/indexing.py`:
- Around line 417-420: The current logger.info is printing the full indexing
result (variable result from
loop.run_until_complete(rag_gateway.index_document(...))) which may include
chunks_data and sensitive/large payloads; change the logging to emit only a
small summary (e.g., document_id or runtime_spec id, index_status, and
chunk_count) by extracting non-sensitive keys from result and explicitly
omitting chunks_data while keeping the full result object intact for later use
(leave rag_gateway.index_document and result assignment unchanged, only replace
the logger.info call to log a sanitized summary).
In `@backend/app/services/rag/local_data_plane/indexing.py`:
- Around line 30-34: The data-plane code in indexing.py currently imports
parse_splitter_config from app.services.knowledge.indexing (via splitter_config
= spec.splitter_config; if splitter_config: from app.services.knowledge.indexing
import parse_splitter_config; splitter_config =
parse_splitter_config(splitter_config)), which creates an unwanted dependency on
the control-plane; fix by moving splitter parsing out of
app.services.knowledge.indexing into a shared low-level helper (e.g., a new
function parse_splitter_config in a neutral module like
app.services.shared.splitter or app.utils.splitter) and update indexing.py to
import parse_splitter_config from that shared helper (or alternatively ensure
spec.splitter_config is fully parsed before it reaches the data plane so
indexing.py simply consumes an already-parsed splitter_config). Ensure you
update any callers that previously used the control-plane parse to use the new
shared helper or pre-parse before constructing spec.
In `@backend/app/services/rag/local_data_plane/retrieval.py`:
- Around line 11-35: The async function query_local is missing the required
OpenTelemetry tracing decorator and hardcodes budget default values inline; add
the `@trace_async`(span_name, tracer_name, extract_attributes) decorator above
query_local (using an appropriate span name like "query_local" and the module
tracer) and replace the inlined magic-number defaults (used_context_tokens=0,
reserved_output_tokens=4096, context_buffer_ratio=0.1, max_direct_chunks=500) by
referencing centralized constants (e.g., DEFAULT_USED_CONTEXT_TOKENS,
DEFAULT_RESERVED_OUTPUT_TOKENS, DEFAULT_CONTEXT_BUFFER_RATIO,
DEFAULT_MAX_DIRECT_CHUNKS) so the call to
RetrievalService.retrieve_for_chat_shell uses those constants when budget is
None.
In `@backend/app/services/rag/local_gateway.py`:
- Around line 13-36: Both index_document and query currently allow a db=None
path but LocalRagGateway's underlying executors (_index_executor and
_retrieval_executor) require a Session; update LocalRagGateway.index_document
and LocalRagGateway.query to fail fast when db is None by raising a clear
exception (e.g., ValueError or RuntimeError) instead of calling the db-less
branch, and ensure the error message references the missing db so callers/tests
know to pass the db sentinel; keep the existing code paths that pass db through
to _index_executor and _retrieval_executor unchanged.
In `@backend/app/services/rag/runtime_resolver.py`:
- Around line 57-90: In build_query_runtime_spec, avoid always instantiating
DirectInjectionBudget when context_window can be None: set the
QueryRuntimeSpec.direct_injection_budget to None if context_window is None,
otherwise create a DirectInjectionBudget(context_window=...,
used_context_tokens=..., reserved_output_tokens=..., context_buffer_ratio=...,
max_direct_chunks=...); this uses the existing symbols
(build_query_runtime_spec, QueryRuntimeSpec, DirectInjectionBudget,
context_window) so downstream code expecting an Optional budget receives None
when appropriate.
In `@backend/tests/services/rag/test_runtime_specs.py`:
- Around line 63-83: The parametrized test cases in test_runtime_specs (the
pytest.mark.parametrize block using ("kwargs", "expected_message")) include
hard-coded "/tmp" paths which trigger Ruff S108; update the two occurrences of
"/tmp/a" in the kwargs for source_type "attachment" and "file_path" to use a
neutral relative path (e.g., "path/to/file" or "tests/fixtures/sample.txt") so
the tests remain semantically identical but lint-safe; locate and change the
string values in the parametrized tuple entries where kwargs contains
"file_path": "/tmp/a".
In `@docs/plans/2026-03-24-rag-service-split-plan.md`:
- Around line 13-20: Convert the inline status notes "状态更新（2026-04-01）" and
"当前补充说明" into subsection headings by replacing their plain text lines with "###
状态更新（2026-04-01）" and "### 当前补充说明" respectively so they follow the document
heading hierarchy; ensure the same change is applied to the other occurrences
mentioned (lines 226-230) so the outline uses consistent `###` subsections under
the main sections referenced in this file.
---
Nitpick comments:
In `@backend/app/api/endpoints/internal/rag.py`:
- Around line 218-236: The endpoint is calling private helpers
RetrievalService._calculate_ratio_based_direct_injection_budget and
RetrievalService._calculate_available_injection_tokens directly; extract that
budget/math logic into a small public helper (e.g., a new function like
RetrievalBudget.calculate_available_injection_tokens and
RetrievalBudget.calculate_ratio_based_direct_injection_budget or a single public
function like RagBudget.calculate_injection_budget) and use this new public
helper from both the endpoint (currently calling the underscored methods) and
the RetrievalService implementation so the endpoint no longer depends on private
internals; update calls in this file to invoke the new public symbol(s) and
replace usages inside RetrievalService to delegate to the same helper.
In `@backend/app/services/knowledge/index_runtime.py`:
- Around line 20-25: Add a brief inline comment above the local import in
is_organization_namespace explaining that the import of
_is_organization_namespace from knowledge_service is deferred to avoid a
circular import between index_runtime and knowledge_service; keep the import
inside the function and the existing return _is_organization_namespace(db,
namespace) call unchanged, just document the reason for the local import
referencing both is_organization_namespace and _is_organization_namespace so
future maintainers understand the pattern.
In `@backend/app/services/rag/local_data_plane/indexing.py`:
- Around line 13-17: The async entrypoint index_document_local is missing the
repo-standard OpenTelemetry decorator; import and apply the trace_async
decorator to that coroutine (e.g., `@trace_async`("local.index_document",
"local_data_plane", extract_attributes)) so the entire function is traced;
implement or reuse an extract_attributes callable that reads useful fields from
IndexRuntimeSpec (the spec param) and pass it as the third argument to
trace_async, and ensure trace_async is imported where index_document_local is
defined.
🪄 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: 107414f3-fbb1-4872-b065-93485b37d9da
📒 Files selected for processing (20)
backend/app/api/endpoints/internal/rag.pybackend/app/services/knowledge/index_runtime.pybackend/app/services/knowledge/indexing.pybackend/app/services/rag/README.mdbackend/app/services/rag/gateway.pybackend/app/services/rag/local_data_plane/__init__.pybackend/app/services/rag/local_data_plane/indexing.pybackend/app/services/rag/local_data_plane/retrieval.pybackend/app/services/rag/local_gateway.pybackend/app/services/rag/runtime_resolver.pybackend/app/services/rag/runtime_specs.pybackend/tests/services/knowledge/test_index_runtime.pybackend/tests/services/rag/test_local_gateway.pybackend/tests/services/rag/test_retrieval_service.pybackend/tests/services/rag/test_runtime_resolver.pybackend/tests/services/rag/test_runtime_specs.pybackend/tests/tasks/test_knowledge_tasks.pydocs/plans/2026-03-24-rag-service-split-plan.mddocs/plans/2026-04-01-rag-modular-data-plane-implementation-plan.mddocs/specs/knowledge/2026-03-31-rag-modular-data-plane-design.md
Comment on lines
+63
to
+83
| @pytest.mark.parametrize( | ||
| ("kwargs", "expected_message"), | ||
| [ | ||
| ( | ||
| {"source_type": "attachment"}, | ||
| "attachment_id is required", | ||
| ), | ||
| ( | ||
| {"source_type": "attachment", "attachment_id": 123, "file_path": "/tmp/a"}, | ||
| "file_path must not be provided", | ||
| ), | ||
| ( | ||
| {"source_type": "file_path"}, | ||
| "file_path is required", | ||
| ), | ||
| ( | ||
| {"source_type": "file_path", "file_path": "/tmp/a", "attachment_id": 123}, | ||
| "attachment_id must not be provided", | ||
| ), | ||
| ], | ||
| ) |
Contributor
There was a problem hiding this comment.
Avoid /tmp placeholders in this parametrized block.
These literals trip Ruff S108, so the new test file will fail lint even though the schema checks are correct. A neutral relative path string is enough here.
🧹 Minimal lint-safe adjustment
(
- {"source_type": "attachment", "attachment_id": 123, "file_path": "/tmp/a"},
+ {
+ "source_type": "attachment",
+ "attachment_id": 123,
+ "file_path": "fixtures/a.txt",
+ },
"file_path must not be provided",
),
(
{"source_type": "file_path"},
"file_path is required",
),
(
- {"source_type": "file_path", "file_path": "/tmp/a", "attachment_id": 123},
+ {
+ "source_type": "file_path",
+ "file_path": "fixtures/a.txt",
+ "attachment_id": 123,
+ },
"attachment_id must not be provided",
),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| @pytest.mark.parametrize( | |
| ("kwargs", "expected_message"), | |
| [ | |
| ( | |
| {"source_type": "attachment"}, | |
| "attachment_id is required", | |
| ), | |
| ( | |
| {"source_type": "attachment", "attachment_id": 123, "file_path": "/tmp/a"}, | |
| "file_path must not be provided", | |
| ), | |
| ( | |
| {"source_type": "file_path"}, | |
| "file_path is required", | |
| ), | |
| ( | |
| {"source_type": "file_path", "file_path": "/tmp/a", "attachment_id": 123}, | |
| "attachment_id must not be provided", | |
| ), | |
| ], | |
| ) | |
| `@pytest.mark.parametrize`( | |
| ("kwargs", "expected_message"), | |
| [ | |
| ( | |
| {"source_type": "attachment"}, | |
| "attachment_id is required", | |
| ), | |
| ( | |
| { | |
| "source_type": "attachment", | |
| "attachment_id": 123, | |
| "file_path": "fixtures/a.txt", | |
| }, | |
| "file_path must not be provided", | |
| ), | |
| ( | |
| {"source_type": "file_path"}, | |
| "file_path is required", | |
| ), | |
| ( | |
| { | |
| "source_type": "file_path", | |
| "file_path": "fixtures/a.txt", | |
| "attachment_id": 123, | |
| }, | |
| "attachment_id must not be provided", | |
| ), | |
| ], | |
| ) |
🧰 Tools
🪛 Ruff (0.15.7)
[error] 71-71: Probable insecure usage of temporary file or directory: "/tmp/a"
(S108)
[error] 79-79: Probable insecure usage of temporary file or directory: "/tmp/a"
(S108)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/tests/services/rag/test_runtime_specs.py` around lines 63 - 83, The
parametrized test cases in test_runtime_specs (the pytest.mark.parametrize block
using ("kwargs", "expected_message")) include hard-coded "/tmp" paths which
trigger Ruff S108; update the two occurrences of "/tmp/a" in the kwargs for
source_type "attachment" and "file_path" to use a neutral relative path (e.g.,
"path/to/file" or "tests/fixtures/sample.txt") so the tests remain semantically
identical but lint-safe; locate and change the string values in the parametrized
tuple entries where kwargs contains "file_path": "/tmp/a".
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
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/knowledge_factory.py (1)
95-106:⚠️ Potential issue | 🔴 CriticalUpdate adapters to populate
model_nameandmodel_namespacein KB tool constructor arguments.The KB tool constructor now requires
current_model_nameandcurrent_model_namespaceto build mediation context for restricted mode (line 104-106). However, neither the WeGent adapter (chat_shell/chat_shell/adapters/wegent.py:166-182) nor the backend adapter populate these fields in themodel_configdict they construct. Both paths extract using(model_config or {}).get("model_name"), which returnsNonesince the adapters omit these keys.When
current_model_nameisNone,KnowledgeBaseTool._build_mediation_context()returnsNone, causing restricted backend mediation to lose model identity completely. The test in test_kb_prompt_deduplication.py:243-262 validates this broken state explicitly.Fix by:
- Update
WeGentToResponseAdapter._build_model_config()to includemodel_nameandmodel_namespacefrom agent_config- Update the backend adapter similarly if it also constructs model_config dicts
- Verify that both call paths (chat_shell and backend) receive these values from their request sources
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@chat_shell/chat_shell/tools/knowledge_factory.py` around lines 95 - 106, The KB tool is receiving None for current_model_name/current_model_namespace because adapters don't populate those keys in the model_config dict; update WeGentToResponseAdapter._build_model_config to copy agent_config["model_name"] and agent_config["model_namespace"] (or fallback to agent_config.get("model", {}).get(...)) into the returned model_config, and make the same change in the backend adapter that constructs model_config (ensure it reads model name/namespace from the incoming request/agent_config); after adding model_name and model_namespace keys confirm KnowledgeBaseTool._build_mediation_context now receives non-None values for current_model_name/current_model_namespace.
♻️ Duplicate comments (1)
backend/app/services/rag/local_data_plane/retrieval.py (1)
11-32: 🛠️ Refactor suggestion | 🟠 MajorThis entrypoint still needs tracing and centralized budget defaults.
query_local()is the new async boundary, but it still has no@trace_async, and the inline0 / 4096 / 0.1 / 500fallbacks can drift from the runtime-spec defaults andRetrievalService.retrieve_for_chat_shell()defaults. Please keep those defaults in one shared place instead of re-declaring them here.As per coding guidelines, "Use
@trace_async(span_name, tracer_name, extract_attributes)decorator to trace entire async functions in OpenTelemetry" and "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/rag/local_data_plane/retrieval.py` around lines 11 - 32, Add OpenTelemetry tracing and remove hardcoded budget fallbacks: decorate the async function query_local with `@trace_async` using an appropriate span name (e.g., "query_local") and tracer name and pass an extract_attributes callable that pulls key attributes from the QueryRuntimeSpec (query, user_id, knowledge_base_ids, etc.); then stop using inline magic numbers and instead read the centralized budget defaults (the shared runtime-spec/default budget object used by RetrievalService.retrieve_for_chat_shell or a single exported DEFAULT_RUNTIME_BUDGET constant) for context_window, used_context_tokens, reserved_output_tokens, context_buffer_ratio, and max_direct_chunks so query_local simply forwards spec values or the canonical defaults to RetrievalService.retrieve_for_chat_shell.
🧹 Nitpick comments (4)
backend/app/services/chat/preprocessing/contexts.py (1)
1197-1208: Pull KB tool assembly out of this module.These lines add another responsibility to
_prepare_kb_tools_from_contexts(), which is already far beyond the 50-line preference and lives in a ~1.4k-line module. Extracting the KB-specific resolution/tool-construction path into a dedicated submodule would make changes like these mediation fields much easier to reason about and test.As per coding guidelines, "File size limit: If a file exceeds 1000 lines, split it into multiple sub-modules" and "Function length: Max 50 lines per function (preferred)".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/chat/preprocessing/contexts.py` around lines 1197 - 1208, The _prepare_kb_tools_from_contexts() function is doing KB-specific assembly (creating KnowledgeBaseTool instances) which violates file/function size rules; extract the KB resolution and KnowledgeBaseTool-construction logic into a new submodule (e.g., chat.preprocess.kb_tools or similar) and replace the in-function block that instantiates KnowledgeBaseTool with a call to a new helper function (e.g., build_kb_tools or resolve_kb_tool) that accepts the same parameters (knowledge_base_ids, document_ids, user_id, db, user_subtask_id, model_config, kb_tool_access_mode) and returns the constructed tool(s); update imports in contexts.py to call that helper and add small unit tests for the new helper to verify injection_mode/current_model_name/current_model_namespace handling while keeping _prepare_kb_tools_from_contexts focused and under the size/function-length guidelines.backend/tests/services/rag/test_local_gateway.py (1)
55-66: Strengthen delete delegation assertion to include forwarded arguments.This test currently verifies only invocation count. Please also assert key delegated kwargs (
knowledge_base_id,document_ref, anddb) to catch wiring regressions.Suggested test tightening
`@pytest.mark.asyncio` async def test_local_gateway_delete_document_index_delegates_to_delete_executor(): gateway = LocalRagGateway() gateway._delete_executor = AsyncMock(return_value={"deleted": True}) + db = object() result = await gateway.delete_document_index( knowledge_base_id=1, document_ref="9", - db=object(), + db=db, ) assert result == {"deleted": True} gateway._delete_executor.assert_awaited_once() + kwargs = gateway._delete_executor.await_args.kwargs + assert kwargs["knowledge_base_id"] == 1 + assert kwargs["document_ref"] == "9" + assert kwargs["db"] is db🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/services/rag/test_local_gateway.py` around lines 55 - 66, The test test_local_gateway_delete_document_index_delegates_to_delete_executor only checks that gateway._delete_executor was awaited once; update it to also assert that it was called with the forwarded kwargs by using gateway._delete_executor.assert_awaited_once_with(...) and pass the same knowledge_base_id, document_ref and db values used in the call to LocalRagGateway.delete_document_index so the test verifies proper argument delegation from LocalRagGateway.delete_document_index to the _delete_executor.backend/app/api/endpoints/internal/rag.py (1)
173-338: Theinternal_retrievefunction exceeds the 50-line guideline (~145 lines).Consider extracting logical sections into helper methods to improve readability and testability:
- Runtime spec construction (lines 209-231)
- Logging/statistics calculation (lines 237-284)
- Persistence handling (lines 289-298)
- Restricted mode mediation (lines 300-315)
♻️ Suggested extraction points
async def _build_runtime_spec_from_request( request: InternalRetrieveRequest, knowledge_base_ids: list[int], restricted_mode: bool, ) -> QueryRuntimeSpec: ... async def _handle_persistence( db: Session, persistence_context: RetrievePersistenceContext, query: str, mode: str, records: list, restricted_mode: bool, ) -> None: ... async def _handle_restricted_mediation( db: Session, request: InternalRetrieveRequest, mode: str, records: list, knowledge_base_ids: list[int], total_estimated_tokens: int, ) -> ProtectedKnowledgeMediationResponse: ...As per coding guidelines, "Function length: Max 50 lines per function (preferred)".
🤖 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 173 - 338, internal_retrieve is too long; split responsibilities into helpers: extract runtime spec construction (the block calling runtime_resolver.build_query_runtime_spec and preparing runtime_spec) into e.g. _build_runtime_spec_from_request(request, knowledge_base_ids, restricted_mode), move logging/statistics (calculating total_content_kb, available_for_kb via RetrievalService._calculate_ratio_based_direct_injection_budget, available_injection_tokens, and the logger.info call) into _log_retrieval_stats(runtime_context, records, result, knowledge_base_ids, request), extract persistence into _handle_persistence(db, persistence_context, request.query, mode, records, restricted_mode) which calls retrieval_persistence_service.persist_retrieval_result, and extract restricted mediation into _handle_restricted_mediation(db, request, mode, records, knowledge_base_ids, total_estimated_tokens) which calls protected_knowledge_mediator.transform; update internal_retrieve to call these helpers and return their results so the main function stays concise and under 50 lines.backend/app/services/knowledge/protected_model_resolver.py (1)
167-196: Consider extracting duplicated user model query logic.The user model query pattern (lines 167-180 and 182-196) is nearly identical. This could be consolidated into a helper method to reduce duplication.
♻️ Optional: Extract common query logic
+ def _query_user_model_kind( + self, + *, + db: Session, + model_name: str, + model_namespace: str, + user_id: int, + ) -> Optional[Kind]: + """Query for a user-owned model Kind.""" + return ( + db.query(Kind) + .filter( + Kind.user_id == user_id, + Kind.kind == "Model", + Kind.name == model_name, + Kind.namespace == model_namespace, + Kind.is_active == True, + ) + .first() + )Then use this helper in both
model_type in {"user", "group"}and the fallback cases.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/protected_model_resolver.py` around lines 167 - 196, The two nearly identical queries that build and execute the Kind lookup (currently producing model_kind and user_kind) should be extracted into a single helper, e.g., _get_user_model_kind(db, user_id, model_name, model_namespace) that performs the db.query(Kind).filter(...).first() and returns the Kind or None; replace both usages in the model resolution logic (the block that checks if model_type in {"user","group"} and the fallback block that runs when user_id is not None) to call this helper, and keep the existing normalized_type logic (set to "user" if model_namespace == "default" else "group") when returning from the fallback call. Ensure the helper uses the same filters (Kind.user_id, Kind.kind == "Model", Kind.name, Kind.namespace, Kind.is_active == True) and maintain variable names model_kind/model_kind_result to preserve behavior.
🤖 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/retrieval_persistence.py`:
- Around line 206-207: The current guard in retrieval_persistence.py treats
user_id == 0 as a no-op which conflicts with RetrievePersistenceContext.user_id
(validated ge=0); change the check so zero is allowed by replacing the condition
"user_id is None or user_id <= 0" with "user_id is None or user_id < 0" (or
otherwise ensure you persist when user_id == 0) so valid requests with
user_id==0 are not silently dropped; update any related tests or callers if they
assumed the old behavior.
In `@backend/app/services/rag/local_data_plane/indexing.py`:
- Around line 27-34: index_document_local currently calls
retriever_kinds_service.get_retriever(...) and passes the result directly into
create_storage_backend(retriever), which will throw if the retriever is None;
mirror delete_document_index_local's behavior by explicitly checking the
retriever returned by get_retriever for None before calling
create_storage_backend, and handle the None case by returning/raising the same
structured skip/error path used in delete_document_index_local (i.e., log or
return a skip result and avoid calling create_storage_backend or DocumentService
when retriever is missing).
- Around line 19-23: The two new async module-level entrypoints
index_document_local and delete_document_index_local are not wrapped with the
OpenTelemetry helper, so their latency and skip reasons are not captured; add
the `@trace_async` decorator to both functions (using appropriate span names like
"index_document_local" and "delete_document_index_local", the correct tracer
name used elsewhere in the module, and the same extract_attributes callable used
for other indexing traces) so the entire async function is traced and attributes
are exported at the OTel boundary.
- Around line 68-76: The lookup for Kind currently uses the surrogate Kind.id
(knowledge_base_id) which violates the requirement to identify a Kind by the
three-part identity; change the filter in the query that constructs kb to use
Kind.namespace, Kind.name, and Kind.user_id (e.g., Kind.namespace ==
knowledge_base_namespace, Kind.name == knowledge_base_name, Kind.user_id ==
knowledge_base_user_id) alongside Kind.kind == "KnowledgeBase" and
Kind.is_active == True; ensure you obtain or pass the namespace/name/user_id
values from the surrounding context instead of relying on the id so delete-index
cleanup uses the canonical three-field identity.
In `@backend/tests/services/knowledge/test_index_runtime.py`:
- Line 11: Add explicit return type hints "-> None" to all test function
definitions in this file: update the signature for
test_build_kb_index_info_uses_current_user_for_default_namespace and the other
test functions referenced (the ones at the comment positions: lines
corresponding to 31, 56, 81, 97, 113) so each test function declares "-> None";
ensure you modify the def signatures (e.g., def test_<name>(...): -> def
test_<name>(...) -> None:) rather than touching body logic or imports.
In `@backend/tests/services/knowledge/test_protected_mediation.py`:
- Around line 60-62: The test currently only checks
mock_resolve.assert_called_once() but not which arguments were passed; update
the assertion to verify that mock_resolve was called with the model identity
supplied by mediation_context (the values set on mediation_context at the lines
above) to ensure precedence over fallback resolution—specifically assert the
call args on mock_resolve include the expected model identity (from
mediation_context) and any relevant context object passed to resolve, while
keeping the existing assertions on result.mode and
result.restricted_safe_summary.summary.
---
Outside diff comments:
In `@chat_shell/chat_shell/tools/knowledge_factory.py`:
- Around line 95-106: The KB tool is receiving None for
current_model_name/current_model_namespace because adapters don't populate those
keys in the model_config dict; update
WeGentToResponseAdapter._build_model_config to copy agent_config["model_name"]
and agent_config["model_namespace"] (or fallback to agent_config.get("model",
{}).get(...)) into the returned model_config, and make the same change in the
backend adapter that constructs model_config (ensure it reads model
name/namespace from the incoming request/agent_config); after adding model_name
and model_namespace keys confirm KnowledgeBaseTool._build_mediation_context now
receives non-None values for current_model_name/current_model_namespace.
---
Duplicate comments:
In `@backend/app/services/rag/local_data_plane/retrieval.py`:
- Around line 11-32: Add OpenTelemetry tracing and remove hardcoded budget
fallbacks: decorate the async function query_local with `@trace_async` using an
appropriate span name (e.g., "query_local") and tracer name and pass an
extract_attributes callable that pulls key attributes from the QueryRuntimeSpec
(query, user_id, knowledge_base_ids, etc.); then stop using inline magic numbers
and instead read the centralized budget defaults (the shared
runtime-spec/default budget object used by
RetrievalService.retrieve_for_chat_shell or a single exported
DEFAULT_RUNTIME_BUDGET constant) for context_window, used_context_tokens,
reserved_output_tokens, context_buffer_ratio, and max_direct_chunks so
query_local simply forwards spec values or the canonical defaults to
RetrievalService.retrieve_for_chat_shell.
---
Nitpick comments:
In `@backend/app/api/endpoints/internal/rag.py`:
- Around line 173-338: internal_retrieve is too long; split responsibilities
into helpers: extract runtime spec construction (the block calling
runtime_resolver.build_query_runtime_spec and preparing runtime_spec) into e.g.
_build_runtime_spec_from_request(request, knowledge_base_ids, restricted_mode),
move logging/statistics (calculating total_content_kb, available_for_kb via
RetrievalService._calculate_ratio_based_direct_injection_budget,
available_injection_tokens, and the logger.info call) into
_log_retrieval_stats(runtime_context, records, result, knowledge_base_ids,
request), extract persistence into _handle_persistence(db, persistence_context,
request.query, mode, records, restricted_mode) which calls
retrieval_persistence_service.persist_retrieval_result, and extract restricted
mediation into _handle_restricted_mediation(db, request, mode, records,
knowledge_base_ids, total_estimated_tokens) which calls
protected_knowledge_mediator.transform; update internal_retrieve to call these
helpers and return their results so the main function stays concise and under 50
lines.
In `@backend/app/services/chat/preprocessing/contexts.py`:
- Around line 1197-1208: The _prepare_kb_tools_from_contexts() function is doing
KB-specific assembly (creating KnowledgeBaseTool instances) which violates
file/function size rules; extract the KB resolution and
KnowledgeBaseTool-construction logic into a new submodule (e.g.,
chat.preprocess.kb_tools or similar) and replace the in-function block that
instantiates KnowledgeBaseTool with a call to a new helper function (e.g.,
build_kb_tools or resolve_kb_tool) that accepts the same parameters
(knowledge_base_ids, document_ids, user_id, db, user_subtask_id, model_config,
kb_tool_access_mode) and returns the constructed tool(s); update imports in
contexts.py to call that helper and add small unit tests for the new helper to
verify injection_mode/current_model_name/current_model_namespace handling while
keeping _prepare_kb_tools_from_contexts focused and under the
size/function-length guidelines.
In `@backend/app/services/knowledge/protected_model_resolver.py`:
- Around line 167-196: The two nearly identical queries that build and execute
the Kind lookup (currently producing model_kind and user_kind) should be
extracted into a single helper, e.g., _get_user_model_kind(db, user_id,
model_name, model_namespace) that performs the
db.query(Kind).filter(...).first() and returns the Kind or None; replace both
usages in the model resolution logic (the block that checks if model_type in
{"user","group"} and the fallback block that runs when user_id is not None) to
call this helper, and keep the existing normalized_type logic (set to "user" if
model_namespace == "default" else "group") when returning from the fallback
call. Ensure the helper uses the same filters (Kind.user_id, Kind.kind ==
"Model", Kind.name, Kind.namespace, Kind.is_active == True) and maintain
variable names model_kind/model_kind_result to preserve behavior.
In `@backend/tests/services/rag/test_local_gateway.py`:
- Around line 55-66: The test
test_local_gateway_delete_document_index_delegates_to_delete_executor only
checks that gateway._delete_executor was awaited once; update it to also assert
that it was called with the forwarded kwargs by using
gateway._delete_executor.assert_awaited_once_with(...) and pass the same
knowledge_base_id, document_ref and db values used in the call to
LocalRagGateway.delete_document_index so the test verifies proper argument
delegation from LocalRagGateway.delete_document_index to the _delete_executor.
🪄 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: 8e12bf04-ca43-40c3-abd4-03125ae74e82
📒 Files selected for processing (34)
backend/app/api/endpoints/internal/rag.pybackend/app/services/chat/preprocessing/contexts.pybackend/app/services/knowledge/document_read_service.pybackend/app/services/knowledge/indexing.pybackend/app/services/knowledge/knowledge_service.pybackend/app/services/knowledge/orchestrator.pybackend/app/services/knowledge/protected_mediation.pybackend/app/services/knowledge/protected_model_resolver.pybackend/app/services/knowledge/retrieval_persistence.pybackend/app/services/rag/README.mdbackend/app/services/rag/gateway.pybackend/app/services/rag/local_data_plane/indexing.pybackend/app/services/rag/local_data_plane/retrieval.pybackend/app/services/rag/local_gateway.pybackend/app/services/rag/retrieval_service.pybackend/app/services/rag/splitter/runtime_config.pybackend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.pybackend/tests/services/knowledge/test_document_read_service.pybackend/tests/services/knowledge/test_index_runtime.pybackend/tests/services/knowledge/test_protected_mediation.pybackend/tests/services/knowledge/test_retrieval_persistence.pybackend/tests/services/rag/test_local_gateway.pybackend/tests/services/rag/test_retrieval_service.pybackend/tests/services/test_summary_service.pybackend/tests/services/test_task_knowledge_base_sync.pychat_shell/chat_shell/tools/builtin/knowledge_base.pychat_shell/chat_shell/tools/knowledge_factory.pychat_shell/tests/test_kb_prompt_deduplication.pychat_shell/tests/test_knowledge_base_restricted_backend_mode.pychat_shell/tests/test_knowledge_injection_strategy.pychat_shell/tests/test_knowledge_injection_strategy_clean.pydocs/plans/2026-03-24-rag-service-split-plan.mddocs/specs/knowledge/2026-03-30-wegent-knowledge-document-read-design.mddocs/specs/knowledge/2026-03-31-rag-modular-data-plane-design.md
💤 Files with no reviewable changes (2)
- backend/app/services/rag/retrieval_service.py
- docs/specs/knowledge/2026-03-30-wegent-knowledge-document-read-design.md
✅ Files skipped from review due to trivial changes (6)
- backend/app/services/knowledge/document_read_service.py
- backend/tests/services/knowledge/test_document_read_service.py
- backend/app/services/knowledge/orchestrator.py
- backend/app/services/rag/splitter/runtime_config.py
- backend/app/services/rag/README.md
- docs/plans/2026-03-24-rag-service-split-plan.md
🚧 Files skipped from review as they are similar to previous changes (4)
- backend/app/services/rag/gateway.py
- backend/tests/services/rag/test_retrieval_service.py
- backend/app/services/rag/local_gateway.py
- backend/app/services/knowledge/indexing.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)
backend/app/services/knowledge/indexing.py (1)
397-438:⚠️ Potential issue | 🟠 MajorDon’t coerce non-success gateway results into
"success".
index_document_local()can legitimately return structured skip results likeretriever_not_found, but this wrapper logsresult["status"]and then always returns"status": "success". That makes the task/state-machine layer treat skipped indexing as succeeded and loses the skip reason.Diff
- return { - "status": "success", - "document_id": document_id, - "knowledge_base_id": knowledge_base_id, - "indexed_count": indexed_count, - "index_name": index_name, - "chunks_data": result.get("chunks_data"), - } + if indexing_status != "success": + return { + "status": indexing_status, + "reason": result.get("reason"), + "document_id": document_id, + "knowledge_base_id": knowledge_base_id, + "indexed_count": indexed_count, + "index_name": index_name, + } + + return { + "status": "success", + "document_id": document_id, + "knowledge_base_id": knowledge_base_id, + "indexed_count": indexed_count, + "index_name": index_name, + "chunks_data": result.get("chunks_data"), + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/indexing.py` around lines 397 - 438, The wrapper is coercing every gateway response into "status":"success"; instead propagate the actual gateway result status and related fields instead of hardcoding "success". Update the return value in the function that calls rag_gateway.index_document (and index_document_local()) to set "status": result.get("status", "unknown") (or similar), ensure indexed_count/index_name/chunks_data use values from result (as already done for indexed_count/index_name) and keep add_span_event sending the real status (use indexing_status/result.get("status")). This preserves skip reasons like "retriever_not_found" and avoids masking non-success outcomes.
♻️ Duplicate comments (2)
backend/app/services/rag/local_data_plane/indexing.py (2)
81-89: 🛠️ Refactor suggestion | 🟠 MajorResolve the KB with the canonical Kind identity, not
Kind.id.This delete path still re-hydrates the
KnowledgeBaseby surrogate id. The backend identity contract requiresnamespace,name, anduser_id, so this helper should take those fields (or an already-resolvedKind) instead of querying byKind.id.As per coding guidelines, "Kind resource is uniquely identified by THREE fields:
namespace,name, anduser_id- always use all three fields to locate a Kind".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/rag/local_data_plane/indexing.py` around lines 81 - 89, The code currently re-hydrates the KnowledgeBase using Kind.id (variable knowledge_base_id) which violates the canonical identity rule; change the helper to accept the canonical identity (namespace, name, user_id) or an already-resolved Kind and replace the DB lookup that uses Kind.id with a query filtering on Kind.namespace == namespace, Kind.name == name, and Kind.user_id == user_id (or skip the query if a Kind instance is passed); update any callers of this helper to provide those three fields (or a Kind) and remove reliance on knowledge_base_id in the function that performs the query (refer to variables/function using kb, Kind, and knowledge_base_id to locate the code).
19-23: 🛠️ Refactor suggestion | 🟠 MajorTrace the two new async data-plane entrypoints.
Both async boundaries currently skip the repo’s
@trace_asyncwrapper, so local indexing/delete latency and skip reasons disappear at the function boundary.Based on learnings, "Applies to **/*.py : Use
trace_async(span_name, tracer_name, extract_attributes)decorator to trace entire async functions in OpenTelemetry".Also applies to: 73-79
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/rag/local_data_plane/indexing.py` around lines 19 - 23, The async entrypoint index_document_local is missing the repo's `@trace_async` wrapper so tracing/latency and skip reasons are lost; add the trace_async decorator (with appropriate span_name, tracer_name and extract_attributes function) above index_document_local and also apply the same trace_async decorator to the other async data-plane entrypoint in this file (the async function around lines 73-79) so the entire async functions are captured by OpenTelemetry.
🧹 Nitpick comments (1)
backend/tests/tasks/test_knowledge_tasks.py (1)
9-9: Assert the gateway call carriesdbto lock in the new contract.
LocalRagGateway.index_document()now requires a session, but this test only checks that it was awaited. It would still pass ifrun_document_indexing()stopped forwarding the task-scopeddb.Diff
-from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ - mock_index.assert_awaited_once() + mock_index.assert_awaited_once_with(mock_resolve.return_value, db=ANY)Also applies to: 188-223
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/tasks/test_knowledge_tasks.py` at line 9, The test currently only asserts that LocalRagGateway.index_document() was awaited but not that it received the task-scoped db; update the test in backend/tests/tasks/test_knowledge_tasks.py to assert the gateway call's arguments include the expected session object passed from run_document_indexing(): capture the AsyncMock/patch for LocalRagGateway.index_document (or the gateway instance used in run_document_indexing()), call run_document_indexing() as before, then assert that index_document was awaited with an argument list/kwargs containing the same db/session object (e.g., via mock_call = mock_index_document.await_args or mock_index_document.assert_awaited_once() + inspect of mock_index_document.call_args) so the test fails if run_document_indexing() does not forward the task-scoped db.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@chat_shell/chat_shell/adapters/wegent.py`:
- Around line 177-187: The current nested dict.get calls for model_name and
model_namespace (involving model_spec, agent_config, model_ref) treat an
explicit None as a present value and thus skip intended fallbacks; change the
selection to explicitly prefer the first non-None value (e.g., check
model_spec.get("model_name") is not None, else agent_config.get("model_name") is
not None, else model_ref.get("model_name")) and do the same for model_namespace
with a final default of "default" so None values fall through to the next
candidate.
---
Outside diff comments:
In `@backend/app/services/knowledge/indexing.py`:
- Around line 397-438: The wrapper is coercing every gateway response into
"status":"success"; instead propagate the actual gateway result status and
related fields instead of hardcoding "success". Update the return value in the
function that calls rag_gateway.index_document (and index_document_local()) to
set "status": result.get("status", "unknown") (or similar), ensure
indexed_count/index_name/chunks_data use values from result (as already done for
indexed_count/index_name) and keep add_span_event sending the real status (use
indexing_status/result.get("status")). This preserves skip reasons like
"retriever_not_found" and avoids masking non-success outcomes.
---
Duplicate comments:
In `@backend/app/services/rag/local_data_plane/indexing.py`:
- Around line 81-89: The code currently re-hydrates the KnowledgeBase using
Kind.id (variable knowledge_base_id) which violates the canonical identity rule;
change the helper to accept the canonical identity (namespace, name, user_id) or
an already-resolved Kind and replace the DB lookup that uses Kind.id with a
query filtering on Kind.namespace == namespace, Kind.name == name, and
Kind.user_id == user_id (or skip the query if a Kind instance is passed); update
any callers of this helper to provide those three fields (or a Kind) and remove
reliance on knowledge_base_id in the function that performs the query (refer to
variables/function using kb, Kind, and knowledge_base_id to locate the code).
- Around line 19-23: The async entrypoint index_document_local is missing the
repo's `@trace_async` wrapper so tracing/latency and skip reasons are lost; add
the trace_async decorator (with appropriate span_name, tracer_name and
extract_attributes function) above index_document_local and also apply the same
trace_async decorator to the other async data-plane entrypoint in this file (the
async function around lines 73-79) so the entire async functions are captured by
OpenTelemetry.
---
Nitpick comments:
In `@backend/tests/tasks/test_knowledge_tasks.py`:
- Line 9: The test currently only asserts that LocalRagGateway.index_document()
was awaited but not that it received the task-scoped db; update the test in
backend/tests/tasks/test_knowledge_tasks.py to assert the gateway call's
arguments include the expected session object passed from
run_document_indexing(): capture the AsyncMock/patch for
LocalRagGateway.index_document (or the gateway instance used in
run_document_indexing()), call run_document_indexing() as before, then assert
that index_document was awaited with an argument list/kwargs containing the same
db/session object (e.g., via mock_call = mock_index_document.await_args or
mock_index_document.assert_awaited_once() + inspect of
mock_index_document.call_args) so the test fails if run_document_indexing() does
not forward the task-scoped db.
🪄 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: f9b1249c-d7c4-42a2-bb35-de41b26665a3
📒 Files selected for processing (14)
backend/app/services/knowledge/indexing.pybackend/app/services/knowledge/retrieval_persistence.pybackend/app/services/rag/local_data_plane/indexing.pybackend/app/services/rag/local_gateway.pybackend/app/services/rag/runtime_resolver.pybackend/tests/services/knowledge/test_index_runtime.pybackend/tests/services/knowledge/test_protected_mediation.pybackend/tests/services/knowledge/test_retrieval_persistence.pybackend/tests/services/rag/test_local_data_plane_indexing.pybackend/tests/services/rag/test_local_gateway.pybackend/tests/services/rag/test_runtime_resolver.pybackend/tests/tasks/test_knowledge_tasks.pychat_shell/chat_shell/adapters/wegent.pychat_shell/tests/test_wegent_adapter.py
✅ Files skipped from review due to trivial changes (4)
- chat_shell/tests/test_wegent_adapter.py
- backend/tests/services/knowledge/test_retrieval_persistence.py
- backend/tests/services/knowledge/test_protected_mediation.py
- backend/tests/services/rag/test_runtime_resolver.py
🚧 Files skipped from review as they are similar to previous changes (3)
- backend/tests/services/rag/test_local_gateway.py
- backend/tests/services/knowledge/test_index_runtime.py
- backend/app/services/knowledge/retrieval_persistence.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
backend/app/services/knowledge/protected_mediation.py (1)
108-151: Trace this async control-plane entrypoint.This path resolves a model, does DB-backed lookups, calls an LLM, and has multiple fallback branches, but none of it is wrapped in an OpenTelemetry span. Please add
@trace_async(...)here and emit span attributes/events for the selected model and fallback reason.As per coding guidelines,
Use@trace_async(span_name, tracer_name, extract_attributes) decorator to trace entire async functions in OpenTelemetry,Use add_span_event(name, attributes) to add event to current OpenTelemetry span, andUse set_span_attribute(key, value) to set attribute on current OpenTelemetry span.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/protected_mediation.py` around lines 108 - 151, Wrap the async transform method with `@trace_async`("protected_mediation.transform", "knowledge_service", extract_attributes=...) to trace the entire control path and include basic attrs (user_id, knowledge_base_ids, retrieval_mode, total_estimated_tokens); after resolving model_config (the call to resolve_model_config) call set_span_attribute("selected_model", <model identifier from model_config>) to record the chosen model; inside the summarization/fallback path (the call to _summarize_records and any returned fallback metadata) call add_span_event("model_fallback", {"reason": <fallback reason>}) and set_span_attribute("fallback_used", True/False) as appropriate; ensure attribute keys reference model_config and the summarization result fields so the span captures both selected model and fallback reason.
🤖 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/protected_mediation.py`:
- Around line 126-143: The code currently calls resolve_model_config(...) and
then passes raw records (which may contain protected chunks) into
_summarize_records(...) that ultimately invokes
LangChainModelFactory.create_from_config(...).ainvoke(...); to fix, enforce a
trust/allowlist check before summarization: either (A) restrict model_config to
an internal/trusted model Kind (e.g., only allow provider kinds in an allowlist)
by validating model_config.kind returned by resolve_model_config(...) and
replacing/overriding it with a safe internal model if not allowed, or (B)
strip/filter out protected chunks from records before calling
_summarize_records(...); implement the check at the call site that builds
model_config/records (use symbols resolve_model_config, mediation_context,
_summarize_records, and LangChainModelFactory.create_from_config) so protected
KB content is never sent to untrusted provider-backed models.
In `@backend/app/services/knowledge/protected_model_resolver.py`:
- Around line 125-132: The returned config currently mixes the requested ref
with the resolved Kind because callers pass the original
model_name/model_namespace into _build_model_config while _lookup_model_kind may
have fallen back; change the call sites (the places that call _lookup_model_kind
and then call _build_model_config) to pass the resolved Kind metadata (use the
resolved object's name/namespace/type returned by
_lookup_model_kind/_lookup_model_kind result) into _build_model_config instead
of the original requested model_name/model_namespace, or alternatively modify
_build_model_config to prefer/overwrite with the resolved Kind values rather
than only using setdefault; update all similar call sites (the other occurrences
that call _lookup_model_kind then _build_model_config) to keep metadata
consistent.
- Around line 76-85: The mediation branch currently resolves the active model
using mediation_context["current_model_name"] and model_namespace but hardcodes
model_type=None in _resolve_named_model, which can collide when a user/group
Kind shares a name with a public model; update the mediation_context to carry
"current_model_type" (from the internal RAG endpoint where current_model_type
should be added alongside current_model_name/current_model_namespace), and
thread that value into the call to _resolve_named_model so it forwards the type
into _lookup_model_kind (replace the hardcoded None with
mediation_context.get("current_model_type")). Apply the same change to the other
similar branch later in this file that uses _resolve_named_model.
- Around line 35-40: Replace the boolean equality checks using "Kind.is_active
== True" with the boolean column directly ("Kind.is_active") in all filter
predicates where it appears (e.g., the db.query(Kind).filter(...) call shown and
the other three occurrences flagged at lines 208, 231, 267); update the
predicates that reference Kind.is_active to use Kind.is_active instead of
comparing to True so the filters and SQLAlchemy expressions use the boolean
column directly.
In `@backend/app/services/rag/local_data_plane/indexing.py`:
- Around line 120-128: Replace the equality comparison on the boolean column
with a direct boolean expression: in the query that builds kb (using
db.query(Kind).filter(...).first()), remove "Kind.is_active == True" and use
"Kind.is_active" instead so the filter reads Kind.id == knowledge_base_id,
Kind.kind == "KnowledgeBase", Kind.is_active; this fixes the Ruff E712 lint
error while keeping the same semantics.
In `@backend/tests/services/knowledge/test_protected_model_resolver.py`:
- Around line 59-107: The test currently masks deduplication because
_resolve_summary_or_system_fallback returns on the first successful lookup;
change the patched _lookup_model_kind to simulate a miss on the first unique ref
and a success on a later unique ref (use patch.object(resolver,
"_lookup_model_kind", side_effect=[None, (model_kind, "public")]) or similar) so
the code must continue past the first ref and exercise the dedupe set
(seen_refs) when encountering the duplicate ref; also update assertions to
expect multiple calls (e.g., assert_called_with for the successful lookup and
assert the call count matches unique refs visited) so the deduplication behavior
is actually verified.
---
Nitpick comments:
In `@backend/app/services/knowledge/protected_mediation.py`:
- Around line 108-151: Wrap the async transform method with
`@trace_async`("protected_mediation.transform", "knowledge_service",
extract_attributes=...) to trace the entire control path and include basic attrs
(user_id, knowledge_base_ids, retrieval_mode, total_estimated_tokens); after
resolving model_config (the call to resolve_model_config) call
set_span_attribute("selected_model", <model identifier from model_config>) to
record the chosen model; inside the summarization/fallback path (the call to
_summarize_records and any returned fallback metadata) call
add_span_event("model_fallback", {"reason": <fallback reason>}) and
set_span_attribute("fallback_used", True/False) as appropriate; ensure attribute
keys reference model_config and the summarization result fields so the span
captures both selected model and fallback reason.
🪄 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: 1b72b242-79ec-4c7a-84ee-9e21fa16374b
📒 Files selected for processing (12)
backend/app/services/knowledge/protected_mediation.pybackend/app/services/knowledge/protected_model_resolver.pybackend/app/services/rag/README.mdbackend/app/services/rag/local_data_plane/indexing.pybackend/app/services/rag/local_data_plane/retrieval.pybackend/app/services/rag/runtime_specs.pybackend/tests/services/knowledge/test_protected_mediation.pybackend/tests/services/knowledge/test_protected_model_resolver.pybackend/tests/services/rag/test_local_data_plane_retrieval.pychat_shell/chat_shell/adapters/wegent.pychat_shell/tests/test_wegent_adapter.pydocs/plans/2026-03-24-rag-service-split-plan.md
✅ Files skipped from review due to trivial changes (1)
- backend/app/services/rag/README.md
🚧 Files skipped from review as they are similar to previous changes (5)
- chat_shell/tests/test_wegent_adapter.py
- chat_shell/chat_shell/adapters/wegent.py
- backend/tests/services/knowledge/test_protected_mediation.py
- docs/plans/2026-03-24-rag-service-split-plan.md
- backend/app/services/rag/local_data_plane/retrieval.py
Comment on lines
+76
to
+85
| if mediation_context and mediation_context.get("current_model_name"): | ||
| current_model = self._resolve_named_model( | ||
| db=db, | ||
| model_name=mediation_context["current_model_name"], | ||
| model_namespace=mediation_context.get( | ||
| "current_model_namespace", "default" | ||
| ), | ||
| user_id=user_id, | ||
| user_name=user_name, | ||
| ) |
Contributor
There was a problem hiding this comment.
The current-model path cannot disambiguate scope collisions.
backend/app/api/endpoints/internal/rag.py:92-103 only supplies current_model_name and current_model_namespace, and _resolve_named_model() hardcodes model_type=None. If a user/group Kind shares a name with the active public model, this branch can silently resolve the wrong config. Please thread current_model_type through the mediation context and pass it into _lookup_model_kind().
Also applies to: 116-121
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/app/services/knowledge/protected_model_resolver.py` around lines 76 -
85, The mediation branch currently resolves the active model using
mediation_context["current_model_name"] and model_namespace but hardcodes
model_type=None in _resolve_named_model, which can collide when a user/group
Kind shares a name with a public model; update the mediation_context to carry
"current_model_type" (from the internal RAG endpoint where current_model_type
should be added alongside current_model_name/current_model_namespace), and
thread that value into the call to _resolve_named_model so it forwards the type
into _lookup_model_kind (replace the hardcoded None with
mediation_context.get("current_model_type")). Apply the same change to the other
similar branch later in this file that uses _resolve_named_model.
Comment on lines
+59
to
+107
| with ( | ||
| patch.object( | ||
| resolver, | ||
| "_lookup_model_kind", | ||
| return_value=(model_kind, "public"), | ||
| ) as mock_lookup, | ||
| patch( | ||
| "app.services.knowledge.protected_model_resolver.extract_and_process_model_config", | ||
| return_value={"model_id": "gpt-summary"}, | ||
| ), | ||
| ): | ||
| result = resolver._resolve_summary_or_system_fallback( | ||
| db=db, | ||
| knowledge_base_ids=[1, 2], | ||
| knowledge_base_snapshots=[ | ||
| { | ||
| "id": 1, | ||
| "name": "KB-1", | ||
| "summary_model_ref": { | ||
| "name": "summary-model", | ||
| "namespace": "default", | ||
| "type": "public", | ||
| }, | ||
| }, | ||
| { | ||
| "id": 2, | ||
| "name": "KB-2", | ||
| "summary_model_ref": { | ||
| "name": "summary-model", | ||
| "namespace": "default", | ||
| "type": "public", | ||
| }, | ||
| }, | ||
| ], | ||
| user_id=7, | ||
| user_name="alice", | ||
| ) | ||
|
|
||
| mock_lookup.assert_called_once_with( | ||
| db=db, | ||
| model_name="summary-model", | ||
| model_namespace="default", | ||
| model_type="public", | ||
| user_id=7, | ||
| ) | ||
| assert result is not None | ||
| assert result["model_id"] == "gpt-summary" | ||
| assert result["model_name"] == "summary-model" | ||
| assert result["model_namespace"] == "default" |
Contributor
There was a problem hiding this comment.
This test still passes if the dedupe set is removed.
_resolve_summary_or_system_fallback() returns on the first successful lookup, so the second identical ref is never visited here. assert_called_once_with() is only proving early return, not seen_refs deduplication. Make the first unique lookup miss, or resolve on a later unique ref, so the duplicate path is actually exercised.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/tests/services/knowledge/test_protected_model_resolver.py` around
lines 59 - 107, The test currently masks deduplication because
_resolve_summary_or_system_fallback returns on the first successful lookup;
change the patched _lookup_model_kind to simulate a miss on the first unique ref
and a success on a later unique ref (use patch.object(resolver,
"_lookup_model_kind", side_effect=[None, (model_kind, "public")]) or similar) so
the code must continue past the first ref and exercise the dedupe set
(seen_refs) when encountering the duplicate ref; also update assertions to
expect multiple calls (e.g., assert_called_with for the successful lookup and
assert the call count matches unique refs visited) so the deduplication behavior
is actually verified.
kissghosts
force-pushed
the
refact/rag-modular-data-plane-design
branch
from
April 3, 2026 07:52
6604ffc to
78d9eca
Compare
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)
backend/app/services/knowledge/indexing.py (1)
337-344:⚠️ Potential issue | 🟡 MinorInconsistent return shape: missing
chunks_datakey.The early skip return here is missing the
chunks_datakey, making it inconsistent with the gateway-path return at lines 431-439. Callers that expect a uniform response shape may break when handling this case.🔧 Proposed fix for consistent return shape
return { "status": "skipped", "reason": skip_reason, "document_id": document_id, "knowledge_base_id": knowledge_base_id, "indexed_count": 0, "index_name": None, + "chunks_data": None, }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/indexing.py` around lines 337 - 344, The early skip return object is missing the chunks_data key, causing an inconsistent response shape compared to the gateway-path return that includes chunks_data; update the early-return dictionary (the block that returns status, reason, document_id, knowledge_base_id, indexed_count, index_name) to also include "chunks_data" (e.g., an empty list) so callers reading skip_reason/document_id/knowledge_base_id receive a uniform response shape; reference the variables skip_reason, document_id, knowledge_base_id in the same return object when making this change.
🧹 Nitpick comments (7)
backend/app/services/knowledge/index_runtime.py (1)
61-69: UseKind.is_activeinstead ofKind.is_active == True.SQLAlchemy column expressions are truthy, so the explicit comparison to
Trueis unnecessary and flagged by Ruff (E712).♻️ Proposed fix
kb = ( db.query(Kind) .filter( Kind.id == kb_id, Kind.kind == "KnowledgeBase", - Kind.is_active == True, + Kind.is_active, ) .first() )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/index_runtime.py` around lines 61 - 69, The filter on the Kind query uses an unnecessary boolean comparison; replace the expression Kind.is_active == True with the column expression Kind.is_active in the kb query (the block that builds kb via db.query(Kind).filter(...).first()) so the filter reads Kind.is_active instead of comparing to True; update any similar filters in the same file that use == True to follow the same pattern and run tests/linting to confirm Ruff E712 is resolved.backend/tests/services/test_summary_service.py (1)
402-430: Add a companion test for the non-default namespace ownership branch.This test only validates the
"default"namespace path forindex_owner_user_id. Please add one more case with a non-default, non-organization namespace and assertindex_owner_user_id == test_knowledge_base.user_idto cover the alternate branch inKnowledgeService.delete_document.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/tests/services/test_summary_service.py` around lines 402 - 430, Add a new test case that exercises the non-default/non-organization namespace branch in KnowledgeService.delete_document: clone the existing block that patches LocalRagGateway.delete_document_index (AsyncMock) and set test_knowledge_base.json["spec"]["retrievalConfig"]["retriever_namespace"] to a non-default value (e.g., "custom-namespace"), commit with flag_modified/test_db.commit(), call KnowledgeService.delete_document with the same test_db, test_document.id, and test_user.id, then assert result.success is True and mock_delete_index was awaited once and that mock_delete_index.await_args.kwargs["index_owner_user_id"] == test_knowledge_base.user_id (also keep assertions for knowledge_base_id, document_ref, and db).backend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.py (1)
83-113: Consider strengthening the gateway call assertion.Line 112 uses
mock_query.assert_awaited_once_with(ANY, db=ANY), which verifies the call happened but doesn't validate that theQueryRuntimeSpecargument contains the expected values (e.g.,knowledge_base_ids=[1],restricted_mode=False). This could mask regressions if the spec construction changes.Example of a more specific assertion
# After the response check, verify the spec argument: call_args = mock_query.await_args query_spec = call_args[0][0] # First positional arg assert query_spec.knowledge_base_ids == [1] assert query_spec.restricted_mode is False🤖 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 83 - 113, The gateway assertion is too loose; after calling the endpoint in test_internal_retrieve_keeps_user_subtask_id_out_of_gateway, inspect the actual QueryRuntimeSpec passed to the patched LocalRagGateway.query (mock_query) instead of only using ANY: retrieve the awaited call arguments from mock_query (e.g., mock_query.await_args or mock_query.assert_awaited_once() then access mock_query.call_args) and assert that the first positional argument (the QueryRuntimeSpec) has the expected fields like knowledge_base_ids == [1] and restricted_mode is False; keep the existing response and persist_retrieval_result assertions but replace or extend mock_query.assert_awaited_once_with(ANY, db=ANY) with these explicit checks on the QueryRuntimeSpec properties.backend/app/services/knowledge/protected_model_resolver.py (2)
76-85: Threadingcurrent_model_typewould improve resolution precision.When
mediation_contextcontainscurrent_model_name, this method calls_resolve_named_model()without passing a model type (line 120 hardcodesmodel_type=None). If the caller can providecurrent_model_typein the mediation context, it should be threaded through here to avoid ambiguous resolution between user/group and public scopes.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/protected_model_resolver.py` around lines 76 - 85, The code path that resolves a current model when mediation_context contains "current_model_name" omits passing through "current_model_type", causing ambiguous resolution across user/group/public scopes; update the call to _resolve_named_model in the mediation_context branch to pass model_type=mediation_context.get("current_model_type") (or default None) so the resolver uses the provided type when present (reference: mediation_context, "current_model_name", "current_model_type", and _resolve_named_model).
43-48: Consider narrowing the exception type.Catching blind
Exceptionis flagged by static analysis (Ruff BLE001). While the defensive approach is reasonable here for a best-effort loader, consider catchingSQLAlchemyErrorspecifically for clarity.Suggested change
+from sqlalchemy.exc import SQLAlchemyError + ... - except Exception: + except SQLAlchemyError: logger.debug( "[protected_model_resolver] Failed to load KB snapshots", exc_info=True, ) return []🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/protected_model_resolver.py` around lines 43 - 48, The except Exception in the KB snapshot loader should be narrowed to SQLAlchemyError to satisfy static analysis and be explicit about expected DB errors; update the exception handler in the function/method that loads KB snapshots in protected_model_resolver.py (the block currently using "except Exception") to "except SQLAlchemyError as e" and add "from sqlalchemy.exc import SQLAlchemyError" to imports, preserving the logger.debug call (you can keep exc_info=True or include e in the message) and still return [] on failure.backend/app/api/endpoints/internal/rag.py (1)
92-102: Consider addingcurrent_model_typeto disambiguate scope collisions.The
RetrieveMediationContextonly suppliescurrent_model_nameandcurrent_model_namespace. If a user/groupKindshares a name with a public model, the resolution inprotected_model_resolver._resolve_named_model()(which hardcodesmodel_type=None) may silently resolve the wrong config. Threadingcurrent_model_typethrough the mediation context would enable precise resolution.Suggested addition
class RetrieveMediationContext(BaseModel): """Model identity used by Backend-side protected mediation.""" current_model_name: Optional[str] = Field( default=None, description="Current answering model name preferred for protected mediation", ) current_model_namespace: Optional[str] = Field( default="default", description="Namespace of the current answering model", ) + current_model_type: Optional[str] = Field( + default=None, + description="Model type (public, user, group) for precise resolution", + )🤖 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 92 - 102, The mediation context lacks model type info causing ambiguous resolution in protected_model_resolver._resolve_named_model() which currently uses model_type=None; add an Optional[str] field current_model_type to the RetrieveMediationContext Pydantic model (alongside current_model_name and current_model_namespace), update any callers/constructors that build RetrieveMediationContext to pass the model type when available, and modify protected_model_resolver._resolve_named_model() (or its invocation sites) to use context.current_model_type instead of hardcoding None so resolution is precise between user/group Kinds and public models.backend/app/services/knowledge/protected_mediation.py (1)
406-423: Consider non-greedy regex for nested JSON robustness.The regex
r"\{.*\}"withDOTALLuses greedy matching, which works for single JSON objects but could have edge cases with malformed output containing multiple braces. For this controlled output format, it's acceptable, but a non-greedy pattern or a proper JSON scanner would be more robust.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/app/services/knowledge/protected_mediation.py` around lines 406 - 423, In _extract_json_payload, replace the greedy regex r"\{.*\}" with a non-greedy pattern (e.g. r"\{.*?\}" with flags=re.DOTALL) so the fallback extraction captures the first balanced-looking JSON object instead of spanning multiple brace groups; update the json.loads(match.group(0)) fallback to use that non-greedy match and keep raising the original JSONDecodeError if no match is found.
🤖 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/tests/services/test_summary_service.py`:
- Around line 390-396: The new test function
test_delete_document_routes_rag_cleanup_through_gateway is missing an explicit
return type; update its signature to include a return annotation "-> None"
(i.e., change the def line for
test_delete_document_routes_rag_cleanup_through_gateway to add "-> None") to
satisfy the project's type-hinting guideline.
---
Outside diff comments:
In `@backend/app/services/knowledge/indexing.py`:
- Around line 337-344: The early skip return object is missing the chunks_data
key, causing an inconsistent response shape compared to the gateway-path return
that includes chunks_data; update the early-return dictionary (the block that
returns status, reason, document_id, knowledge_base_id, indexed_count,
index_name) to also include "chunks_data" (e.g., an empty list) so callers
reading skip_reason/document_id/knowledge_base_id receive a uniform response
shape; reference the variables skip_reason, document_id, knowledge_base_id in
the same return object when making this change.
---
Nitpick comments:
In `@backend/app/api/endpoints/internal/rag.py`:
- Around line 92-102: The mediation context lacks model type info causing
ambiguous resolution in protected_model_resolver._resolve_named_model() which
currently uses model_type=None; add an Optional[str] field current_model_type to
the RetrieveMediationContext Pydantic model (alongside current_model_name and
current_model_namespace), update any callers/constructors that build
RetrieveMediationContext to pass the model type when available, and modify
protected_model_resolver._resolve_named_model() (or its invocation sites) to use
context.current_model_type instead of hardcoding None so resolution is precise
between user/group Kinds and public models.
In `@backend/app/services/knowledge/index_runtime.py`:
- Around line 61-69: The filter on the Kind query uses an unnecessary boolean
comparison; replace the expression Kind.is_active == True with the column
expression Kind.is_active in the kb query (the block that builds kb via
db.query(Kind).filter(...).first()) so the filter reads Kind.is_active instead
of comparing to True; update any similar filters in the same file that use ==
True to follow the same pattern and run tests/linting to confirm Ruff E712 is
resolved.
In `@backend/app/services/knowledge/protected_mediation.py`:
- Around line 406-423: In _extract_json_payload, replace the greedy regex
r"\{.*\}" with a non-greedy pattern (e.g. r"\{.*?\}" with flags=re.DOTALL) so
the fallback extraction captures the first balanced-looking JSON object instead
of spanning multiple brace groups; update the json.loads(match.group(0))
fallback to use that non-greedy match and keep raising the original
JSONDecodeError if no match is found.
In `@backend/app/services/knowledge/protected_model_resolver.py`:
- Around line 76-85: The code path that resolves a current model when
mediation_context contains "current_model_name" omits passing through
"current_model_type", causing ambiguous resolution across user/group/public
scopes; update the call to _resolve_named_model in the mediation_context branch
to pass model_type=mediation_context.get("current_model_type") (or default None)
so the resolver uses the provided type when present (reference:
mediation_context, "current_model_name", "current_model_type", and
_resolve_named_model).
- Around line 43-48: The except Exception in the KB snapshot loader should be
narrowed to SQLAlchemyError to satisfy static analysis and be explicit about
expected DB errors; update the exception handler in the function/method that
loads KB snapshots in protected_model_resolver.py (the block currently using
"except Exception") to "except SQLAlchemyError as e" and add "from
sqlalchemy.exc import SQLAlchemyError" to imports, preserving the logger.debug
call (you can keep exc_info=True or include e in the message) and still return
[] on failure.
In `@backend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.py`:
- Around line 83-113: The gateway assertion is too loose; after calling the
endpoint in test_internal_retrieve_keeps_user_subtask_id_out_of_gateway, inspect
the actual QueryRuntimeSpec passed to the patched LocalRagGateway.query
(mock_query) instead of only using ANY: retrieve the awaited call arguments from
mock_query (e.g., mock_query.await_args or mock_query.assert_awaited_once() then
access mock_query.call_args) and assert that the first positional argument (the
QueryRuntimeSpec) has the expected fields like knowledge_base_ids == [1] and
restricted_mode is False; keep the existing response and
persist_retrieval_result assertions but replace or extend
mock_query.assert_awaited_once_with(ANY, db=ANY) with these explicit checks on
the QueryRuntimeSpec properties.
In `@backend/tests/services/test_summary_service.py`:
- Around line 402-430: Add a new test case that exercises the
non-default/non-organization namespace branch in
KnowledgeService.delete_document: clone the existing block that patches
LocalRagGateway.delete_document_index (AsyncMock) and set
test_knowledge_base.json["spec"]["retrievalConfig"]["retriever_namespace"] to a
non-default value (e.g., "custom-namespace"), commit with
flag_modified/test_db.commit(), call KnowledgeService.delete_document with the
same test_db, test_document.id, and test_user.id, then assert result.success is
True and mock_delete_index was awaited once and that
mock_delete_index.await_args.kwargs["index_owner_user_id"] ==
test_knowledge_base.user_id (also keep assertions for knowledge_base_id,
document_ref, and db).
🪄 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: 8a2bfc5c-9725-48f7-a40b-260d3cac1c09
📒 Files selected for processing (47)
backend/app/api/endpoints/internal/rag.pybackend/app/services/chat/preprocessing/contexts.pybackend/app/services/knowledge/document_read_service.pybackend/app/services/knowledge/index_runtime.pybackend/app/services/knowledge/indexing.pybackend/app/services/knowledge/knowledge_service.pybackend/app/services/knowledge/orchestrator.pybackend/app/services/knowledge/protected_mediation.pybackend/app/services/knowledge/protected_model_resolver.pybackend/app/services/knowledge/retrieval_persistence.pybackend/app/services/rag/README.mdbackend/app/services/rag/gateway.pybackend/app/services/rag/local_data_plane/__init__.pybackend/app/services/rag/local_data_plane/indexing.pybackend/app/services/rag/local_data_plane/retrieval.pybackend/app/services/rag/local_gateway.pybackend/app/services/rag/retrieval_service.pybackend/app/services/rag/runtime_resolver.pybackend/app/services/rag/runtime_specs.pybackend/app/services/rag/splitter/runtime_config.pybackend/tests/api/endpoints/internal/test_rag_retrieve_endpoint.pybackend/tests/services/knowledge/test_document_read_service.pybackend/tests/services/knowledge/test_index_runtime.pybackend/tests/services/knowledge/test_indexing.pybackend/tests/services/knowledge/test_protected_mediation.pybackend/tests/services/knowledge/test_protected_model_resolver.pybackend/tests/services/knowledge/test_retrieval_persistence.pybackend/tests/services/rag/test_local_data_plane_indexing.pybackend/tests/services/rag/test_local_data_plane_retrieval.pybackend/tests/services/rag/test_local_gateway.pybackend/tests/services/rag/test_retrieval_service.pybackend/tests/services/rag/test_runtime_resolver.pybackend/tests/services/rag/test_runtime_specs.pybackend/tests/services/test_summary_service.pybackend/tests/services/test_task_knowledge_base_sync.pybackend/tests/tasks/test_knowledge_tasks.pychat_shell/chat_shell/adapters/wegent.pychat_shell/chat_shell/tools/builtin/knowledge_base.pychat_shell/chat_shell/tools/knowledge_factory.pychat_shell/tests/test_kb_prompt_deduplication.pychat_shell/tests/test_knowledge_base_restricted_backend_mode.pychat_shell/tests/test_knowledge_injection_strategy.pychat_shell/tests/test_knowledge_injection_strategy_clean.pychat_shell/tests/test_wegent_adapter.pydocs/plans/2026-03-24-rag-service-split-plan.mddocs/specs/knowledge/2026-03-30-wegent-knowledge-document-read-design.mddocs/specs/knowledge/2026-03-31-rag-modular-data-plane-design.md
💤 Files with no reviewable changes (2)
- backend/app/services/rag/retrieval_service.py
- docs/specs/knowledge/2026-03-30-wegent-knowledge-document-read-design.md
✅ Files skipped from review due to trivial changes (14)
- backend/app/services/rag/local_data_plane/init.py
- backend/tests/services/knowledge/test_document_read_service.py
- backend/app/services/knowledge/document_read_service.py
- backend/tests/services/rag/test_local_data_plane_indexing.py
- backend/app/services/knowledge/orchestrator.py
- backend/tests/services/rag/test_local_gateway.py
- backend/tests/services/rag/test_local_data_plane_retrieval.py
- chat_shell/tests/test_wegent_adapter.py
- backend/tests/services/rag/test_runtime_resolver.py
- backend/app/services/rag/gateway.py
- chat_shell/tests/test_knowledge_base_restricted_backend_mode.py
- backend/app/services/rag/README.md
- backend/tests/services/knowledge/test_index_runtime.py
- backend/app/services/rag/runtime_specs.py
🚧 Files skipped from review as they are similar to previous changes (14)
- chat_shell/chat_shell/tools/knowledge_factory.py
- backend/tests/tasks/test_knowledge_tasks.py
- backend/tests/services/test_task_knowledge_base_sync.py
- chat_shell/tests/test_knowledge_injection_strategy_clean.py
- backend/tests/services/rag/test_retrieval_service.py
- backend/app/services/rag/local_data_plane/retrieval.py
- backend/app/services/rag/splitter/runtime_config.py
- docs/plans/2026-03-24-rag-service-split-plan.md
- backend/app/services/knowledge/knowledge_service.py
- backend/app/services/rag/local_data_plane/indexing.py
- backend/tests/services/knowledge/test_retrieval_persistence.py
- backend/app/services/rag/runtime_resolver.py
- backend/app/services/rag/local_gateway.py
- backend/tests/services/knowledge/test_protected_model_resolver.py
Comment on lines
+390
to
+396
| def test_delete_document_routes_rag_cleanup_through_gateway( | ||
| self, | ||
| test_db: Session, | ||
| test_user: User, | ||
| test_knowledge_base: Kind, | ||
| test_document: KnowledgeDocument, | ||
| ): |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify missing return type annotations in this reviewed file using Python AST.
python - <<'PY'
import ast
from pathlib import Path
path = Path("backend/tests/services/test_summary_service.py")
tree = ast.parse(path.read_text(encoding="utf-8"))
missing = []
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.returns is None:
missing.append((node.lineno, node.name))
for lineno, name in sorted(missing):
print(f"{path}:{lineno}: function '{name}' has no return type annotation")
PYRepository: wecode-ai/Wegent
Length of output: 1541
Add an explicit return type to the new test function.
Line 390 introduces a new Python function without a return annotation. Add -> None to comply with the coding guideline requiring type hints.
Proposed fix
- def test_delete_document_routes_rag_cleanup_through_gateway(
+ def test_delete_document_routes_rag_cleanup_through_gateway(
self,
test_db: Session,
test_user: User,
test_knowledge_base: Kind,
test_document: KnowledgeDocument,
- ):
+ ) -> None:Per coding guidelines: "Python: Follow PEP 8, Black formatter (line length: 88), isort, type hints required".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_delete_document_routes_rag_cleanup_through_gateway( | |
| self, | |
| test_db: Session, | |
| test_user: User, | |
| test_knowledge_base: Kind, | |
| test_document: KnowledgeDocument, | |
| ): | |
| def test_delete_document_routes_rag_cleanup_through_gateway( | |
| self, | |
| test_db: Session, | |
| test_user: User, | |
| test_knowledge_base: Kind, | |
| test_document: KnowledgeDocument, | |
| ) -> None: |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@backend/tests/services/test_summary_service.py` around lines 390 - 396, The
new test function test_delete_document_routes_rag_cleanup_through_gateway is
missing an explicit return type; update its signature to include a return
annotation "-> None" (i.e., change the def line for
test_delete_document_routes_rag_cleanup_through_gateway to add "-> None") to
satisfy the project's type-hinting guideline.
qdaxb
approved these changes
Apr 3, 2026
This was referenced Apr 5, 2026
This was referenced Apr 19, 2026
cocowh
pushed a commit
to cocowh/Wegent
that referenced
this pull request
May 11, 2026
fix(backend): avoid storing binary_data twice in context upload (wecode-ai#903) See merge request weibo_rd/common/wecode/wegent!1070
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
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
Behavior Changes / Bug Fixes
Documentation
Tests