crewplus 0.2.25__py3-none-any.whl → 0.2.27__py3-none-any.whl

crewplus/services/azure_chat_model.py

@@ -4,34 +4,29 @@ from typing import Any, Optional
 
 from langchain_openai.chat_models.azure import AzureChatOpenAI
 from pydantic import Field
-
-# Langfuse imports with graceful fallback
-try:
-    from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
-    LANGFUSE_AVAILABLE = True
-except ImportError:
-    LANGFUSE_AVAILABLE = False
-    LangfuseCallbackHandler = None
+from .tracing_manager import TracingManager, TracingContext
 
 class TracedAzureChatOpenAI(AzureChatOpenAI):
     """
-    Wrapper for AzureChatOpenAI that integrates with Langfuse for tracing.
+    Wrapper for AzureChatOpenAI that integrates with tracing services like Langfuse.
     
-    This class automatically handles Langfuse callback integration, making it easier
+    This class automatically handles callback integration, making it easier
     to trace and debug your interactions with the Azure OpenAI service.
 
-    **Langfuse Integration:**
-    Langfuse tracing is automatically enabled when environment variables are set:
+    **Tracing Integration (e.g., Langfuse):**
+    Tracing is automatically enabled when the respective environment variables are set.
+    For Langfuse:
     - LANGFUSE_PUBLIC_KEY: Your Langfuse public key
     - LANGFUSE_SECRET_KEY: Your Langfuse secret key  
     - LANGFUSE_HOST: Langfuse host URL (optional, defaults to https://cloud.langfuse.com)
     
-    You can also configure it explicitly or disable it. Session and user tracking 
-    can be set per call via metadata in the `config` argument.
+    You can explicitly control this with the `enable_tracing` parameter or disable
+    it for specific calls by adding `{"metadata": {"tracing_disabled": True}}`
+    to the `config` argument.
 
     Attributes:
         logger (Optional[logging.Logger]): An optional logger instance.
-        enable_langfuse (Optional[bool]): Enable/disable Langfuse tracing (auto-detect if None).
+        enable_tracing (Optional[bool]): Enable/disable tracing (auto-detect if None).
 
     Example:
         .. code-block:: python
@@ -54,21 +49,21 @@ class TracedAzureChatOpenAI(AzureChatOpenAI):
             response = model.invoke("Hello, how are you?")
             print("Text response:", response.content)
 
-            # --- Langfuse tracing with session/user tracking ---
+            # --- Tracing with session/user tracking (for Langfuse) ---
             response = model.invoke(
                 "What is AI?",
                 config={
                     "metadata": {
                         "langfuse_session_id": "chat-session-123",
-                        "langfuse_user_id": "user-456"
+                        "user_id": "user-456"
                     }
                 }
             )
             
-            # --- Disable Langfuse for specific calls ---
+            # --- Disable tracing for a specific call ---
             response = model.invoke(
                 "Hello without tracing",
-                config={"metadata": {"langfuse_disabled": True}}
+                config={"metadata": {"tracing_disabled": True}}
             )
 
             # --- Asynchronous Streaming Usage ---
@@ -86,9 +81,9 @@ class TracedAzureChatOpenAI(AzureChatOpenAI):
             # asyncio.run(main())
     """
     logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance", exclude=True)
-    enable_langfuse: Optional[bool] = Field(default=None, description="Enable Langfuse tracing (auto-detect if None)")
+    enable_tracing: Optional[bool] = Field(default=None, description="Enable tracing (auto-detect if None)")
     
-    langfuse_handler: Optional[LangfuseCallbackHandler] = Field(default=None, exclude=True)
+    _tracing_manager: Optional[TracingManager] = None
 
     def __init__(self, **kwargs: Any):
         super().__init__(**kwargs)
@@ -100,102 +95,35 @@ class TracedAzureChatOpenAI(AzureChatOpenAI):
                 self.logger.addHandler(logging.StreamHandler())
                 self.logger.setLevel(logging.INFO)
         
