crewplus 0.2.68__py3-none-any.whl → 0.2.71__py3-none-any.whl

crewplus/services/tracing_manager.py

@@ -9,16 +9,11 @@ import logging
 try:
     from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
     from ..callbacks.async_langfuse_handler import AsyncLangfuseCallbackHandler
-    # Import the new custom handlers
-    from ..callbacks.run_id_handler import RunIdCallbackHandler, AsyncRunIdCallbackHandler
     LANGFUSE_AVAILABLE = True
 except ImportError:
     LANGFUSE_AVAILABLE = False
     LangfuseCallbackHandler = None
     AsyncLangfuseCallbackHandler = None
-    # Define dummy classes for the custom handlers to prevent errors if langfuse is not installed
-    class RunIdCallbackHandler: pass
-    class AsyncRunIdCallbackHandler: pass
 
 class TracingContext(Protocol):
     """
@@ -101,24 +96,20 @@ class TracingManager:
         # by detecting its environment variables.
         enable_langfuse = self.context.enable_tracing
         if enable_langfuse is None: # Auto-detect if not explicitly set
-            self.context.logger.debug("enable_tracing is None, auto-detecting via environment variables.")
             langfuse_env_vars = ["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"]
             enable_langfuse = any(os.getenv(var) for var in langfuse_env_vars)
-            self.context.logger.debug(f"Auto-detection result for enable_langfuse: {enable_langfuse}")
         
         if enable_langfuse:
             try:
-                self.context.logger.debug("Langfuse is enabled, creating handlers.")
-                # Create and add both the standard and the run_id-capturing handlers.
-                # The standard handler creates the trace, and the custom one updates it.
-                self._sync_handlers.append(LangfuseCallbackHandler())
-                self._sync_handlers.append(RunIdCallbackHandler())
+                # Create both sync and async handlers. We'll pick one at runtime.
+                sync_handler = LangfuseCallbackHandler()
+                self._sync_handlers.append(sync_handler)
 
                 if AsyncLangfuseCallbackHandler:
-                    self._async_handlers.append(AsyncLangfuseCallbackHandler())
-                    self._async_handlers.append(AsyncRunIdCallbackHandler())
+                    async_handler = AsyncLangfuseCallbackHandler()
+                    self._async_handlers.append(async_handler)
                 
-                self.context.logger.info(f"Langfuse tracing enabled for {self.context.get_model_identifier()}. Sync handlers loaded: {len(self._sync_handlers)}")
+                self.context.logger.info(f"Langfuse tracing enabled for {self.context.get_model_identifier()}")
             except Exception as e:
                 self.context.logger.warning(f"Failed to initialize Langfuse: {e}", exc_info=True)
         else:

crewplus/vectorstores/milvus/vdb_service.py

@@ -5,14 +5,15 @@
 # @Last Modified time: 2025-10-09
 
 import logging
-from typing import List, Dict, Union, Optional
+from typing import List, Dict, Union, Optional, Any
 from langchain_milvus import Milvus
 from langchain_core.embeddings import Embeddings
 from langchain_openai import AzureOpenAIEmbeddings
-from pymilvus import MilvusClient, AsyncMilvusClient
+from pymilvus import MilvusClient, AsyncMilvusClient, connections
 import time
 import asyncio
 import uuid
+from collections import defaultdict
 
 from ...services.init_services import get_model_balancer
 from .schema_milvus import SchemaMilvus, DEFAULT_SCHEMA
@@ -96,6 +97,7 @@ class VDBService(object):
     _async_client: Optional[AsyncMilvusClient] = None
     _instances: Dict[str, Milvus] = {}
     _async_instances: Dict[str, Milvus] = {}
+    _async_instance_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
 
     schema: str
     embedding_function: Embeddings
@@ -168,6 +170,10 @@ class VDBService(object):
         # lazy-initialize async milvus
         # self._async_client = self._initialize_async_milvus_client(provider)
 
+        # Do not initialize the async client here.
+        # It must be lazily initialized within an async context.
+        self._async_client: Optional[AsyncMilvusClient] = None
+        
         self.schema = schema
         self.index_params = self.settings.get("index_params")
         
@@ -256,12 +262,16 @@ class VDBService(object):
 
     async def aget_async_vector_client(self) -> AsyncMilvusClient:
         """
