crewplus 0.2.15__py3-none-any.whl → 0.2.21__py3-none-any.whl

crewplus/services/__init__.py

@@ -1,6 +1,13 @@
 from .gemini_chat_model import GeminiChatModel
 from .init_services import init_load_balancer, get_model_balancer
 from .model_load_balancer import ModelLoadBalancer
+from .azure_chat_model import TracedAzureChatOpenAI
 
-
-__all__ = ["GeminiChatModel", "init_load_balancer", "get_model_balancer", "ModelLoadBalancer"]
+__all__ = [
+    "GeminiChatModel", 
+    "init_load_balancer", 
+    "get_model_balancer", 
+    "ModelLoadBalancer", 
+    "init_services", 
+    "TracedAzureChatOpenAI"
+]

crewplus/services/azure_chat_model.py

@@ -0,0 +1,201 @@
+import os
+import logging
+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
+
+class TracedAzureChatOpenAI(AzureChatOpenAI):
+    """
+    Wrapper for AzureChatOpenAI that integrates with Langfuse for tracing.
+    
+    This class automatically handles Langfuse 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:
+    - 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.
+
+    Attributes:
+        logger (Optional[logging.Logger]): An optional logger instance.
+        enable_langfuse (Optional[bool]): Enable/disable Langfuse tracing (auto-detect if None).
+
+    Example:
+        .. code-block:: python
+
+            # Set Langfuse environment variables (optional)
+            import os
+            os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-..."
+            os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-..."
+
+            from crewplus.services.azure_chat_model import TracedAzureChatOpenAI
+            from langchain_core.messages import HumanMessage
+
+            # Initialize the model
+            model = TracedAzureChatOpenAI(
+                azure_deployment="your-deployment",
+                api_version="2024-05-01-preview",
+            )
+
+            # --- Text-only usage (automatically traced if env vars set) ---
+            response = model.invoke("Hello, how are you?")
+            print("Text response:", response.content)
+
+            # --- Langfuse tracing with session/user tracking ---
+            response = model.invoke(
+                "What is AI?",
+                config={
+                    "metadata": {
+                        "langfuse_session_id": "chat-session-123",
+                        "langfuse_user_id": "user-456"
+                    }
+                }
+            )
+            
+            # --- Disable Langfuse for specific calls ---
+            response = model.invoke(
+                "Hello without tracing",
+                config={"metadata": {"langfuse_disabled": True}}
+            )
+
+            # --- Asynchronous Streaming Usage ---
+            import asyncio
+            from langchain_core.messages import HumanMessage
+
+            async def main():
+                messages = [HumanMessage(content="Tell me a short story about a brave robot.")]
+                print("\nAsync Streaming response:")
+                async for chunk in model.astream(messages):
+                    print(chunk.content, end="", flush=True)
+                print()
+
+            # In a real application, you would run this with:
+            # 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)")
+    
+    langfuse_handler: Optional[LangfuseCallbackHandler] = Field(default=None, exclude=True)
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        
+        # Initialize logger
+        if self.logger is None:
+            self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
+            if not self.logger.handlers:
+                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
+
+    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)
+
+            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}
+        
+        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}
+        
+        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]
+        
+        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]
+
+        async for chunk in super().astream(input, config=config, **kwargs):
+            yield chunk

crewplus/services/gemini_chat_model.py

@@ -22,6 +22,14 @@ 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
+
 class GeminiChatModel(BaseChatModel):
     """Custom chat model for Google Gemini, supporting text, image, and video.
     
@@ -33,6 +41,15 @@ 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:
+    - 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.
+
     Attributes:
         model_name (str): The Google model name to use (e.g., "gemini-1.5-flash").
         google_api_key (Optional[SecretStr]): Your Google API key.
@@ -41,10 +58,18 @@ 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).
 
     Example:
         .. code-block:: python
 
