crewplus 0.2.15__tar.gz → 0.2.19__tar.gz

{crewplus-0.2.15 → crewplus-0.2.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.15
+Version: 0.2.19
 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 → crewplus-0.2.19}/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-0.2.15 → crewplus-0.2.19}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "crewplus"
-version = "0.2.15"
+version = "0.2.19"
 description = "Base services for CrewPlus AI applications"
 authors = [
     { name = "Tim Liu", email = "tim@opsmateai.com" },
@@ -21,6 +21,7 @@ dependencies = [
     "mkdocs-material (>=9.6.14,<10.0.0)",
     "mkdocstrings-python (>=1.16.12,<2.0.0)",
     "langchain-milvus (>=0.2.1,<0.3.0)",
+    "langfuse (>=3.1.3,<4.0.0)",
 ]
 
 [project.license]