feat: filter RAG sources by LLM citations to reduce false positives - #246
Conversation
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (12)
📝 WalkthroughWalkthroughThis PR implements end-of-response source citation extraction and filtering for the RAG pipeline, refactors context formatting to return included document indices, adds streaming-safe source filtering utilities, integrates them into streaming and non-streaming routes, and adds unit/integration tests plus mock streaming support. Changes
Sequence DiagramsequenceDiagram
participant Client
participant API as API Router
participant Pipeline
participant Utils
participant LLM
participant Filter as Source Filter
Client->>API: Chat completion request
API->>Pipeline: prepare_for_chat_completion(...)
Pipeline->>Utils: format_context(docs, number_sources=True)
Utils-->>Pipeline: (formatted_context, included_indices)
Pipeline-->>API: (payload, filtered_docs)
API->>LLM: Send payload + prompt
alt Streaming
LLM-->>API: SSE chunks (includes final [Sources: ...])
API->>Filter: stream_with_source_filtering(llm_stream, all_sources_json)
Filter->>Filter: buffer chunks, detect sources tag
Filter->>Utils: extract_and_strip_sources_block(buffered_text)
Utils-->>Filter: (clean_text, citations)
Filter->>Utils: filter_sources_by_citations(all_sources, citations)
Utils-->>Filter: filtered_sources
Filter-->>API: cleaned SSE chunks + final extra with filtered_sources
API-->>Client: Streamed cleaned response
else Non-Streaming
LLM-->>API: Full response with [Sources: ...]
API->>Utils: extract_and_strip_sources_block(content)
Utils-->>API: (clean_content, citations)
API->>Utils: filter_sources_by_citations(all_sources, citations)
Utils-->>API: filtered_sources
API-->>Client: Response with cleaned content and extra filtered_sources
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
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 |
51816a3 to
d94f0ae
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (4)
openrag/components/pipeline.py (1)
179-191: Source citation instruction duplicated inline for completions path.The source instruction on Line 187 duplicates what the prompt templates already specify for the chat path. This is acceptable since completions don't use system prompt templates, but consider extracting this instruction string to a shared constant in
components/promptsto keep the expected format in sync.openrag/routers/openai.py (1)
191-192:all_sources_jsonis computed but effectively unused.
all_sources_jsonis passed tostream_with_source_filteringbut the parameter is annotatednoqa: ARG001("no longer sent to clients"). Consider removing both the variable here and the unused parameter fromstream_with_source_filteringto avoid confusion.tests/api_tests/conftest.py (1)
160-170:wait_for_indexingsilently falls back to a 5-second sleep and returnsNone.When neither
task_idnortask_status_urlis present, the function silently sleeps and returnsNoneinstead of a status dict. This may mask unexpected API response formats. Consider logging a warning or raising, so test authors are alerted when the response structure changes unexpectedly.Also, the existing
indexed_folder_partitionfixture (lines 188–206) duplicates the same poll-and-wait logic manually — it could be simplified by callingwait_for_indexinginstead.Proposed improvement for the silent fallback
else: + import warnings + warnings.warn("No task_id or task_status_url in response; falling back to fixed sleep", stacklevel=2) time.sleep(5)tests/api_tests/test_openai_compat.py (1)
125-149: SSE parsing is duplicated across streaming tests — consider a small helper.The pattern of streaming, filtering
data:lines, checking[DONE], and parsing JSON chunks is repeated verbatim intest_streaming_content_clean,test_streaming_has_role_delta, andtest_streaming_has_finish_reason. A small helper (e.g.iter_sse_chunks(response)) that yields parsed chunk dicts would reduce the boilerplate and make the tests easier to maintain.Example helper
def iter_sse_chunks(response): """Yield parsed SSE chunk dicts from a streaming httpx response.""" for line in response.iter_lines(): if not line.startswith("data:"): continue if "[DONE]" in line: return yield json.loads(line[len("data: "):])
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
LGTM. This PR highlights a broader issue: we need better filtering before sending chunks to the LLM. Unused chunks consume context space, reducing available generation tokens. Filtering them would free up room for more relevant chunks and improve their ranking.
| for n_tokens, doc in docs_with_tokens: | ||
|
|
||
| for i, doc in enumerate(docs): | ||
| prefix = f"[Source {len(reduced_docs) + 1}]\n" if number_sources else "" |
There was a problem hiding this comment.
This is a small thing but instead of len(reduced_docs) + 1 the counter variable i could be used directly:
- prefix = f"[Source {len(reduced_docs) + 1}]\n" if number_sources else ""
+ prefix = f"[Source {i + 1}]\n" if number_sources else ""| last_chunk_template = data | ||
| content_buffer += content | ||
|
|
||
| if len(content_buffer) > buffer_size: |
There was a problem hiding this comment.
My main concern here is that we are not sending chunks anymore but an accumulated value to_emit = content_buffer[:-buffer_size] in the place the current streamed chunk content data["choices"][0]["delta"]["content"] and metadata of that chunks doesn't go along witht the content which is problematic.
There was a problem hiding this comment.
I've proposed another version in this PR #257 that leaves the original chunks intact.
There was a problem hiding this comment.
Ok, I considered it was ok to have potentially incorrect metadata chunk (or at least, shifted), but your fix avoids that, so it's nice 👍
And good use of deque!
660fadd to
b09e9fa
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/api_tests/api_run/mock_vllm.py (1)
1-16:⚠️ Potential issue | 🟡 MinorFix Ruff formatting to unblock CI.
The pipeline reports that this file would be reformatted by
ruff format. Runruff format tests/api_tests/api_run/mock_vllm.pyto fix.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/api_run/mock_vllm.py` around lines 1 - 16, Run ruff format on this module (or apply equivalent changes) to match the project's formatting rules: reflow the module docstring, sort and group imports, ensure single blank lines between import blocks and top-level code, fix spacing around commas and parentheses in the import lines (hashlib, json, time, uuid, typing import), and ensure a trailing newline; the primary identifiers to check are the module-level imports and the FastAPI app instantiation (app = FastAPI()) so CI no longer reports the file as reformattable by ruff.
🧹 Nitpick comments (3)
tests/api_tests/api_run/mock_vllm.py (1)
128-154: Mock always cites[Sources: 1]— sufficient for testing but worth noting.The mock appends
\n[Sources: 1]whenever the system prompt contains[Source 1]. This exercises the extraction/filtering path but won't test multi-source citation filtering or the "no match → fallback to all" path. Consider parameterizing or adding a second mock variant if broader coverage is needed later.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/api_run/mock_vllm.py` around lines 128 - 154, The mock in this file unconditionally appends a single-source citation whenever the system prompt contains “[Source 1]” (see variables system_msg and has_numbered_sources and the final response mutation), which limits test coverage for multi-source and no-match fallback paths; update the mock to allow configurable citation behavior (e.g., accept an argument or environment flag like cited_sources or citation_mode) and implement logic in the mock function to: 1) append multiple citations when cited_sources > 1, 2) return no citations when citation_mode == "none", and 3) simulate the fallback-to-all behavior when citation_mode == "fallback"; use the same response construction flow (content handling, content_lower checks) and only alter the final citation-appending block so tests can toggle between single-source, multi-source, and no-match scenarios.tests/api_tests/conftest.py (1)
177-213:indexed_folder_partitiondoesn't use the newwait_for_task/wait_for_indexinghelpers.The fixture still has its own inline polling loop (lines 193-211) with different timeout logic (30 iterations × 2s = 60s) vs.
TASK_TIMEOUT(180s), and checks for both upper/lowercase state strings. Consider refactoring to usewait_for_indexingfor consistency and reduced duplication.Proposed refactor
data = response.json() - # Wait for each file to be indexed - if "task_status_url" in data: - task_url = data["task_status_url"] - task_path = "/" + "/".join(task_url.split("/")[3:]) - elif "task_id" in data: - task_path = f"/indexer/task/{data['task_id']}" - else: - time.sleep(3) - continue - - for _ in range(30): - task_response = api_client.get(task_path) - task_data = task_response.json() - state = task_data.get("task_state", "") - if state in ["SUCCESS", "COMPLETED", "success", "completed"]: - break - elif state in ["FAILED", "failed", "FAILURE", "failure"]: - pytest.skip(f"Indexing failed for {filename}: {task_data}") - time.sleep(2) + try: + wait_for_indexing(api_client, data) + except AssertionError: + pytest.skip(f"Indexing failed for {filename}")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/conftest.py` around lines 177 - 213, The indexed_folder_partition fixture contains an inline polling loop to wait for indexing that duplicates logic and timeout values; replace the per-file polling (the block that inspects "task_status_url" / "task_id" and loops up to 30×2s checking task_state) with the existing helper wait_for_indexing (or wait_for_task) to centralize timeout and status handling. Specifically, in indexed_folder_partition use the response to derive task_path as currently done, then call wait_for_indexing(api_client, task_path) (or wait_for_task) instead of the manual for-loop and state casing checks so the fixture uses TASK_TIMEOUT and consistent success/failure logic. Ensure you preserve the existing behavior of skipping the test on failure by letting the helper raise or return failure in the same way the callers expect.openrag/components/pipeline.py (1)
183-188: Inline source instruction for completions — consider extracting to a shared constant.The source citation instruction on line 187 duplicates the intent of lines 10-12 in
spoken_style_answer_tmpl.txt. If the citation format changes (e.g., different tag syntax), this inline string would need a separate update.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/pipeline.py` around lines 183 - 188, The inline citation instruction string appended into the local variable prompt (the multiline f-string that includes "At the very end of your response, on a new line, list which source numbers you used: [Sources: 1, 3]") should be extracted into a shared constant (e.g., SOURCE_CITATION_INSTRUCTION) and referenced from the prompt construction; update the existing spoken_style_answer_tmpl.txt usage to import or load that same constant so both the pipeline.py prompt builder and the spoken_style template share a single source of truth, and replace the duplicated literal in pipeline.py (where prompt is created) with that constant.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@openrag/components/utils.py`:
- Around line 165-195: The async loop over llm_stream can exit without hitting
the "data: [DONE]" branch and drop buffered content in chunk_buffer; wrap the
iteration in try/finally (or add a sentinel after the loop) to detect premature
termination and flush chunk_buffer the same way the "[DONE]" branch does: build
buffered_text from chunk_buffer, call
extract_and_strip_sources_block(buffered_text) to get clean_text and citations,
trim each chunk's original content to the surviving slice, set chunk["extra"] =
"{}", yield those chunks, and if last_chunk_template exists produce the
finish_chunk with filtered = filter_sources_by_citations(sources, citations) and
include last_finish_reason and the extra sources, then yield "data: [DONE]\n\n".
In `@tests/api_tests/conftest.py`:
- Around line 165-175: The wait_for_indexing function currently falls through
and implicitly returns None when response_data lacks "task_id" and
"task_status_url", which hides failures; update wait_for_indexing to log a
warning (use the module logger or testing logger) indicating missing task info
(include response_data for context) and make the return explicit (return None or
a clearly documented sentinel) so callers can detect the condition; reference
wait_for_indexing and wait_for_task when making the change so callers still use
wait_for_task when a task_id is available.
---
Outside diff comments:
In `@tests/api_tests/api_run/mock_vllm.py`:
- Around line 1-16: Run ruff format on this module (or apply equivalent changes)
to match the project's formatting rules: reflow the module docstring, sort and
group imports, ensure single blank lines between import blocks and top-level
code, fix spacing around commas and parentheses in the import lines (hashlib,
json, time, uuid, typing import), and ensure a trailing newline; the primary
identifiers to check are the module-level imports and the FastAPI app
instantiation (app = FastAPI()) so CI no longer reports the file as
reformattable by ruff.
---
Nitpick comments:
In `@openrag/components/pipeline.py`:
- Around line 183-188: The inline citation instruction string appended into the
local variable prompt (the multiline f-string that includes "At the very end of
your response, on a new line, list which source numbers you used: [Sources: 1,
3]") should be extracted into a shared constant (e.g.,
SOURCE_CITATION_INSTRUCTION) and referenced from the prompt construction; update
the existing spoken_style_answer_tmpl.txt usage to import or load that same
constant so both the pipeline.py prompt builder and the spoken_style template
share a single source of truth, and replace the duplicated literal in
pipeline.py (where prompt is created) with that constant.
In `@tests/api_tests/api_run/mock_vllm.py`:
- Around line 128-154: The mock in this file unconditionally appends a
single-source citation whenever the system prompt contains “[Source 1]” (see
variables system_msg and has_numbered_sources and the final response mutation),
which limits test coverage for multi-source and no-match fallback paths; update
the mock to allow configurable citation behavior (e.g., accept an argument or
environment flag like cited_sources or citation_mode) and implement logic in the
mock function to: 1) append multiple citations when cited_sources > 1, 2) return
no citations when citation_mode == "none", and 3) simulate the fallback-to-all
behavior when citation_mode == "fallback"; use the same response construction
flow (content handling, content_lower checks) and only alter the final
citation-appending block so tests can toggle between single-source,
multi-source, and no-match scenarios.
In `@tests/api_tests/conftest.py`:
- Around line 177-213: The indexed_folder_partition fixture contains an inline
polling loop to wait for indexing that duplicates logic and timeout values;
replace the per-file polling (the block that inspects "task_status_url" /
"task_id" and loops up to 30×2s checking task_state) with the existing helper
wait_for_indexing (or wait_for_task) to centralize timeout and status handling.
Specifically, in indexed_folder_partition use the response to derive task_path
as currently done, then call wait_for_indexing(api_client, task_path) (or
wait_for_task) instead of the manual for-loop and state casing checks so the
fixture uses TASK_TIMEOUT and consistent success/failure logic. Ensure you
preserve the existing behavior of skipping the test on failure by letting the
helper raise or return failure in the same way the callers expect.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
CLAUDE.mdopenrag/app_front.pyopenrag/components/pipeline.pyopenrag/components/test_source_filtering.pyopenrag/components/utils.pyopenrag/routers/openai.pyprompts/example1/spoken_style_answer_tmpl.txtprompts/example1/sys_prompt_tmpl.txttests/api_tests/api_run/mock_vllm.pytests/api_tests/conftest.pytests/api_tests/test_indexer.pytests/api_tests/test_openai_compat.py
🚧 Files skipped from review as they are similar to previous changes (5)
- prompts/example1/sys_prompt_tmpl.txt
- openrag/components/test_source_filtering.py
- openrag/app_front.py
- tests/api_tests/test_indexer.py
- tests/api_tests/test_openai_compat.py
| async for line in llm_stream: | ||
| if not line.startswith("data:"): | ||
| continue | ||
|
|
||
| if line.strip() == "data: [DONE]": | ||
| buffered_text = "".join( | ||
| (c.get("choices", [{}])[0].get("delta", {}).get("content", "") or "") for c in chunk_buffer | ||
| ) | ||
| clean_text, citations = extract_and_strip_sources_block(buffered_text) | ||
|
|
||
| remaining = clean_text | ||
| for chunk in chunk_buffer: | ||
| original_content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "") or "" | ||
| if not remaining: | ||
| break | ||
| surviving = remaining[: len(original_content)] | ||
| remaining = remaining[len(original_content) :] | ||
| if surviving != original_content: | ||
| chunk["choices"][0]["delta"]["content"] = surviving | ||
| chunk["extra"] = "{}" | ||
| yield f"data: {json.dumps(chunk)}\n\n" | ||
|
|
||
| if last_chunk_template: | ||
| filtered = filter_sources_by_citations(sources, citations) | ||
| finish_chunk = copy.deepcopy(last_chunk_template) | ||
| finish_chunk["choices"][0]["delta"] = {} | ||
| finish_chunk["choices"][0]["finish_reason"] = last_finish_reason or "stop" | ||
| finish_chunk["extra"] = json.dumps({"sources": filtered}) | ||
| yield f"data: {json.dumps(finish_chunk)}\n\n" | ||
|
|
||
| yield "data: [DONE]\n\n" |
There was a problem hiding this comment.
Buffered content is silently lost if the stream ends without data: [DONE].
If the upstream LLM stream is interrupted (connection drop, timeout, exception), the async for loop exits without hitting the [DONE] branch. Up to buffer_size (100) characters of legitimate content in chunk_buffer will never be yielded to the client. The client sees a truncated response with no indication that content was withheld.
Consider adding a try/finally or sentinel to flush the buffer on unexpected stream termination:
Sketch of a safer flush
+ try:
async for line in llm_stream:
... # existing logic
+ finally:
+ # Flush remaining buffer on unexpected exit (no [DONE] received)
+ if chunk_buffer:
+ for chunk in chunk_buffer:
+ chunk["extra"] = "{}"
+ yield f"data: {json.dumps(chunk)}\n\n"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@openrag/components/utils.py` around lines 165 - 195, The async loop over
llm_stream can exit without hitting the "data: [DONE]" branch and drop buffered
content in chunk_buffer; wrap the iteration in try/finally (or add a sentinel
after the loop) to detect premature termination and flush chunk_buffer the same
way the "[DONE]" branch does: build buffered_text from chunk_buffer, call
extract_and_strip_sources_block(buffered_text) to get clean_text and citations,
trim each chunk's original content to the surviving slice, set chunk["extra"] =
"{}", yield those chunks, and if last_chunk_template exists produce the
finish_chunk with filtered = filter_sources_by_citations(sources, citations) and
include last_finish_reason and the extra sources, then yield "data: [DONE]\n\n".
| def wait_for_indexing(api_client, response_data: dict, timeout: int = TASK_TIMEOUT): | ||
| """Wait for file indexing task to complete, extracting task_id from response.""" | ||
| if "task_id" in response_data: | ||
| return wait_for_task(api_client, response_data["task_id"], timeout) | ||
| elif "task_status_url" in response_data: | ||
| # Extract task_id from URL like http://host/indexer/task/{task_id} | ||
| task_id = response_data["task_status_url"].split("/")[-1] | ||
| return wait_for_task(api_client, task_id, timeout) | ||
| else: | ||
| time.sleep(5) | ||
|
|
There was a problem hiding this comment.
wait_for_indexing silently returns None when response lacks task info.
When neither task_id nor task_status_url is present, the function sleeps 5 seconds and implicitly returns None. Callers have no way to distinguish "completed successfully" from "couldn't track the task." Consider logging a warning or returning a sentinel.
Proposed fix
else:
+ # No task tracking info available — best-effort wait
time.sleep(5)
+ return NoneThe explicit return None doesn't change behavior but documents the intent. A logger.warning(...) here would help diagnose test failures.
📝 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.
| def wait_for_indexing(api_client, response_data: dict, timeout: int = TASK_TIMEOUT): | |
| """Wait for file indexing task to complete, extracting task_id from response.""" | |
| if "task_id" in response_data: | |
| return wait_for_task(api_client, response_data["task_id"], timeout) | |
| elif "task_status_url" in response_data: | |
| # Extract task_id from URL like http://host/indexer/task/{task_id} | |
| task_id = response_data["task_status_url"].split("/")[-1] | |
| return wait_for_task(api_client, task_id, timeout) | |
| else: | |
| time.sleep(5) | |
| def wait_for_indexing(api_client, response_data: dict, timeout: int = TASK_TIMEOUT): | |
| """Wait for file indexing task to complete, extracting task_id from response.""" | |
| if "task_id" in response_data: | |
| return wait_for_task(api_client, response_data["task_id"], timeout) | |
| elif "task_status_url" in response_data: | |
| # Extract task_id from URL like http://host/indexer/task/{task_id} | |
| task_id = response_data["task_status_url"].split("/")[-1] | |
| return wait_for_task(api_client, task_id, timeout) | |
| else: | |
| # No task tracking info available — best-effort wait | |
| time.sleep(5) | |
| return None |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/api_tests/conftest.py` around lines 165 - 175, The wait_for_indexing
function currently falls through and implicitly returns None when response_data
lacks "task_id" and "task_status_url", which hides failures; update
wait_for_indexing to log a warning (use the module logger or testing logger)
indicating missing task info (include response_data for context) and make the
return explicit (return None or a clearly documented sentinel) so callers can
detect the condition; reference wait_for_indexing and wait_for_task when making
the change so callers still use wait_for_task when a task_id is available.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
tests/api_tests/conftest.py (1)
165-175:⚠️ Potential issue | 🟡 MinorMake the “no task info” path explicit and observable.
This still falls through and returns
Nonesilently; the prior review already flagged adding an explicit return + warning. Please make the return explicit and log the condition so failures are diagnosable.
As per coding guidelines, useget_logger()fromopenrag.utils.loggerwith.bind()for context.🛈 Proposed fix
+from openrag.utils.logger import get_logger + +logger = get_logger(__name__) + def wait_for_indexing(api_client, response_data: dict, timeout: int = TASK_TIMEOUT): """Wait for file indexing task to complete, extracting task_id from response.""" if "task_id" in response_data: return wait_for_task(api_client, response_data["task_id"], timeout) elif "task_status_url" in response_data: # Extract task_id from URL like http://host/indexer/task/{task_id} task_id = response_data["task_status_url"].split("/")[-1] return wait_for_task(api_client, task_id, timeout) else: - time.sleep(5) + logger.bind(response_data=response_data).warning("No task tracking info in indexing response") + time.sleep(5) + return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api_tests/conftest.py` around lines 165 - 175, The "no task info" branch in wait_for_indexing currently falls through and returns None silently; update wait_for_indexing to create a logger via get_logger().bind(context={'function':'wait_for_indexing'}) and when neither "task_id" nor "task_status_url" are present, log a warning with the response_data (e.g., logger.warning("No task info in response", response=response_data)) and then explicitly return None so the absence of task info is observable and diagnosable.openrag/components/utils.py (1)
165-195:⚠️ Potential issue | 🟡 MinorStream may drop buffered content if
[DONE]never arrives.
If the upstream stream terminates early, the buffered tail is never flushed to the client. Consider atry/finallyflush path.Possible fix sketch
async def stream_with_source_filtering(...): ... - async for line in llm_stream: - ... + try: + async for line in llm_stream: + ... + finally: + if chunk_buffer: + buffered_text = "".join( + (c.get("choices", [{}])[0].get("delta", {}).get("content", "") or "") + for c in chunk_buffer + ) + clean_text, citations = extract_and_strip_sources_block(buffered_text) + remaining = clean_text + for chunk in chunk_buffer: + original_content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "") or "" + if not remaining: + break + surviving = remaining[: len(original_content)] + remaining = remaining[len(original_content) :] + if surviving != original_content: + chunk["choices"][0]["delta"]["content"] = surviving + chunk["extra"] = "{}" + yield f"data: {json.dumps(chunk)}\n\n" + if last_chunk_template: + filtered = filter_sources_by_citations(sources, citations) + finish_chunk = copy.deepcopy(last_chunk_template) + finish_chunk["choices"][0]["delta"] = {} + finish_chunk["choices"][0]["finish_reason"] = last_finish_reason or "stop" + finish_chunk["extra"] = json.dumps({"sources": filtered}) + yield f"data: {json.dumps(finish_chunk)}\n\n" + yield "data: [DONE]\n\n"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@openrag/components/utils.py` around lines 165 - 195, The loop over llm_stream currently only flushes buffered chunk_buffer when a "data: [DONE]" arrives, so if the upstream stream closes early you must flush the tail in a finally/after-loop path; implement a post-loop flush that mirrors the "data: [DONE]" handling: build buffered_text from chunk_buffer, call extract_and_strip_sources_block(buffered_text) to get clean_text and citations, reconcile surviving content with each chunk in chunk_buffer (updating chunk["choices"][0]["delta"]["content"] and chunk["extra"]) and yield them, then if last_chunk_template create finish_chunk (set choices[0].delta = {}, choices[0].finish_reason = last_finish_reason or "stop", extra = json.dumps({"sources": filter_sources_by_citations(sources, citations)})) and yield it, and finally yield "data: [DONE]\n\n"; ensure this runs even if the async iterator ends without a DONE by using try/finally or handling StopAsyncIteration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/api_tests/conftest.py`:
- Around line 130-163: The wait_for_task helper currently only treats
"COMPLETED" as success and can miss "SUCCESS" or "success" states; update the
wait_for_task function to normalize the returned task_state (e.g., state =
status.get("task_state", "").upper()) and consider both "COMPLETED" and
"SUCCESS" as terminal success values before returning the status, while keeping
the existing "FAILED" handling and timeout behavior.
---
Duplicate comments:
In `@openrag/components/utils.py`:
- Around line 165-195: The loop over llm_stream currently only flushes buffered
chunk_buffer when a "data: [DONE]" arrives, so if the upstream stream closes
early you must flush the tail in a finally/after-loop path; implement a
post-loop flush that mirrors the "data: [DONE]" handling: build buffered_text
from chunk_buffer, call extract_and_strip_sources_block(buffered_text) to get
clean_text and citations, reconcile surviving content with each chunk in
chunk_buffer (updating chunk["choices"][0]["delta"]["content"] and
chunk["extra"]) and yield them, then if last_chunk_template create finish_chunk
(set choices[0].delta = {}, choices[0].finish_reason = last_finish_reason or
"stop", extra = json.dumps({"sources": filter_sources_by_citations(sources,
citations)})) and yield it, and finally yield "data: [DONE]\n\n"; ensure this
runs even if the async iterator ends without a DONE by using try/finally or
handling StopAsyncIteration.
In `@tests/api_tests/conftest.py`:
- Around line 165-175: The "no task info" branch in wait_for_indexing currently
falls through and returns None silently; update wait_for_indexing to create a
logger via get_logger().bind(context={'function':'wait_for_indexing'}) and when
neither "task_id" nor "task_status_url" are present, log a warning with the
response_data (e.g., logger.warning("No task info in response",
response=response_data)) and then explicitly return None so the absence of task
info is observable and diagnosable.
ℹ️ Review info
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
CLAUDE.mdopenrag/app_front.pyopenrag/components/pipeline.pyopenrag/components/test_source_filtering.pyopenrag/components/utils.pyopenrag/routers/openai.pyprompts/example1/spoken_style_answer_tmpl.txtprompts/example1/sys_prompt_tmpl.txttests/api_tests/api_run/mock_vllm.pytests/api_tests/conftest.pytests/api_tests/test_indexer.pytests/api_tests/test_openai_compat.py
🚧 Files skipped from review as they are similar to previous changes (5)
- prompts/example1/spoken_style_answer_tmpl.txt
- openrag/components/test_source_filtering.py
- openrag/components/pipeline.py
- tests/api_tests/api_run/mock_vllm.py
- tests/api_tests/test_openai_compat.py
Number sources in context ([Source 1], [Source 2], ...) and instruct the LLM to append a [Sources: N, ...] tag listing which it actually used. The server strips this tag from the response and filters the source metadata accordingly, for both streaming and non-streaming modes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
b09e9fa to
a89d1cc
Compare
2 participants
Number sources in context ([Source 1], [Source 2], ...) and instruct the LLM to append a [Sources: N, ...] tag listing which it actually used. The server strips this tag from the response and filters the source metadata accordingly, for both streaming and non-streaming modes.
See before/after example (do not mind the old info in the example)
Before
After
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Documentation