-        Asynchronously returns the active AsyncMilvusClient instance, initializing it if necessary.
-
-        Returns:
-            AsyncMilvusClient: The initialized async client for interacting with the vector database.
+        Lazily initializes and returns the AsyncMilvusClient.
+        This ensures the client is created within the running event loop.
         """
-        return await self._get_or_create_async_client()
+        if self._async_client is None:
+            self.logger.info("Lazily initializing AsyncMilvusClient...")
+            client_args = self._get_milvus_client_args(self._provider)
+            # Use the dedicated async alias
+            client_args['alias'] = self.async_alias
+            self._async_client = AsyncMilvusClient(**client_args)
+        return self._async_client
 
     def get_vector_field(self, collection_name: str) -> str:
         """
@@ -370,6 +380,7 @@ class VDBService(object):
         Asynchronously checks if a collection exists and creates it if it doesn't.
         """
         try:
+            # Call the new lazy initializer for the async client
             client = await self.aget_async_vector_client()
             if check_existence and not await client.has_collection(collection_name):
                 self.logger.info(f"Collection '{collection_name}' does not exist. Creating it.")
@@ -498,130 +509,117 @@ class VDBService(object):
 
     async def _get_or_create_async_client(self) -> AsyncMilvusClient:
         """
-        Lazily initializes and returns the AsyncMilvusClient.
-        This runs the blocking constructor in a separate thread, but also creates
-        a temporary event loop inside that thread to satisfy the client's
-        initialization requirements.
+        Lazily initializes the AsyncMilvusClient.
+        Based on grpcio source, the client MUST be initialized in a thread
+        with a running event loop. Therefore, we initialize it directly in the
+        main async context. The synchronous __init__ is fast enough not to
+        block the event loop meaningfully.
         """
         if self._async_client is None:
-            self.logger.info("Lazily initializing AsyncMilvusClient...")
-            
-            def _create_with_loop():
-                # This function runs in a separate thread via asyncio.to_thread
-                try:
-                    # Check if an event loop exists in this new thread
-                    asyncio.get_running_loop()
-                except RuntimeError:  # 'RuntimeError: There is no current event loop...'
-                    # If not, create and set a new one
-                    loop = asyncio.new_event_loop()
-                    asyncio.set_event_loop(loop)
-
-                # Now, with an event loop present in this thread, initialize the client.
-                # This is still a blocking call, but it's contained in the thread.
-                provider = self.settings.get("vector_store", {}).get("provider")
-                return self._initialize_async_milvus_client(provider)
-
-            self._async_client = await asyncio.to_thread(_create_with_loop)
+            self.logger.info("Lazily initializing AsyncMilvusClient directly in the main event loop...")
+            provider = self.settings.get("vector_store", {}).get("provider")
+            # This is a synchronous call, but it's lightweight and must run here.
+            self._async_client = self._initialize_async_milvus_client(provider)
             
         return self._async_client
 
     async def aget_vector_store(self, collection_name: str, embeddings: Embeddings = None, metric_type: str = "IP") -> Milvus:
         """
         Asynchronously gets a vector store instance, creating it if it doesn't exist.