-        # Initialize Langfuse handler
-        self._initialize_langfuse()
-
-    def _initialize_langfuse(self):
-        """Initialize Langfuse handler if enabled and available."""
-        if not LANGFUSE_AVAILABLE:
-            if self.enable_langfuse is True:
-                self.logger.warning("Langfuse is not installed. Install with: pip install langfuse")
-            return
-        
-        # Auto-detect if Langfuse should be enabled
-        if self.enable_langfuse is None:
-            langfuse_env_vars = ["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"]
-            self.enable_langfuse = any(os.getenv(var) for var in langfuse_env_vars)
-        
-        if not self.enable_langfuse:
-            return
-        
-        try:
-            self.langfuse_handler = LangfuseCallbackHandler()
-            self.logger.info(f"Langfuse tracing enabled for TracedAzureChatOpenAI with deployment: {self.deployment_name}")
-        except Exception as e:
-            self.logger.warning(f"Failed to initialize Langfuse: {e}")
-            self.langfuse_handler = None
+        self._tracing_manager = TracingManager(self)
 
-    def invoke(self, input, config=None, **kwargs):
-        """Override invoke to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        if self.langfuse_handler:
-            # Do not trace if disabled via metadata
-            if config.get("metadata", {}).get("langfuse_disabled"):
-                return super().invoke(input, config=config, **kwargs)
+    def get_model_identifier(self) -> str:
+        """Return a string identifying this model for tracing and logging."""
+        return f"{self.__class__.__name__} (deployment='{self.deployment_name}')"
 
-            callbacks = config.get("callbacks", [])
-            has_langfuse = any(isinstance(callback, LangfuseCallbackHandler) for callback in callbacks)
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self.langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
+    def invoke(self, input, config=None, **kwargs):
+        config = self._tracing_manager.add_callbacks_to_config(config)
         return super().invoke(input, config=config, **kwargs)
 
     async def ainvoke(self, input, config=None, **kwargs):
-        """Override ainvoke to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        if self.langfuse_handler:
-            # Do not trace if disabled via metadata
-            if config.get("metadata", {}).get("langfuse_disabled"):
-                return await super().ainvoke(input, config=config, **kwargs)
-
-            callbacks = config.get("callbacks", [])
-            has_langfuse = any(isinstance(callback, LangfuseCallbackHandler) for callback in callbacks)
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self.langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
+        config = self._tracing_manager.add_callbacks_to_config(config)
         return await super().ainvoke(input, config=config, **kwargs)
 
     def stream(self, input, config=None, **kwargs):
-        """Override stream to add Langfuse callback and request usage metadata."""
-        if config is None:
-            config = {}
-
         # Add stream_options to get usage data for Langfuse
         stream_options = kwargs.get("stream_options", {})
         stream_options["include_usage"] = True
         kwargs["stream_options"] = stream_options
-
-        # Add Langfuse callback if enabled and not already present
-        if self.langfuse_handler and not config.get("metadata", {}).get("langfuse_disabled"):
-            callbacks = config.get("callbacks", [])
-            if not any(isinstance(c, LangfuseCallbackHandler) for c in callbacks):
-                config["callbacks"] = callbacks + [self.langfuse_handler]
         
+        config = self._tracing_manager.add_callbacks_to_config(config)
         yield from super().stream(input, config=config, **kwargs)
 
     async def astream(self, input, config=None, **kwargs) :
-        """Override astream to add Langfuse callback and request usage metadata."""
-        if config is None:
-            config = {}
-
         # Add stream_options to get usage data for Langfuse
         stream_options = kwargs.get("stream_options", {})
         stream_options["include_usage"] = True
         kwargs["stream_options"] = stream_options
         
-        # Add Langfuse callback if enabled and not already present
-        if self.langfuse_handler and not config.get("metadata", {}).get("langfuse_disabled"):
-            callbacks = config.get("callbacks", [])
-            if not any(isinstance(c, LangfuseCallbackHandler) for c in callbacks):
-                config["callbacks"] = callbacks + [self.langfuse_handler]
-
+        config = self._tracing_manager.add_callbacks_to_config(config)
         async for chunk in super().astream(input, config=config, **kwargs):
             yield chunk

crewplus/services/gemini_chat_model.py

@@ -21,14 +21,7 @@ from langchain_core.callbacks import (
 )
 from pydantic import Field, SecretStr
 from langchain_core.utils import convert_to_secret_str
