From 8eee3990f7a41215e8a2aa7123a5c0f0269553f9 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 20 Jan 2026 11:55:41 +0100 Subject: [PATCH] fix: track followup question in cot retriever --- cognee/modules/retrieval/graph_completion_cot_retriever.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cognee/modules/retrieval/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index d42dbda0f..be28cc97c 100644 --- a/cognee/modules/retrieval/graph_completion_cot_retriever.py +++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py @@ -277,6 +277,12 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever): for batched_query_state in query_state_tracker.values() ] ) + + for batched_query, batched_followup_question in zip( + query_batch, followup_question_batch + ): + query_state_tracker[batched_query].followup_question = batched_followup_question + logger.info( f"Chain-of-thought: round {round_idx} - follow-up questions: {followup_question_batch}" )