+        This version is optimized to handle high concurrency using a lock.
         """
         if not collection_name:
             self.logger.error("aget_vector_store called with no collection_name.")
             raise ValueError("collection_name must be provided.")
 
-        # Check for a cached instance. If found, return it immediately.
-        if collection_name in self._async_instances:
-            self.logger.info(f"Returning existing async vector store instance for collection: {collection_name}")
-            return self._async_instances[collection_name]
-
-        self.logger.info(f"Creating new async vector store instance for collection: {collection_name}")
-        if embeddings is None:
-            embeddings = self.get_embeddings()
+        lock = self._async_instance_locks[collection_name]
+        async with lock:
+            if collection_name in self._async_instances:
+                self.logger.info(f"Returning existing async vector store instance for collection: {collection_name} (post-lock)")
+                return self._async_instances[collection_name]
 
-        await self._aensure_collection_exists(collection_name, embeddings, check_existence=True)
+            self.logger.info(f"Creating new async vector store instance for collection: {collection_name}")
+            if embeddings is None:
+                embeddings = self.get_embeddings()
 
-        # try:
-        #     self.logger.info(f"Testing embedding function for collection '{collection_name}'...")
-        #     await embeddings.aembed_query("validation_test_string")
-        #     self.logger.info("Embedding function is valid.")
-        # except Exception as e:
-        #     self.logger.error(
-        #         f"The provided embedding function is invalid and failed with error: {e}. "
-        #         f"Cannot create a vector store for collection '{collection_name}'."
-        #     )
-        #     raise RuntimeError(f"Invalid embedding function provided.") from e
+            # CRITICAL: Ensure the shared async client is initialized *under the lock*
+            # before any operation that might use it.
+            await self._get_or_create_async_client()
 
-        index_params = self.index_params or {
-            "metric_type": metric_type,
-            "index_type": "AUTOINDEX",
-            "params": {}
-        }
-        
-        # Create a dedicated connection_args for the async path with the correct alias
-        async_conn_args = self.connection_args.copy()
-        async_conn_args['alias'] = self.async_alias
-        
-        # For async operations, we MUST instantiate the Milvus object using the SYNCHRONOUS alias
-        # because its __init__ method is synchronous. This is now done in a separate thread.
-        vdb = await self._acreate_milvus_instance_with_retry(
-            collection_name=collection_name,
-            embeddings=embeddings,
-            index_params=index_params,
-            connection_args=async_conn_args  # Pass the async-specific connection args
-        )
+            await self._aensure_collection_exists(collection_name, embeddings, check_existence=True)
 
-        # After successful synchronous initialization, we hot-swap the alias on the
-        # ASYNCHRONOUS client to ensure future async operations use the correct connection.
-        self.logger.info(f"Swapping to async alias for instance of collection {collection_name}")
-        # DO NOT get the async client here, get it outside this function
-        #await self._get_or_create_async_client()
-        vdb.aclient._using = self.async_alias
+            vdb = await self._acreate_milvus_instance_with_retry(
+                collection_name=collection_name,
+                embeddings=embeddings,
+                metric_type=metric_type
+            )
 
-        self._async_instances[collection_name] = vdb
+            self.logger.info(f"Swapping to async alias for instance of collection {collection_name}")
+            vdb.aclient._using = self.async_alias
 
-        return vdb
+            self._async_instances[collection_name] = vdb
+            return vdb
 
-    async def _acreate_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict, connection_args: Optional[dict] = None) -> Milvus:
+    async def _acreate_milvus_instance_with_retry(
+        self,
+        embeddings: Embeddings,
+        collection_name: str,
+        metric_type: str = "IP",
+    ) -> Milvus:
         """