-
-# Langfuse imports with graceful fallback
-try:
-    from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
-    LANGFUSE_AVAILABLE = True
-except ImportError:
-    LANGFUSE_AVAILABLE = False
-    LangfuseCallbackHandler = None
+from .tracing_manager import TracingManager, TracingContext
 
 class GeminiChatModel(BaseChatModel):
     """Custom chat model for Google Gemini, supporting text, image, and video.
@@ -41,8 +34,9 @@ class GeminiChatModel(BaseChatModel):
     API keys can be provided directly or loaded from the `GOOGLE_API_KEY`
     environment variable.
 
-    **Langfuse Integration:**
-    Langfuse tracing is automatically enabled when environment variables are set:
+    **Tracing Integration:**
+    Tracing (e.g., with Langfuse) is automatically enabled when the respective
+    environment variables are set. For Langfuse:
     - LANGFUSE_PUBLIC_KEY: Your Langfuse public key
     - LANGFUSE_SECRET_KEY: Your Langfuse secret key  
     - LANGFUSE_HOST: Langfuse host URL (optional, defaults to https://cloud.langfuse.com)
@@ -58,7 +52,7 @@ class GeminiChatModel(BaseChatModel):
         top_p (Optional[float]): The top-p (nucleus) sampling parameter.
         top_k (Optional[int]): The top-k sampling parameter.
         logger (Optional[logging.Logger]): An optional logger instance.
-        enable_langfuse (Optional[bool]): Enable/disable Langfuse tracing (auto-detect if None).
+        enable_tracing (Optional[bool]): Enable/disable all tracing (auto-detect if None).
 
     Example:
         .. code-block:: python
@@ -83,7 +77,7 @@ class GeminiChatModel(BaseChatModel):
             response = model.invoke("Hello, how are you?")
             print("Text response:", response.content)
 
-            # --- Langfuse tracing with session/user tracking ---
+            # --- Tracing with session/user tracking (for Langfuse) ---
             response = model.invoke(
                 "What is AI?",
                 config={
@@ -189,25 +183,25 @@ class GeminiChatModel(BaseChatModel):
             # --- Disable Langfuse for specific calls ---
             response = model.invoke(
                 "Hello without tracing",
-                config={"metadata": {"langfuse_disabled": True}}
+                config={"metadata": {"tracing_disabled": True}}
             )
     """
     
     # Model configuration
-    model_name: str = Field(default="gemini-2.0-flash", description="The Google model name to use")
+    model_name: str = Field(default="gemini-2.5-flash", description="The Google model name to use")
     google_api_key: Optional[SecretStr] = Field(default=None, description="Google API key")
     temperature: Optional[float] = Field(default=0.7, description="Sampling temperature")
     max_tokens: Optional[int] = Field(default=None, description="Maximum tokens to generate")
     top_p: Optional[float] = Field(default=None, description="Top-p sampling parameter")
     top_k: Optional[int] = Field(default=None, description="Top-k sampling parameter")
-    logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance")
     
-    # Langfuse configuration
-    enable_langfuse: Optional[bool] = Field(default=None, description="Enable Langfuse tracing (auto-detect if None)")
+    # Configuration for tracing and logging
+    logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance", exclude=True)
+    enable_tracing: Optional[bool] = Field(default=None, description="Enable tracing (auto-detect if None)")
     
-    # Internal clients
+    # Internal clients and managers
     _client: Optional[genai.Client] = None
-    _langfuse_handler: Optional[LangfuseCallbackHandler] = None
+    _tracing_manager: Optional[TracingManager] = None
     
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -215,7 +209,7 @@ class GeminiChatModel(BaseChatModel):
         # Initialize logger
         if self.logger is None:
             self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
-            if not self.logger.handlers: # and not getattr(self.logger, 'propagate', True):
+            if not self.logger.handlers:
                 self.logger.addHandler(logging.StreamHandler())
                 self.logger.setLevel(logging.INFO)
         
@@ -227,143 +221,40 @@ class GeminiChatModel(BaseChatModel):
         
         # Initialize the Google GenAI client
         if self.google_api_key:
-            self._client = genai.Client(
-                api_key=self.google_api_key.get_secret_value()
-            )
+            self._client = genai.Client(api_key=self.google_api_key.get_secret_value())
             self.logger.info(f"Initialized GeminiChatModel with model: {self.model_name}")
         else:
             error_msg = "Google API key is required. Set GOOGLE_API_KEY environment variable or pass google_api_key parameter."
             self.logger.error(error_msg)
             raise ValueError(error_msg)
         