+            # Set Langfuse environment variables (optional)
+            import os
+            os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-..."
+            os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-..."
+            os.environ["LANGFUSE_HOST"] = "https://cloud.langfuse.com"  # EU region or self-hosted
+            # os.environ["LANGFUSE_HOST"] = "https://us.cloud.langfuse.com"  # US region
+
             from crewplus.services import GeminiChatModel
             from langchain_core.messages import HumanMessage
             import base64
@@ -54,10 +79,21 @@ class GeminiChatModel(BaseChatModel):
             logger = logging.getLogger("my_app.gemini")
             model = GeminiChatModel(model_name="gemini-2.0-flash", logger=logger)
 
-            # --- Text-only usage ---
+            # --- Text-only usage (automatically traced if env vars set) ---
             response = model.invoke("Hello, how are you?")
             print("Text response:", response.content)
 
+            # --- Langfuse tracing with session/user tracking ---
+            response = model.invoke(
+                "What is AI?",
+                config={
+                    "metadata": {
+                        "langfuse_session_id": "chat-session-123",
+                        "langfuse_user_id": "user-456"
+                    }
+                }
+            )
+
             # --- Image processing with base64 data URI ---
             # Replace with a path to your image
             image_path = "path/to/your/image.jpg"
@@ -138,6 +174,23 @@ class GeminiChatModel(BaseChatModel):
             print("Streaming response:")
             for chunk in model.stream([url_message]):
                 print(chunk.content, end="", flush=True)