-        Asynchronously creates a Milvus instance with a retry mechanism, running the synchronous
-        constructor in a separate thread to avoid blocking the event loop.
+        Asynchronously creates a Milvus instance with retry logic, ensuring the connection
+        is established in the target thread.
         """
-        retries = 2
-        conn_args = connection_args if connection_args is not None else self.connection_args
-        
-        def _create_instance():
-            # This synchronous function will be run in a thread
-            return Milvus(
-                embedding_function=embeddings,
-                collection_name=collection_name,
-                connection_args=conn_args,
-                index_params=index_params
-            )
-
-        self.logger.info(f"Creating Milvus instance for collection '{collection_name}' in a separate thread...")
-        self.logger.info(f"Connection args: {conn_args}")
+        retries = 3
+        last_exception = None
 
-        for attempt in range(retries + 1):
+        for attempt in range(retries):
             try:
-                # Run the blocking constructor in a separate thread
-                vdb = await asyncio.to_thread(_create_instance)
-                
-                self.logger.info(f"Successfully connected to Milvus for collection '{collection_name}' on attempt {attempt + 1}.")
-                return vdb  # Return on success
+                conn_args = self.connection_args.copy()
+                # Langchain's Milvus class will use the alias to find the connection.
+                conn_args["alias"] = self.sync_alias
+
+                def _create_instance_in_thread():
+                    # --- START: CRITICAL FIX ---
+                    # Manually connect within the thread before creating the Milvus instance.
+                    # This ensures pymilvus registers the connection details for the current thread.
+                    try:
+                        connections.connect(**conn_args)
+                        self.logger.info(f"Successfully connected to Milvus with alias '{self.sync_alias}' in thread.")
+                    except Exception as e:
+                        self.logger.error(f"Failed to manually connect in thread: {e}")
+                        raise
+
+                    # Now, creating the Milvus instance will find the existing connection via the alias.
+                    instance = Milvus(
+                        embedding_function=embeddings,
+                        collection_name=collection_name,
+                        connection_args=conn_args, # Pass args for completeness
+                        # metric_type=metric_type,  # <-- CRITICAL FIX: REMOVE THIS LINE
+                        consistency_level="Strong",
+                        # --- START: CRITICAL FIX ---
+                        # Pass self.index_params to the Milvus constructor here
+                        index_params=self.index_params,
+                        # --- END: CRITICAL FIX ---
+                    )
+                    return instance
+                    # --- END: CRITICAL FIX ---
+
+                self.logger.info(f"Attempt {attempt + 1}/{retries}: Creating Milvus instance for collection '{collection_name}' in a separate thread...")
+                vdb = await asyncio.to_thread(_create_instance_in_thread)
+                self.logger.info("Successfully created Milvus instance.")
+                return vdb
+
             except Exception as e:
+                last_exception = e
                 self.logger.warning(
-                    f"Attempt {attempt + 1}/{retries + 1} to connect to Milvus for collection '{collection_name}' failed: {e}"
+                    f"Attempt {attempt + 1}/{retries} failed to create Milvus instance: {e}. Retrying in {2 ** attempt}s..."
                 )
-                if attempt < retries:
-                    self.logger.info("Retrying in 3 seconds...")
-                    await asyncio.sleep(3) # Use async sleep
-                else:
-                    self.logger.error(f"Failed to connect to Milvus for collection '{collection_name}' after {retries + 1} attempts.")
-                    raise RuntimeError(f"Could not connect to Milvus after {retries + 1} attempts.") from e
+                await asyncio.sleep(2 ** attempt)
+        
+        raise RuntimeError(
+            f"Failed to create Milvus instance after {retries} retries."
+        ) from last_exception
 
     def _create_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict, connection_args: Optional[dict] = None) -> Milvus:
         """