-        # Initialize Langfuse handler
-        self._initialize_langfuse()
-    
-    def _initialize_langfuse(self):
-        """Initialize Langfuse handler if enabled and available."""
-        if not LANGFUSE_AVAILABLE:
-            if self.enable_langfuse is True:
-                self.logger.warning("Langfuse is not installed. Install with: pip install langfuse")
-            return
-        
-        # Auto-detect if Langfuse should be enabled
-        if self.enable_langfuse is None:
-            # Check if Langfuse environment variables are set
-            langfuse_env_vars = ["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"]
-            self.enable_langfuse = any(os.getenv(var) for var in langfuse_env_vars)
-        
-        if not self.enable_langfuse:
-            return
-        
-        try:
-            # Initialize Langfuse handler with minimal config
-            # Session/user tracking will be handled per call via metadata
-            self._langfuse_handler = LangfuseCallbackHandler()
-            self.logger.info("Langfuse tracing enabled for GeminiChatModel")
-            
-        except Exception as e:
-            self.logger.warning(f"Failed to initialize Langfuse: {e}")
-            self._langfuse_handler = None
-    
-    def _should_add_langfuse_callback(self, run_manager: Optional[CallbackManagerForLLMRun] = None) -> bool:
-        """Check if Langfuse callback should be added."""
-        if not self._langfuse_handler:
-            return False
-        
-        # Check if Langfuse is already in the callback manager
-        if run_manager and hasattr(run_manager, 'handlers'):
-            has_langfuse = any(
-                isinstance(handler, LangfuseCallbackHandler) 
-                for handler in run_manager.handlers
-            )
-            if has_langfuse:
-                return False
-        
-        return True
+        self._tracing_manager = TracingManager(self)
+
+    def get_model_identifier(self) -> str:
+        """Return a string identifying this model for tracing and logging."""
+        return f"{self.__class__.__name__} (model='{self.model_name}')"
 
     def invoke(self, input, config=None, **kwargs):
-        """Override invoke to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        # Add Langfuse callback if enabled and not already present
-        if self._langfuse_handler:
-            callbacks = config.get("callbacks", [])
-            
-            # Check if Langfuse callback is already present
-            has_langfuse = any(
-                isinstance(callback, LangfuseCallbackHandler) 
-                for callback in callbacks
-            )
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self._langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
+        """Override invoke to add tracing callbacks automatically."""
+        config = self._tracing_manager.add_callbacks_to_config(config)
         return super().invoke(input, config=config, **kwargs)
 
     async def ainvoke(self, input, config=None, **kwargs):
-        """Override ainvoke to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        # Add Langfuse callback if enabled and not already present
-        if self._langfuse_handler:
-            callbacks = config.get("callbacks", [])
-            
-            # Check if Langfuse callback is already present
-            has_langfuse = any(
-                isinstance(callback, LangfuseCallbackHandler) 
-                for callback in callbacks
-            )
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self._langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
+        """Override ainvoke to add tracing callbacks automatically."""
+        config = self._tracing_manager.add_callbacks_to_config(config)
         return await super().ainvoke(input, config=config, **kwargs)
 
     def stream(self, input, config=None, **kwargs):