+
+            # --- Traditional Langfuse callback approach still works ---
+            from langfuse.langchain import CallbackHandler
+            langfuse_handler = CallbackHandler(
+                session_id="session-123",
+                user_id="user-456"
+            )
+            response = model.invoke(
+                "Hello with manual callback",
+                config={"callbacks": [langfuse_handler]}
+            )
+
+            # --- Disable Langfuse for specific calls ---
+            response = model.invoke(
+                "Hello without tracing",
+                config={"metadata": {"langfuse_disabled": True}}
+            )
     """
     
     # Model configuration
@@ -149,8 +202,12 @@ class GeminiChatModel(BaseChatModel):
     top_k: Optional[int] = Field(default=None, description="Top-k sampling parameter")
     logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance")
     
-    # Internal client
+    # Langfuse configuration
+    enable_langfuse: Optional[bool] = Field(default=None, description="Enable Langfuse tracing (auto-detect if None)")
+    
+    # Internal clients
     _client: Optional[genai.Client] = None
+    _langfuse_handler: Optional[LangfuseCallbackHandler] = None
     
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -178,6 +235,135 @@ class GeminiChatModel(BaseChatModel):
             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
+
+    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}
+        
+        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}
+        
+        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}
+        
+        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)
     
     @property
     def _llm_type(self) -> str:
@@ -461,6 +647,66 @@ class GeminiChatModel(BaseChatModel):
             }
         return part_dict
 
+    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:
+            self.logger.debug(f"[_extract_usage_metadata] Found usage_metadata: {response.usage_metadata}")
+            return response.usage_metadata
+        return None
+
+    def _create_chat_generation_chunk(self, chunk_response) -> ChatGenerationChunk:
+        """Creates a ChatGenerationChunk for streaming."""
+        # For streaming, we do not include usage metadata in individual chunks
+        # to prevent merge conflicts. The final, aggregated response will contain
+        # the full usage details for callbacks like Langfuse.
+        return ChatGenerationChunk(
+            message=AIMessageChunk(
+                content=chunk_response.text,
+                response_metadata={"model_name": self.model_name},
+            ),
+            generation_info=None,
+        )
+
+    def _create_chat_result_with_usage(self, response) -> ChatResult:
+        """Creates a ChatResult with usage metadata for Langfuse tracking."""
+        generated_text = response.text
+        finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
+        
+        # Extract usage metadata for token tracking
+        usage_metadata = self._extract_usage_metadata(response)
+        usage_dict = usage_metadata.dict() if usage_metadata and hasattr(usage_metadata, "dict") else {}
+
+        # Create AIMessage with usage information in response_metadata
+        message = AIMessage(
+            content=generated_text,
+            response_metadata={
+                "model_name": self.model_name,
+                "finish_reason": finish_reason,
+                **usage_dict
+            }
+        )
+        
+        # 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
+        )
+        
+        # We also construct the llm_output dictionary in the format expected
+        # by LangChain callback handlers, with a specific "token_usage" key.
+        chat_result = ChatResult(
+            generations=[generation],
+            llm_output={
+                "token_usage": usage_dict,
+                "model_name": self.model_name
+            } if usage_dict else {
+                "model_name": self.model_name
+            }
+        )
+        
+        return chat_result
+
     def _generate(
         self,
         messages: List[BaseMessage],
@@ -471,6 +717,8 @@ class GeminiChatModel(BaseChatModel):
         """Generates a chat response from a list of messages."""
         self.logger.info(f"Generating response for {len(messages)} messages.")
         
+        # Remove the problematic add_handler call - callbacks are now handled in invoke methods
+        
         contents = self._convert_messages(messages)
         config = self._prepare_generation_config(messages, stop)
 
@@ -482,14 +730,7 @@ class GeminiChatModel(BaseChatModel):
                 **kwargs,
             )
             
-            generated_text = response.text
-            finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
-            
-            message = AIMessage(
-                content=generated_text,
-                response_metadata={"model_name": self.model_name, "finish_reason": finish_reason},
-            )
-            return ChatResult(generations=[ChatGeneration(message=message)])
+            return self._create_chat_result_with_usage(response)
             
         except Exception as e:
             self.logger.error(f"Error generating content with Google GenAI: {e}", exc_info=True)
@@ -516,14 +757,7 @@ class GeminiChatModel(BaseChatModel):
                 **kwargs,
             )
             
-            generated_text = response.text
-            finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
-            
-            message = AIMessage(
-                content=generated_text,
-                response_metadata={"model_name": self.model_name, "finish_reason": finish_reason},
-            )
-            return ChatResult(generations=[ChatGeneration(message=message)])
+            return self._create_chat_result_with_usage(response)
 
         except Exception as e:
             self.logger.error(f"Error during async generation: {e}", exc_info=True)
@@ -536,7 +770,7 @@ class GeminiChatModel(BaseChatModel):
         run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> Iterator[ChatGenerationChunk]:
-        """Streams the chat response."""
+        """Streams the chat response and properly handles final usage metadata."""
         self.logger.info(f"Streaming response for {len(messages)} messages.")
         
         contents = self._convert_messages(messages)
@@ -549,12 +783,35 @@ class GeminiChatModel(BaseChatModel):
                 config=config,
                 **kwargs,
             )
+            
+            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 = ChatGenerationChunk(message=AIMessageChunk(content=text_content))
+                    chunk = self._create_chat_generation_chunk(chunk_response)
                     if run_manager:
                         run_manager.on_llm_new_token(text_content, chunk=chunk)
                     yield chunk
+            
+            # 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.
+            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
+                )
+
         except Exception as e:
             self.logger.error(f"Error streaming content: {e}", exc_info=True)
             raise ValueError(f"Error during streaming: {e}")
@@ -566,7 +823,7 @@ class GeminiChatModel(BaseChatModel):
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
-        """Asynchronously streams the chat response."""
+        """Asynchronously streams the chat response and properly handles final usage metadata."""
         self.logger.info(f"Async streaming response for {len(messages)} messages.")
         
         contents = self._convert_messages(messages)
@@ -579,12 +836,35 @@ class GeminiChatModel(BaseChatModel):
                 config=config,
                 **kwargs,
             )
+            
+            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 = ChatGenerationChunk(message=AIMessageChunk(content=text_content))
+                    chunk = self._create_chat_generation_chunk(chunk_response)
                     if run_manager:
                         await run_manager.on_llm_new_token(text_content, chunk=chunk)
                     yield chunk
+            
+            # 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.
+            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
+                )
+                    
         except Exception as e:
             self.logger.error(f"Error during async streaming: {e}", exc_info=True)
             raise ValueError(f"Error during async streaming: {e}")

crewplus/services/model_load_balancer.py

@@ -3,8 +3,9 @@ import random
 import logging
 from typing import Dict, List, Optional, Union
 from collections import defaultdict
-from langchain_openai import AzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings
+from langchain_openai import ChatOpenAI, AzureOpenAIEmbeddings
 from .gemini_chat_model import GeminiChatModel
+from .azure_chat_model import TracedAzureChatOpenAI
 
 
 class ModelLoadBalancer:
@@ -30,7 +31,7 @@ class ModelLoadBalancer:
         self.config_data = config_data
         self.logger = logger or logging.getLogger(__name__)
         self.models_config: List[Dict] = []
-        self.models: Dict[int, Union[AzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings, GeminiChatModel]] = {}
+        self.models: Dict[int, Union[TracedAzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings, GeminiChatModel]] = {}
         self._initialize_state()
         self._config_loaded = False  # Flag to check if config is loaded
 
@@ -131,7 +132,7 @@ class ModelLoadBalancer:
                 kwargs['temperature'] = model_config['temperature']
             if model_config.get('deployment_name') == 'o1-mini':
                 kwargs['disable_streaming'] = True
-            return AzureChatOpenAI(**kwargs)
+            return TracedAzureChatOpenAI(**kwargs)
         elif provider == 'openai':
             kwargs = {
                 'openai_api_key': model_config['api_key']

{crewplus-0.2.15.dist-info → crewplus-0.2.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.15
+Version: 0.2.21
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -16,6 +16,7 @@ Requires-Dist: mkdocs<2.0.0,>=1.6.1
 Requires-Dist: mkdocs-material<10.0.0,>=9.6.14
 Requires-Dist: mkdocstrings-python<2.0.0,>=1.16.12
 Requires-Dist: langchain-milvus<0.3.0,>=0.2.1
+Requires-Dist: langfuse<4.0.0,>=3.1.3
 Description-Content-Type: text/markdown
 
 # CrewPlus

{crewplus-0.2.15.dist-info → crewplus-0.2.21.dist-info}/RECORD

@@ -1,12 +1,13 @@
-crewplus-0.2.15.dist-info/METADATA,sha256=UsIzsZEpV38aiw_SP-wLsjA3bcEYf4Q3lp3oEzcLSaI,5087
-crewplus-0.2.15.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.15.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.15.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.21.dist-info/METADATA,sha256=_cUtSY9ZifWXNKbJ7j3zQjP0gqKY5RPvExND8z-XOB0,5125
+crewplus-0.2.21.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.21.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.21.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
-crewplus/services/__init__.py,sha256=ra_ciHcJN_sbv7q8UCP2kY91SbD32-QBpQLRgIosEcE,267
-crewplus/services/gemini_chat_model.py,sha256=i9p5KvSJYaHSUBLPKM_bpyGVLWCDQoNeah_WjQVJRXs,26227
+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/init_services.py,sha256=U91zoMNJlOEKyldarNnATjeZDT2V-0CrXPAwI64hZkw,758
-crewplus/services/model_load_balancer.py,sha256=AKwgxZauXGcSQvFxYQT6tR57ZSYRVnWWZvLWAWTOewU,8743
+crewplus/services/model_load_balancer.py,sha256=6JvmqmHz52KmVdBqF8nt45rjwymUCiF6fqSDFZcQuJ0,8791
 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
@@ -18,4 +19,4 @@ docs/GeminiChatModel.md,sha256=_IQyup3ofAa2HxfSurO1GYUEezTHYYt5Q1khYNVThGM,8040
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.15.dist-info/RECORD,,
+crewplus-0.2.21.dist-info/RECORD,,