{crewplus-0.2.68.dist-info → crewplus-0.2.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.68
+Version: 0.2.71
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

{crewplus-0.2.68.dist-info → crewplus-0.2.71.dist-info}/RECORD

@@ -1,26 +1,25 @@
-crewplus-0.2.68.dist-info/METADATA,sha256=sqbjQMZRzh7U0AajzZZ1lvgXuZzRsYDI3_pC9AKYLYg,5424
-crewplus-0.2.68.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.68.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.68.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.71.dist-info/METADATA,sha256=x45o-KC-K6zW8OIsL3eEwcS-mQ2wupMQYHVo5JANcLs,5424
+crewplus-0.2.71.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.71.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.71.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/callbacks/__init__.py,sha256=YG7ieeb91qEjp1zF0-inEN7mjZ7yT_D2yzdWFT8Z1Ws,63
 crewplus/callbacks/async_langfuse_handler.py,sha256=A4uFeLpvOUdc58M7sZoE65_C1V98u0QCvx5jUquM0pM,7006
-crewplus/callbacks/run_id_handler.py,sha256=T_YqmGmUp2_DIc01dlVY-aC43NKcevfRejAWN-lft6M,5266
 crewplus/services/__init__.py,sha256=V1CG8b2NOmRzNgQH7BPl4KVxWSYJH5vfEsW1wVErKNE,375
 crewplus/services/azure_chat_model.py,sha256=iWzJ2GQFSNmwJx-2O5_xKPSB6VVc-7T6bcfFI8_WezA,5521
 crewplus/services/gemini_chat_model.py,sha256=DYqz01H2TIHiCDQesSozVfOsMigno6QGwOtIweg7UHk,40103
 crewplus/services/init_services.py,sha256=tc1ti8Yufo2ixlJpwg8uH0KmoyQ4EqxCOe4uTEWnlRM,2413
 crewplus/services/model_load_balancer.py,sha256=Q9Gx3GrbKworU-Ytxeqp0ggHSgZ1Q6brtTk-nCl4sak,12095
-crewplus/services/tracing_manager.py,sha256=oyKoPAfirAvb9M6Kca6HMNQYo3mBtXCLM3jFkPYU55s,8368
+crewplus/services/tracing_manager.py,sha256=_C4zYj6o_k5mDWY7vM8UeaVaXp8SqrkvYOb0Jj-y3sY,7566
 crewplus/utils/__init__.py,sha256=2Gk1n5srFJQnFfBuYTxktdtKOVZyNrFcNaZKhXk35Pw,142
 crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
 crewplus/utils/schema_document_updater.py,sha256=frvffxn2vbi71fHFPoGb9hq7gH2azmmdq17p-Fumnvg,7322
 crewplus/vectorstores/milvus/__init__.py,sha256=OeYv2rdyG7tcREIjBJPyt2TbE54NvyeRoWMe7LwopRE,245
 crewplus/vectorstores/milvus/milvus_schema_manager.py,sha256=-QRav-hzu-XWeJ_yKUMolal_EyMUspSg-nvh5sqlrlQ,11442
 crewplus/vectorstores/milvus/schema_milvus.py,sha256=wwNpfqsKS0xeozZES40IvB0iNwUtpCall_7Hkg0dL1g,27223
-crewplus/vectorstores/milvus/vdb_service.py,sha256=Jo2GWzdEuDSPlADCWA7wTgDbYQ65QiYWgFGbtkG9vG8,37789
+crewplus/vectorstores/milvus/vdb_service.py,sha256=U8I6IUYZK0gCe1R9rTnVezvZfEcUS9UEbKEoeJPX8kY,37528
 docs/GeminiChatModel.md,sha256=zZYyl6RmjZTUsKxxMiC9O4yV70MC4TD-IGUmWhIDBKA,8677
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.68.dist-info/RECORD,,
+crewplus-0.2.71.dist-info/RECORD,,

crewplus/callbacks/run_id_handler.py

@@ -1,131 +0,0 @@
-# File: crewplus/callbacks/run_id_handler.py
-
-from typing import Any, Dict, List
-from uuid import UUID
-import logging
-
-# Langfuse imports with graceful fallback
-from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
-from .async_langfuse_handler import AsyncLangfuseCallbackHandler
-from langfuse import get_client
-from langchain_core.messages import BaseMessage
-#from langchain_core.outputs import LLMResult
-LANGFUSE_AVAILABLE = True
-
-# --- Custom Callback Handlers to capture the run_id ---
-
-class RunIdCallbackHandler(LangfuseCallbackHandler):
-    """
-    A custom handler that injects the LangChain run_id into the metadata
-    before the Langfuse observation is created.
-    """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.run_trace_map = {}
-        # Use a named logger for better context
-        self.logger = logging.getLogger(__name__)
-        self.logger.info("RunIdCallbackHandler initialized.")
-
-    def _inject_run_id_to_metadata(self, run_id: UUID, kwargs: Dict[str, Any]) -> Dict[str, Any]:
-        """Helper to safely add the run_id to the metadata in kwargs."""
-        metadata = kwargs.get("metadata") or {}
-        metadata["langchain_run_id"] = str(run_id)
-        kwargs["metadata"] = metadata
-        return kwargs
-
-    def on_chain_start(
-        self,
-        serialized: Dict[str, Any],
-        inputs: Dict[str, Any],
-        *,
-        run_id: UUID,
-        **kwargs: Any,
-    ) -> Any:
-        self.logger.debug(f"[on_chain_start] Intercepted run_id: {run_id}")
-        kwargs = self._inject_run_id_to_metadata(run_id, kwargs)
-        
-        # Call the base handler with the modified kwargs
-        result = super().on_chain_start(serialized, inputs, run_id=run_id, **kwargs)
-        
-        # We still map the trace_id for easy retrieval in tests
-        if self.last_trace_id:
-            self.run_trace_map[str(run_id)] = self.last_trace_id
-            self.logger.info(f"[on_chain_start] Mapped run_id '{run_id}' to trace_id '{self.last_trace_id}'.")
-            
-        return result
-
-    def on_chat_model_start(
-        self,
-        serialized: Dict[str, Any],
-        messages: List[List[BaseMessage]],
-        *,
-        run_id: UUID,
-        parent_run_id: UUID | None = None,
-        **kwargs: Any,
-    ) -> Any:
-        self.logger.debug(f"[on_chat_model_start] Intercepted run_id: {run_id}")
-        
-        # Only inject the run_id if this is the root of the trace
-        if parent_run_id is None:
-            kwargs = self._inject_run_id_to_metadata(run_id, kwargs)
-
-        # Call the base handler with potentially modified kwargs
-        result = super().on_chat_model_start(serialized, messages, run_id=run_id, parent_run_id=parent_run_id, **kwargs)
-
-        if parent_run_id is None and self.last_trace_id:
-            self.run_trace_map[str(run_id)] = self.last_trace_id
-            self.logger.info(f"[on_chat_model_start] Mapped root run_id '{run_id}' to trace_id '{self.last_trace_id}'.")
-
-        return result
-
-# You would similarly update the AsyncRunIdCallbackHandler if you use it
-class AsyncRunIdCallbackHandler(AsyncLangfuseCallbackHandler):
-    """
-    An async custom handler that injects the LangChain run_id into the metadata
-    before the Langfuse observation is created.
-    """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.run_trace_map = {}
-        self.logger = logging.getLogger(__name__)
-        self.logger.info("AsyncRunIdCallbackHandler initialized.")
-
-    def _inject_run_id_to_metadata(self, run_id: UUID, kwargs: Dict[str, Any]) -> Dict[str, Any]:
-        """Helper to safely add the run_id to the metadata in kwargs."""
-        metadata = kwargs.get("metadata") or {}
-        metadata["langchain_run_id"] = str(run_id)
-        kwargs["metadata"] = metadata
-        return kwargs
-
-    async def on_chain_start(
-        self,
-        serialized: Dict[str, Any],
-        inputs: Dict[str, Any],
-        *,
-        run_id: UUID,
-        **kwargs: Any,
-    ) -> Any:
-        self.logger.debug(f"Async [on_chain_start] Intercepted run_id: {run_id}")
-        kwargs = self._inject_run_id_to_metadata(run_id, kwargs)
-        result = await super().on_chain_start(serialized, inputs, run_id=run_id, **kwargs)
-        if self.last_trace_id:
-            self.run_trace_map[str(run_id)] = self.last_trace_id
-        return result
-
-    async def on_chat_model_start(
-        self,
-        serialized: Dict[str, Any],
-        messages: List[List[BaseMessage]],
-        *,
-        run_id: UUID,
-        parent_run_id: UUID | None = None,
-        **kwargs: Any,
-    ) -> Any:
-        self.logger.debug(f"Async [on_chat_model_start] Intercepted run_id: {run_id}")
-        if parent_run_id is None:
-            kwargs = self._inject_run_id_to_metadata(run_id, kwargs)
-        result = await super().on_chat_model_start(serialized, messages, run_id=run_id, parent_run_id=parent_run_id, **kwargs)
-        
-        if parent_run_id is None and self.last_trace_id:
-            self.run_trace_map[str(run_id)] = self.last_trace_id
-        return result