feat: AB-938 - add citations to ask graph block#1272
Conversation
📝 WalkthroughWalkthroughThe WriterAskGraphQuestion block's Changes
Sequence Diagram(s)(Skipped — changes are internal refactor + default change without multi-component sequential interaction requiring visualization.) Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/writer/blocks/writeraskkg.py`:
- Around line 139-150: The code calls self._set_state(state_element, ...)
unconditionally which will fail when state_element is None; update the branches
around graph_citations to only call self._set_state if state_element is truthy
(e.g., if state_element: self._set_state(...)), otherwise skip the call but
still set self.result appropriately; ensure this check is applied in both the
graph_citations and non-graph_citations branches (references: _set_state,
state_element, result_dict, answer_so_far, graph_citations, self.result).
src/writer/blocks/writeraskkg.py
Outdated
| answer_so_far = response.answer | ||
| result_dict["answer"] = answer_so_far | ||
|
|
||
|
|
||
| if graph_citations: | ||
| citations_so_far = response.sources or [] | ||
| result_dict["citations"] = citations_so_far | ||
|
|
||
| self._set_state(state_element, result_dict) | ||
| self.result = answer_so_far | ||
| result_dict["answer"] = answer_so_far | ||
| result_dict["citations"] = response.sources or [] | ||
|
|
||
| if graph_citations: | ||
| self._set_state(state_element, result_dict) | ||
| self.result = result_dict | ||
| else: | ||
| self._set_state(state_element, answer_so_far) | ||
| self.result = answer_so_far |
There was a problem hiding this comment.
Guard _set_state when stateElement is absent.
Non-streaming allows stateElement to be omitted, but Line 145-150 now calls _set_state unconditionally, which will likely error when state_element is None. This breaks previously valid configurations.
✅ Suggested fix
- if graph_citations:
- self._set_state(state_element, result_dict)
- self.result = result_dict
- else:
- self._set_state(state_element, answer_so_far)
- self.result = answer_so_far
+ if graph_citations:
+ if state_element:
+ self._set_state(state_element, result_dict)
+ self.result = result_dict
+ else:
+ if state_element:
+ self._set_state(state_element, answer_so_far)
+ self.result = answer_so_far📝 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.
| answer_so_far = response.answer | |
| result_dict["answer"] = answer_so_far | |
| if graph_citations: | |
| citations_so_far = response.sources or [] | |
| result_dict["citations"] = citations_so_far | |
| self._set_state(state_element, result_dict) | |
| self.result = answer_so_far | |
| result_dict["answer"] = answer_so_far | |
| result_dict["citations"] = response.sources or [] | |
| if graph_citations: | |
| self._set_state(state_element, result_dict) | |
| self.result = result_dict | |
| else: | |
| self._set_state(state_element, answer_so_far) | |
| self.result = answer_so_far | |
| answer_so_far = response.answer | |
| if graph_citations: | |
| result_dict["answer"] = answer_so_far | |
| result_dict["citations"] = response.sources or [] | |
| if graph_citations: | |
| if state_element: | |
| self._set_state(state_element, result_dict) | |
| self.result = result_dict | |
| else: | |
| if state_element: | |
| self._set_state(state_element, answer_so_far) | |
| self.result = answer_so_far |
🤖 Prompt for AI Agents
In `@src/writer/blocks/writeraskkg.py` around lines 139 - 150, The code calls
self._set_state(state_element, ...) unconditionally which will fail when
state_element is None; update the branches around graph_citations to only call
self._set_state if state_element is truthy (e.g., if state_element:
self._set_state(...)), otherwise skip the call but still set self.result
appropriately; ensure this check is applied in both the graph_citations and
non-graph_citations branches (references: _set_state, state_element,
result_dict, answer_so_far, graph_citations, self.result).
There was a problem hiding this comment.
@UladzislauK-Writer @madeindjs I am a bit confused, I thought the definition of state element was not dependent on useStreaming, and was just used to update the agent state. I believe before this feature implementation, state element was used to set state regardless of whether streaming was enabled. Any insights into this suggestion would be awesome
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
Yes, state_element was used regardless of use_streaming
This is talking about the fact that you can have use_streaming=False and state_element=None. In this case execution will reach these lines and calling self._set_state(None, ...) will (most likely) throw an error
| self._set_state(state_element, result_dict) | ||
|
|
There was a problem hiding this comment.
Removing this breaks the "Use streaming" functionality. The whole purpose of it is to update the state variable as we iterate over the response which is arriving over the network in multiple parts. With it the user can get more "AI-chatbot-like" experience when you are not waiting for the whole answer to appear at once but instead see it being generated
There was a problem hiding this comment.
@UladzislauK-Writer I see, so when streaming is enabled, I should include the set state inside the for loop?
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/writer/blocks/writeraskkg.py`:
- Around line 150-151: The except block catching json.JSONDecodeError in
writeraskkg.py is incorrect because model_extra.get() doesn't parse JSON and
that exception will never be raised; change the handler to catch a broader
exception (e.g., Exception as e) around the section that accesses
model_extra.get(...) / graph.question and replace logging.error(...) with
logging.exception(...) so the traceback is captured and the error message
includes context (e.g., "Could not parse stream chunk from graph.question"). If
you determine the try/except is unnecessary (dead code), remove it instead.
Ensure you update the exception block associated with the
model_extra.get()/graph.question access in the same function where this diff
appears.
🧹 Nitpick comments (3)
src/writer/blocks/writeraskkg.py (3)
118-121: Minor redundancy in streaming mode.For streaming,
_set_stateis called on every chunk inside_parse_response(lines 146, 148), then again here after the method returns. The final call is redundant for streaming since the last chunk already set the final state.Consider gating this to avoid the duplicate call:
self.result = self._parse_response(response, state_element, use_streaming, graph_citations) -if state_element: +if state_element and not use_streaming: self._set_state(state_element, self.result)
123-128: Use bareraiseto preserve full traceback.Using
raise ecan lose traceback context in some Python versions. A bareraisere-raises the current exception with its original traceback intact.Suggested fix
except BaseException as e: self.outcome = "error" - raise e - - + raise
143-146: Use list as default for sources to ensure type safety.If
sourcesis missing frommodel_extra, the default""(empty string) won't cause issues with.extend(), but using[]is more semantically correct and defensive against unexpected types.- delta_sources = chunk.model_extra.get("sources", "") + delta_sources = chunk.model_extra.get("sources", [])
| except json.JSONDecodeError: | ||
| logging.error("Could not parse stream chunk from graph.question") |
There was a problem hiding this comment.
Exception handler appears mismatched and should use logging.exception.
Two issues here:
-
Wrong exception type: The try block doesn't perform any JSON parsing—
model_extra.get()returns dict values without parsing.json.JSONDecodeErrorwill never be raised. Consider catching a broader exception type or removing this handler if it's dead code. -
Use
logging.exception: When logging from an exception handler,logging.exceptionautomatically includes the traceback.
Suggested fix
- except json.JSONDecodeError:
- logging.error("Could not parse stream chunk from graph.question")
+ except Exception:
+ logging.exception("Could not process stream chunk from graph.question")📝 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.
| except json.JSONDecodeError: | |
| logging.error("Could not parse stream chunk from graph.question") | |
| except Exception: | |
| logging.exception("Could not process stream chunk from graph.question") |
🧰 Tools
🪛 Ruff (0.14.14)
151-151: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
🤖 Prompt for AI Agents
In `@src/writer/blocks/writeraskkg.py` around lines 150 - 151, The except block
catching json.JSONDecodeError in writeraskkg.py is incorrect because
model_extra.get() doesn't parse JSON and that exception will never be raised;
change the handler to catch a broader exception (e.g., Exception as e) around
the section that accesses model_extra.get(...) / graph.question and replace
logging.error(...) with logging.exception(...) so the traceback is captured and
the error message includes context (e.g., "Could not parse stream chunk from
graph.question"). If you determine the try/except is unnecessary (dead code),
remove it instead. Ensure you update the exception block associated with the
model_extra.get()/graph.question access in the same function where this diff
appears.
This reverts commit d1a4b4e.
3 participants
Changes to response format when citations are not enabled to make sure existing agents don't break
Summary by CodeRabbit
Changes
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.