From cf00909eeb8784585bb1a9f700d0cffcf66c756e Mon Sep 17 00:00:00 2001
From: Neelasha Bhattacharjee <46113280+Neelashab@users.noreply.github.com>
Date: Mon, 26 Jan 2026 17:33:03 +0000
Subject: [PATCH 1/2] feat: add citations to kg block - fixed return types

---
 src/writer/blocks/writeraskkg.py | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/writer/blocks/writeraskkg.py b/src/writer/blocks/writeraskkg.py
index 1981bb1ca..fc4516ee8 100644
--- a/src/writer/blocks/writeraskkg.py
+++ b/src/writer/blocks/writeraskkg.py
@@ -58,7 +58,7 @@ def register(cls, type: str):
                         "name": "Add inline graph citations",
                         "type": "Boolean",
                         "desc": "Shows what specific graph sources were used to answer the question.",
-                        "default": "yes",
+                        "default": "no",
                         "validator": {
                             "type": "boolean",
                         },
@@ -103,7 +103,7 @@ def run(self):
             subqueries = self._get_field(
                 "subqueries", default_field_value="yes") == "yes"
             graph_citations = self._get_field(
-                "graphCitations", default_field_value="yes") == "yes"
+                "graphCitations", default_field_value="no") == "yes"
 
             response = client.graphs.question(
                 graph_ids=graph_ids,
@@ -124,30 +124,34 @@ def run(self):
                     try:
                         delta_answer = chunk.model_extra.get("answer", "")
                         answer_so_far += delta_answer
-                        result_dict["answer"] = answer_so_far
 
                         if graph_citations:
                             delta_sources = chunk.model_extra.get("sources", "")
                             citations_so_far.extend(delta_sources)
+                            result_dict["answer"] = answer_so_far
                             result_dict["citations"] = citations_so_far
                         
-                        self._set_state(state_element, result_dict)
-
                     except json.JSONDecodeError:
                         logging.error(
                             "Could not parse stream chunk from graph.question")
 
             else:
                 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
+    
             self.outcome = "success"
+            
+
 
         except BaseException as e:
             self.outcome = "error"

From 9663cdbb52a91f01b7d57e69f839848f12473d5c Mon Sep 17 00:00:00 2001
From: Neelasha Bhattacharjee <46113280+Neelashab@users.noreply.github.com>
Date: Mon, 26 Jan 2026 19:54:39 +0000
Subject: [PATCH 2/2] feat: add citations to kg - Vlad's updates to improve
 linearity

---
 src/writer/blocks/writeraskkg.py | 67 +++++++++++++++-----------------
 1 file changed, 32 insertions(+), 35 deletions(-)

diff --git a/src/writer/blocks/writeraskkg.py b/src/writer/blocks/writeraskkg.py
index fc4516ee8..4e5cd4593 100644
--- a/src/writer/blocks/writeraskkg.py
+++ b/src/writer/blocks/writeraskkg.py
@@ -115,45 +115,42 @@ def run(self):
                 }
             )
 
-            answer_so_far = ""
-            result_dict = {}
-            citations_so_far = []
-
-            if use_streaming:
-                for chunk in response:
-                    try:
-                        delta_answer = chunk.model_extra.get("answer", "")
-                        answer_so_far += delta_answer
-
-                        if graph_citations:
-                            delta_sources = chunk.model_extra.get("sources", "")
-                            citations_so_far.extend(delta_sources)
-                            result_dict["answer"] = answer_so_far
-                            result_dict["citations"] = citations_so_far
-                        
-                    except json.JSONDecodeError:
-                        logging.error(
-                            "Could not parse stream chunk from graph.question")
+            self.result = self._parse_response(response, state_element, use_streaming, graph_citations)
+            if state_element:
+                self._set_state(state_element, self.result)
+            self.outcome = "success"
 
-            else:
-                answer_so_far = response.answer
+        except BaseException as e:
+            self.outcome = "error"
+            raise e
             
-                if graph_citations:
-                    result_dict["answer"] = answer_so_far
-                    result_dict["citations"] = response.sources or []
 
+
+    def _parse_response(self, response, state_element, use_streaming: bool, graph_citations: bool):
+        if not use_streaming:
             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
-    
-            self.outcome = "success"
-            
+                return {"answer": response.answer, "citations": response.sources or []}
+            return response.answer
 
+        answer = ""
+        citations = []
 
-        except BaseException as e:
-            self.outcome = "error"
-            raise e
+        for chunk in response:
+            try:
+                delta_answer = chunk.model_extra.get("answer", "")
+                answer += delta_answer
+
+                if graph_citations:
+                    delta_sources = chunk.model_extra.get("sources", "")
+                    citations.extend(delta_sources)
+                    self._set_state(state_element, {"answer": answer, "citations": citations})
+                else:
+                    self._set_state(state_element, answer)
 
+            except json.JSONDecodeError:
+                logging.error("Could not parse stream chunk from graph.question")
+        
+        if graph_citations:
+            return {"answer": answer, "citations": citations}
+        return answer
+        
\ No newline at end of file