From 0f9aaaaa48cd92aaae315881a480478aebe3cbdb Mon Sep 17 00:00:00 2001
From: Logan <logan.markewich@live.com>
Date: Wed, 6 Mar 2024 13:08:56 -0600
Subject: [PATCH] inheirt callback manager (#11662)

---
 .../core/indices/struct_store/sql_query.py    | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/llama-index-core/llama_index/core/indices/struct_store/sql_query.py b/llama-index-core/llama_index/core/indices/struct_store/sql_query.py
index e48e07b72a162..271911b718867 100644
--- a/llama-index-core/llama_index/core/indices/struct_store/sql_query.py
+++ b/llama-index-core/llama_index/core/indices/struct_store/sql_query.py
@@ -6,6 +6,7 @@
 
 from llama_index.core.base.base_query_engine import BaseQueryEngine
 from llama_index.core.base.response.schema import Response
+from llama_index.core.callbacks import CallbackManager
 from llama_index.core.indices.struct_store.container_builder import (
     SQLContextContainerBuilder,
 )
@@ -305,6 +306,7 @@ def __init__(
         llm: Optional[LLM] = None,
         synthesize_response: bool = True,
         response_synthesis_prompt: Optional[BasePromptTemplate] = None,
+        callback_manager: Optional[CallbackManager] = None,
         refine_synthesis_prompt: Optional[BasePromptTemplate] = None,
         verbose: bool = False,
         # deprecated
@@ -314,6 +316,9 @@ def __init__(
         """Initialize params."""
         self._service_context = service_context
         self._llm = llm or llm_from_settings_or_context(Settings, service_context)
+        if callback_manager is not None:
+            self._llm.callback_manager = callback_manager
+
         self._response_synthesis_prompt = (
             response_synthesis_prompt or DEFAULT_RESPONSE_SYNTHESIS_PROMPT_V2
         )
@@ -328,9 +333,8 @@ def __init__(
         self._synthesize_response = synthesize_response
         self._verbose = verbose
         super().__init__(
-            callback_manager=callback_manager_from_settings_or_context(
-                Settings, service_context
-            ),
+            callback_manager=callback_manager
+            or callback_manager_from_settings_or_context(Settings, service_context),
             **kwargs,
         )
 
@@ -434,6 +438,7 @@ def __init__(
         service_context: Optional[ServiceContext] = None,
         context_str_prefix: Optional[str] = None,
         sql_only: bool = False,
+        callback_manager: Optional[CallbackManager] = None,
         verbose: bool = False,
         **kwargs: Any,
     ) -> None:
@@ -448,6 +453,7 @@ def __init__(
             context_str_prefix=context_str_prefix,
             service_context=service_context,
             sql_only=sql_only,
+            callback_manager=callback_manager,
             verbose=verbose,
         )
         super().__init__(
@@ -456,6 +462,7 @@ def __init__(
             refine_synthesis_prompt=refine_synthesis_prompt,
             llm=llm,
             service_context=service_context,
+            callback_manager=callback_manager,
             verbose=verbose,
             **kwargs,
         )
@@ -489,6 +496,7 @@ def __init__(
         service_context: Optional[ServiceContext] = None,
         context_str_prefix: Optional[str] = None,
         sql_only: bool = False,
+        callback_manager: Optional[CallbackManager] = None,
         **kwargs: Any,
     ) -> None:
         """Initialize params."""
@@ -503,6 +511,7 @@ def __init__(
             context_str_prefix=context_str_prefix,
             service_context=service_context,
             sql_only=sql_only,
+            callback_manager=callback_manager,
         )
         super().__init__(
             synthesize_response=synthesize_response,
@@ -510,6 +519,7 @@ def __init__(
             refine_synthesis_prompt=refine_synthesis_prompt,
             llm=llm,
             service_context=service_context,
+            callback_manager=callback_manager,
             **kwargs,
         )
 
@@ -535,6 +545,7 @@ def __init__(
         service_context: Optional[ServiceContext] = None,
         context_str_prefix: Optional[str] = None,
         sql_only: bool = False,
+        callback_manager: Optional[CallbackManager] = None,
         **kwargs: Any,
     ) -> None:
         """Initialize params."""
@@ -547,6 +558,7 @@ def __init__(
             context_str_prefix=context_str_prefix,
             service_context=service_context,
             sql_only=sql_only,
+            callback_manager=callback_manager,
         )
         super().__init__(
             synthesize_response=synthesize_response,
@@ -554,6 +566,7 @@ def __init__(
             refine_synthesis_prompt=refine_synthesis_prompt,
             llm=llm,
             service_context=service_context,
+            callback_manager=callback_manager,
             **kwargs,
         )