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

crewplus/callbacks/run_id_handler.py

@@ -1,106 +1,131 @@
 # File: crewplus/callbacks/run_id_handler.py
 
-from typing import Any, Dict
+from typing import Any, Dict, List
 from uuid import UUID
 import logging
 
 # Langfuse imports with graceful fallback
-try:
-    from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
-    from .async_langfuse_handler import AsyncLangfuseCallbackHandler
-    from langfuse import get_client
-    LANGFUSE_AVAILABLE = True
-except ImportError:
-    LANGFUSE_AVAILABLE = False
-    # Define dummy base classes if langfuse is not available to avoid runtime errors
-    class LangfuseCallbackHandler: pass
-    class AsyncLangfuseCallbackHandler: pass
-    get_client = None
+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 ---
-if LANGFUSE_AVAILABLE:
-    class RunIdCallbackHandler(LangfuseCallbackHandler):
-        """
-        A custom synchronous callback handler that captures the Langchain run_id
-        and adds it to the Langfuse trace metadata. It also stores the mapping
-        between the run_id and the generated trace_id.
-        """
-        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 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 triggered for run_id: {run_id}")
-            # First, let the base handler do its work, which includes creating the trace
-            run = super().on_chain_start(serialized, inputs, run_id, **kwargs)
-            self.logger.debug(f"Base handler's on_chain_start completed. Current trace_id: {self.last_trace_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.")
 
-            try:
-                # The base handler sets `self.last_trace_id` when a trace is created
-                if self.last_trace_id:
-                    self.run_trace_map[str(run_id)] = self.last_trace_id
-                    self.logger.info(f"Mapping run_id '{run_id}' to trace_id '{self.last_trace_id}'.")
+    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
 
-                    # Then, update the created trace with the run_id as metadata
-                    langfuse_client = get_client()
-                    langfuse_client.update_current_trace(
-                        metadata={"langchain_run_id": str(run_id)}
-                    )
-                    self.logger.info(f"Successfully updated trace '{self.last_trace_id}' with metadata for run_id '{run_id}'.")
-                else:
-                    self.logger.warning(f"Could not find 'last_trace_id' for run_id '{run_id}'. Metadata will not be updated.")
+    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
 
-            except Exception as e:
-                self.logger.error(f"Failed to process run_id '{run_id}' in RunIdCallbackHandler: {e}", exc_info=True)
+    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)
 
-            return run
+        # 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)
 
-    class AsyncRunIdCallbackHandler(AsyncLangfuseCallbackHandler):
-        """
-        A custom asynchronous callback handler that captures the Langchain run_id
-        and adds it to the Langfuse trace metadata.
-        """
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-            self.run_trace_map = {}
-            self.logger = logging.getLogger(__name__)
-            self.logger.info("AsyncRunIdCallbackHandler initialized.")
+        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}'.")
 
-        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 triggered for run_id: {run_id}")
-            # First, let the base handler do its work
-            run = await super().on_chain_start(serialized, inputs, run_id, **kwargs)
-            self.logger.debug(f"Async base handler's on_chain_start completed. Current trace_id: {self.last_trace_id}")
-            
-            try:
-                if self.last_trace_id:
-                    self.run_trace_map[str(run_id)] = self.last_trace_id
-                    self.logger.info(f"Async: Mapping 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.")
 
-                    langfuse_client = get_client()
-                    langfuse_client.update_current_trace(
-                        metadata={"langchain_run_id": str(run_id)}
-                    )
-                    self.logger.info(f"Async: Successfully updated trace '{self.last_trace_id}' with metadata for run_id '{run_id}'.")
-                else:
-                    self.logger.warning(f"Async: Could not find 'last_trace_id' for run_id '{run_id}'. Metadata will not be updated.")
+    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
 
-            except Exception as e:
-                self.logger.error(f"Async: Failed to process run_id '{run_id}' in AsyncRunIdCallbackHandler: {e}", exc_info=True)
+    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
 
-            return run
+    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

crewplus/services/tracing_manager.py

@@ -89,20 +89,26 @@ class TracingManager:
     
     def _initialize_langfuse(self):
         """Initializes the Langfuse handler if it's available and enabled."""
+        self.context.logger.debug("Attempting to initialize Langfuse handlers.")
         if not LANGFUSE_AVAILABLE:
             if self.context.enable_tracing is True:
                 self.context.logger.warning("Langfuse is not installed; tracing will be disabled. Install with: pip install langfuse")
+            else:
+                self.context.logger.debug("Langfuse is not installed, skipping handler initialization.")
             return
         
         # Determine if Langfuse should be enabled via an explicit flag or
         # 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())
@@ -112,27 +118,37 @@ class TracingManager:
                     self._async_handlers.append(AsyncLangfuseCallbackHandler())
                     self._async_handlers.append(AsyncRunIdCallbackHandler())
                 
-                self.context.logger.info(f"Langfuse tracing enabled for {self.context.get_model_identifier()}")
+                self.context.logger.info(f"Langfuse tracing enabled for {self.context.get_model_identifier()}. Sync handlers loaded: {len(self._sync_handlers)}")
             except Exception as e:
-                self.context.logger.warning(f"Failed to initialize Langfuse: {e}")
+                self.context.logger.warning(f"Failed to initialize Langfuse: {e}", exc_info=True)
+        else:
+            self.context.logger.info("Langfuse is not enabled, skipping handler initialization.")
     
     def _add_callbacks_to_config(self, config: Optional[dict], handlers: List[Any]) -> dict:
         """A generic helper to add a list of handlers to a config object."""
         if config is None:
             config = {}
         
+        self.context.logger.debug(f"Adding callbacks to config. Have {len(handlers)} handlers to add.")
+
         if not handlers or config.get("metadata", {}).get("tracing_disabled"):
+            self.context.logger.debug("No handlers to add or tracing is disabled for this run.")
             return config
         
         callbacks = config.get("callbacks")
         
         if hasattr(callbacks, 'add_handler') and hasattr(callbacks, 'handlers'):
+            # This block is for CallbackManager instances
+            self.context.logger.debug(f"Config has a CallbackManager with {len(callbacks.handlers)} existing handlers.")
             for handler in handlers:
                 if not any(isinstance(cb, type(handler)) for cb in callbacks.handlers):
                     callbacks.add_handler(handler, inherit=True)
+            self.context.logger.debug(f"CallbackManager now has {len(callbacks.handlers)} handlers.")
             return config
         
+        # This block is for simple lists of callbacks
         current_callbacks = callbacks or []
+        self.context.logger.debug(f"Config has a list with {len(current_callbacks)} existing callbacks.")
         new_callbacks = list(current_callbacks)
         
         for handler in handlers:

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

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

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

@@ -1,17 +1,17 @@
-crewplus-0.2.66.dist-info/METADATA,sha256=nfG_3A6490ii2Hmp9cfDZgWd_aIOebzVa7cyhAsCDao,5424
-crewplus-0.2.66.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.66.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.66.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+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/__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=UurImqJYJg_EnOKB8nzBmG377KL4wI7fl8lz-98_VXY,4964
+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=tTha3NGoQUvO0gwBrx_pJzQ3-2eEz2ZsJtICtI6OeiA,7033
+crewplus/services/tracing_manager.py,sha256=oyKoPAfirAvb9M6Kca6HMNQYo3mBtXCLM3jFkPYU55s,8368
 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
@@ -23,4 +23,4 @@ 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.66.dist-info/RECORD,,
+crewplus-0.2.68.dist-info/RECORD,,