-        """Override stream to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        # Add Langfuse callback if enabled and not already present
-        if self._langfuse_handler:
-            callbacks = config.get("callbacks", [])
-            
-            # Check if Langfuse callback is already present
-            has_langfuse = any(
-                isinstance(callback, LangfuseCallbackHandler) 
-                for callback in callbacks
-            )
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self._langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
+        """Override stream to add tracing callbacks automatically."""
+        config = self._tracing_manager.add_callbacks_to_config(config)
         return super().stream(input, config=config, **kwargs)
 
     async def astream(self, input, config=None, **kwargs):
-        """Override astream to add Langfuse callback automatically."""
-        if config is None:
-            config = {}
-        
-        # Add Langfuse callback if enabled and not already present
-        if self._langfuse_handler:
-            callbacks = config.get("callbacks", [])
-            
-            # Check if Langfuse callback is already present
-            has_langfuse = any(
-                isinstance(callback, LangfuseCallbackHandler) 
-                for callback in callbacks
-            )
-            
-            if not has_langfuse:
-                callbacks = callbacks + [self._langfuse_handler]
-                config = {**config, "callbacks": callbacks}
-        
-        return super().astream(input, config=config, **kwargs)
+        """Override astream to add tracing callbacks automatically."""
+        config = self._tracing_manager.add_callbacks_to_config(config)
+        # We must call an async generator,
+        async for chunk in super().astream(input, config=config, **kwargs):
+            yield chunk
     
     @property
     def _llm_type(self) -> str:
@@ -647,6 +538,59 @@ class GeminiChatModel(BaseChatModel):
             }
         return part_dict
 
+    def _map_usage_metadata(self, usage_metadata: Any) -> Optional[dict]:
+        """
+        Maps Google's rich usage metadata to LangChain's expected format,
+        including detailed breakdowns by modality.
+        """
+        if not usage_metadata:
+            return None
+        
+        # --- Basic Token Counts ---
+        input_tokens = getattr(usage_metadata, "prompt_token_count", 0)
+        output_tokens = getattr(usage_metadata, "candidates_token_count", 0)
+        thoughts_tokens = getattr(usage_metadata, "thoughts_token_count", 0)
+        total_tokens = getattr(usage_metadata, "total_token_count", 0)
+
+        # In some cases, total_tokens is not provided, so we calculate it
+        if total_tokens == 0 and (input_tokens > 0 or output_tokens > 0):
+            total_tokens = input_tokens + output_tokens
+
+        # --- Detailed Token Counts (The Fix) ---
+        input_details = {}
+        # The `prompt_tokens_details` is a list of ModalityTokenCount objects.
+        # We convert it to a dictionary.
+        if prompt_details_list := getattr(usage_metadata, "prompt_tokens_details", None):
+            for detail in prompt_details_list:
+                # Convert enum e.g., <MediaModality.TEXT: 'TEXT'> to "text"
+                modality_key = detail.modality.name.lower()
+                input_details[modality_key] = detail.token_count
+        
+        # Add cached tokens to input details if present
+        #if cached_tokens := getattr(usage_metadata, "cached_content_token_count", 0):
+        #    input_details["cached_content"] = cached_tokens
+
+        output_details = {}
+        # The `candidates_tokens_details` is also a list, so we convert it.
+        if candidate_details_list := getattr(usage_metadata, "candidates_tokens_details", None):
+            for detail in candidate_details_list:
+                modality_key = detail.modality.name.lower()
+                output_details[modality_key] = detail.token_count
+
+        # --- Construct the final dictionary ---
+        final_metadata = {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "thoughts_tokens": thoughts_tokens,
+            "total_tokens": total_tokens,
+        }
+        if input_details:
+            final_metadata["input_token_details"] = input_details
+        if output_details:
+            final_metadata["output_token_details"] = output_details
+            
+        return final_metadata
+
     def _extract_usage_metadata(self, response) -> Optional[Any]:
         """Extracts the raw usage_metadata object from a Google GenAI response."""
         if hasattr(response, 'usage_metadata') and response.usage_metadata:
@@ -672,11 +616,10 @@ class GeminiChatModel(BaseChatModel):
         generated_text = response.text
         finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
         
-        # Extract usage metadata for token tracking
+        # Use the new mapping function here for invoke calls
         usage_metadata = self._extract_usage_metadata(response)
-        usage_dict = usage_metadata.dict() if usage_metadata and hasattr(usage_metadata, "dict") else {}
+        usage_dict = self._map_usage_metadata(usage_metadata) or {}
 
-        # Create AIMessage with usage information in response_metadata
         message = AIMessage(
             content=generated_text,
             response_metadata={
@@ -686,11 +629,9 @@ class GeminiChatModel(BaseChatModel):
             }
         )
         
-        # For non-streaming, we include the usage dict in generation_info.
-        # This is another field that callback handlers like Langfuse might inspect.
         generation = ChatGeneration(
             message=message,
-            generation_info=usage_dict if usage_dict else None
+            generation_info={"token_usage": usage_dict} if usage_dict else None
         )
         
         # We also construct the llm_output dictionary in the format expected
@@ -786,31 +727,19 @@ class GeminiChatModel(BaseChatModel):
             
             final_usage_metadata = None
             for chunk_response in stream:
-                # The usage metadata is on the chunk response itself. We update
-                # our variable on each chunk that has it to ensure we get the
-                # final, cumulative count at the end of the stream.
                 if chunk_response.usage_metadata:
                     final_usage_metadata = self._extract_usage_metadata(chunk_response)
 
-                if text_content := chunk_response.text:
-                    chunk = self._create_chat_generation_chunk(chunk_response)
-                    if run_manager:
-                        run_manager.on_llm_new_token(text_content, chunk=chunk)
-                    yield chunk
+                if chunk_response.text:
+                    yield self._create_chat_generation_chunk(chunk_response)
             
-            # After the stream is exhausted, we yield a final, empty chunk
-            # containing the full usage details. LangChain merges this into the
-            # final result, making it available to callback handlers.
+            # **FIX:** Yield a final chunk with the mapped usage data
             if final_usage_metadata:
-                usage_dict = final_usage_metadata.dict() if hasattr(final_usage_metadata, "dict") else {}
-                final_generation_info = {
-                    "token_usage": usage_dict,
-                    "model_name": self.model_name
-                }
-                yield ChatGenerationChunk(
-                    message=AIMessageChunk(content=""),
-                    generation_info=final_generation_info
-                )
+                lc_usage_metadata = self._map_usage_metadata(final_usage_metadata)
+                if lc_usage_metadata:
+                    yield ChatGenerationChunk(
+                        message=AIMessageChunk(content="", usage_metadata=lc_usage_metadata)
+                    )
 
         except Exception as e:
             self.logger.error(f"Error streaming content: {e}", exc_info=True)
@@ -839,31 +768,19 @@ class GeminiChatModel(BaseChatModel):
             
             final_usage_metadata = None
             async for chunk_response in stream:
-                # The usage metadata is on the chunk response itself. We update
-                # our variable on each chunk that has it to ensure we get the
-                # final, cumulative count at the end of the stream.
                 if chunk_response.usage_metadata:
                     final_usage_metadata = self._extract_usage_metadata(chunk_response)
                 
-                if text_content := chunk_response.text:
-                    chunk = self._create_chat_generation_chunk(chunk_response)
-                    if run_manager:
-                        await run_manager.on_llm_new_token(text_content, chunk=chunk)
-                    yield chunk
+                if chunk_response.text:
+                    yield self._create_chat_generation_chunk(chunk_response)
             
-            # After the stream is exhausted, we yield a final, empty chunk
-            # containing the full usage details. LangChain merges this into the
-            # final result, making it available to callback handlers.
+            # **FIX:** Yield a final chunk with the mapped usage data
             if final_usage_metadata:
-                usage_dict = final_usage_metadata.dict() if hasattr(final_usage_metadata, "dict") else {}
-                final_generation_info = {
-                    "token_usage": usage_dict,
-                    "model_name": self.model_name
-                }
-                yield ChatGenerationChunk(
-                    message=AIMessageChunk(content=""),
-                    generation_info=final_generation_info
-                )
+                lc_usage_metadata = self._map_usage_metadata(final_usage_metadata)
+                if lc_usage_metadata:
+                    yield ChatGenerationChunk(
+                        message=AIMessageChunk(content="", usage_metadata=lc_usage_metadata)
+                    )
                     
         except Exception as e:
             self.logger.error(f"Error during async streaming: {e}", exc_info=True)

crewplus/services/tracing_manager.py

@@ -0,0 +1,148 @@
+# File: crewplus/services/tracing_manager.py
+
+from typing import Any, Optional, List, Protocol
+import os
+import logging
+
+# Langfuse imports with graceful fallback. This allows the application to run
+# even if the langfuse library is not installed.
+try:
+    from langfuse.langchain import CallbackHandler as LangfuseCallbackHandler
+    LANGFUSE_AVAILABLE = True
+except ImportError:
+    LANGFUSE_AVAILABLE = False
+    LangfuseCallbackHandler = None
+
+class TracingContext(Protocol):
+    """
+    A protocol that defines a formal contract for a model to be "traceable."
+
+    This protocol ensures that any class using the TracingManager provides the
+    necessary attributes and methods for the manager to function correctly. By
+    using a Protocol, we leverage Python's static analysis tools (like mypy)
+    to enforce this contract, preventing runtime errors and making the system
+    more robust and self-documenting.
+
+    It allows the TracingManager to be completely decoupled from any specific
+    model implementation, promoting clean, compositional design.
+
+    A class that implements this protocol must provide:
+    - A `logger` attribute for logging.
+    - An `enable_tracing` attribute to control tracing.
+    - A `get_model_identifier` method to describe itself for logging purposes.
+    """
+    logger: logging.Logger
+    enable_tracing: Optional[bool]
+    
+    def get_model_identifier(self) -> str:
+        """
+        Return a string that uniquely identifies the model instance for logging.
+        
+        Example:
+            "GeminiChatModel (model='gemini-1.5-flash')"
+            
+        Note:
+            The '...' (Ellipsis) is the standard way in a Protocol to indicate
+            that this method must be implemented by any class that conforms to
+            this protocol, but has no implementation in the protocol itself.
+        """
+        ...
+
+class TracingManager:
+    """
+    Manages the initialization and injection of tracing handlers for chat models.
+    
+    This class uses a composition-based approach, taking a context object that
+    fulfills the TracingContext protocol. This design is highly extensible,
+    allowing new tracing providers (e.g., Helicone, OpenTelemetry) to be added
+    with minimal, isolated changes.
+    """
+    
+    def __init__(self, context: TracingContext):
+        """
+        Args:
+            context: An object (typically a chat model instance) that conforms
+                     to the TracingContext protocol.
+        """
+        self.context = context
+        self._handlers: List[Any] = []
+        self._initialize_handlers()
+    
+    def _initialize_handlers(self):
+        """
+        Initializes all supported tracing handlers. This is the central point
+        for adding new observability tools.
+        """
+        self._handlers = []
+        self._initialize_langfuse()
+        # To add a new handler (e.g., Helicone), you would add a call to
+        # self._initialize_helicone() here.
+    
+    def _initialize_langfuse(self):
+        """Initializes the Langfuse handler if it's available and enabled."""
+        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")
+            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
+            langfuse_env_vars = ["LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY"]
+            enable_langfuse = any(os.getenv(var) for var in langfuse_env_vars)
+        
+        if enable_langfuse:
+            try:
+                handler = LangfuseCallbackHandler()
+                self._handlers.append(handler)
+                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}")
+    
+    def add_callbacks_to_config(self, config: Optional[dict]) -> dict:
+        """
+        Adds all registered tracing handlers to the request configuration.
+        
+        This method is robust and handles three scenarios for the 'callbacks' key:
+        1. A list of callbacks.
+        2. A LangChain CallbackManager instance.
+        3. None or a missing key.
+        
+        Args:
+            config: The request configuration dictionary from a LangChain call.
+            
+        Returns:
+            The updated configuration dictionary with tracing callbacks added.
+        """
+        if config is None:
+            config = {}
+        
+        # Respect a global disable flag for this specific call.
+        if not self._handlers or config.get("metadata", {}).get("tracing_disabled"):
+            return config
+        
+        callbacks = config.get("callbacks")
+        
+        # Case 1: The 'callbacks' key holds a CallbackManager instance
+        if hasattr(callbacks, 'add_handler') and hasattr(callbacks, 'handlers'):
+            for handler in self._handlers:
+                if not any(isinstance(cb, type(handler)) for cb in callbacks.handlers):
+                    callbacks.add_handler(handler, inherit=True)
+            return config # Return the original, now-mutated config
+        
+        # Case 2: The 'callbacks' key holds a list or is None
+        current_callbacks = callbacks or []
+        new_callbacks = list(current_callbacks)
+        
+        for handler in self._handlers:
+            if not any(isinstance(cb, type(handler)) for cb in new_callbacks):
+                new_callbacks.append(handler)
+        
+        if len(new_callbacks) > len(current_callbacks):
+            # Create a new dictionary with the updated callbacks list.
+            # This is a safe operation that overwrites the existing 'callbacks'
+            # key and avoids mutating the original config object.
+            return {**config, "callbacks": new_callbacks}
+        
+        return config

