crewplus 0.2.19__py3-none-any.whl → 0.2.22__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/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']
@@ -146,6 +147,7 @@ class ModelLoadBalancer:
                 openai_api_version=model_config['api_version'],
                 api_key=model_config['api_key'],
                 azure_endpoint=model_config['api_base'],
+                dimensions=model_config['dimensions'],
                 chunk_size=16, request_timeout=60, max_retries=2
             )
         elif provider == 'google-genai':

{crewplus-0.2.19.dist-info → crewplus-0.2.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.19
+Version: 0.2.22
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -12,9 +12,6 @@ Requires-Python: <4.0,>=3.11
 Requires-Dist: langchain==0.3.25
 Requires-Dist: langchain-openai==0.3.24
 Requires-Dist: google-genai==1.21.1
-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-0.2.19.dist-info → crewplus-0.2.22.dist-info}/RECORD

@@ -1,12 +1,13 @@
-crewplus-0.2.19.dist-info/METADATA,sha256=7Y8TPTHtxcJZFMF5GJgeCIt5JK_bJMxeYR0p7Z_QRDY,5125
-crewplus-0.2.19.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.19.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.19.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.22.dist-info/METADATA,sha256=0Mhxz-ArK28dpPEBsh835Tw2IWxEXsNEH6uLwbrMJkw,4991
+crewplus-0.2.22.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.22.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.22.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/__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=OvWgyGurmhWGRZY3tpAHbZ3ab_FSaR7tEKTmT53ikCI,8847
 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.19.dist-info/RECORD,,
+crewplus-0.2.22.dist-info/RECORD,,