{crewplus-0.2.25.dist-info → crewplus-0.2.27.dist-info}/METADATA

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

{crewplus-0.2.25.dist-info → crewplus-0.2.27.dist-info}/RECORD

@@ -1,13 +1,14 @@
-crewplus-0.2.25.dist-info/METADATA,sha256=NynftFjnRM1sFSQfe8PUWMRDuo1YOWyyyjgODaegtBs,4991
-crewplus-0.2.25.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.25.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.25.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.27.dist-info/METADATA,sha256=bYZxj9fhuzOqc7YDs-7Ef5msU_-ZXktY9ruuAsJjevM,4991
+crewplus-0.2.27.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.27.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.27.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/services/__init__.py,sha256=zUM4ZwUfGMBDx-j7Wehf_KC5yYXPTK8BK_oeO5veIXQ,398
-crewplus/services/azure_chat_model.py,sha256=xPuIsQpLV5Y3Ntwe3eqvquhBjh35g65VlF22AWJdEcU,8648
-crewplus/services/gemini_chat_model.py,sha256=HMDt7TKlLpQ43ZPxY9omG64EGFkP846BXT_SfyBeM0I,38415
+crewplus/services/azure_chat_model.py,sha256=WMSf4BDO8UcP7ZASNGRJxdTEnuWBmCRSY_4yx_VMbok,5499
+crewplus/services/gemini_chat_model.py,sha256=oVLL07VEjwrHOeP56YSPLnldbfjvTVRrsTI6xcPNt1E,35224
 crewplus/services/init_services.py,sha256=U91zoMNJlOEKyldarNnATjeZDT2V-0CrXPAwI64hZkw,758
 crewplus/services/model_load_balancer.py,sha256=HH_eHxFfxgarPWFGpANg7dgShnWca4q46Jz0b1vJ4Sw,9405
+crewplus/services/tracing_manager.py,sha256=aCU9N4Jvh8pDD3h8kWX4O-Ax8xwdLHnQ4wJ3sf-vLwA,6289
 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
@@ -15,8 +16,8 @@ crewplus/vectorstores/milvus/__init__.py,sha256=egGncAdjlXG6ekTQvKMKqhvKBifrUrPl
 crewplus/vectorstores/milvus/milvus_schema_manager.py,sha256=2IZT61LVui21Pt5Z3y8YYS2dYcwzkgUKxMq2NA0-lQE,9222
 crewplus/vectorstores/milvus/schema_milvus.py,sha256=IvKdUCH451HJ-F3TUR5jDjqwQlQs4SEXAQ_th4JAnfc,12117
 crewplus/vectorstores/milvus/vdb_service.py,sha256=wCltxZc0aD27iTu7wjveHqQWPEF2VyO4B2WGQCheeVs,21118
-docs/GeminiChatModel.md,sha256=_IQyup3ofAa2HxfSurO1GYUEezTHYYt5Q1khYNVThGM,8040
+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.25.dist-info/RECORD,,
+crewplus-0.2.27.dist-info/RECORD,,

docs/GeminiChatModel.md

@@ -61,9 +61,22 @@ response = model.invoke("Hello, how are you?")
 print(response.content)
 
 # Using stream for a chunked response
-print("\\n--- Streaming Response ---")
-for chunk in model.stream("Tell me a short story."):
+print("\n--- Streaming Response ---")
+for chunk in model.stream("Tell me a short story about a brave robot."):
     print(chunk.content, end="", flush=True)
+
+# Using astream for an asynchronous chunked response
+import asyncio
+
+async def main():
+    print("\n--- Async Streaming Response ---")
+    async for chunk in model.astream("Tell me a short story about a brave robot."):
+        print(chunk.content, end="", flush=True)
+
+# To run the async function in a Jupyter Notebook or a script:
+# await main()
+# Or, if not in an async context:
+# asyncio.run(main())
 ```
 
 ## 5. Image Understanding
@@ -120,9 +133,17 @@ try:
     print("Image response (base64):", image_response.content)
 except FileNotFoundError:
     print(f"Image file not found at {image_path}, skipping base64 example.")
+
+### Example 3: Streaming a Multimodal Response
+
+Streaming also works with complex, multimodal inputs. This is useful for getting faster time-to-first-token while the model processes all the data.
+
+```python
+# The url_message is from the previous example
+print("\n--- Streaming Multimodal Response ---")
+for chunk in model.stream([url_message]):
+    print(chunk.content, end="", flush=True)
 ```
-> **Sample Output:**
-> This image is a movie still from the 2017 Japanese thriller "22 Year Old's Confession: I am the Murderer"... The four women in the photo are the victims of a serial killer...
 
 ## 6. Video Understanding