crewplus 0.2.89__py3-none-any.whl

crewplus/services/init_services.py

@@ -0,0 +1,57 @@
+import logging
+import os
+from typing import Optional
+
+from .model_load_balancer import ModelLoadBalancer
+
+model_balancer = None
+
+def init_load_balancer(
+    config_path: Optional[str] = None, 
+    logger: Optional[logging.Logger] = None
+):
+    """
+    Initializes the global ModelLoadBalancer instance.
+
+    This function is idempotent. If the balancer is already initialized,
+    it does nothing. It follows a safe initialization pattern where the
+    global instance is only assigned after successful configuration loading.
+    
+    Args:
+        config_path (Optional[str]): The path to the model configuration file.
+            If not provided, it's determined by the `MODEL_CONFIG_PATH`
+            environment variable, or defaults to "config/models_config.json".
+        logger (Optional[logging.Logger]): An optional logger instance to be
+            used by the model balancer.
+    """
+    global model_balancer
+    if model_balancer is None:
+        # Use parameter if provided, otherwise check env var, then default
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        base_package_dir = os.path.dirname(os.path.dirname(current_dir))
+        default_config_path = os.path.join(base_package_dir, "_config", "models_config.json")
+
+        final_config_path = config_path or os.getenv(
+            "MODEL_CONFIG_PATH",
+            default_config_path
+        )
+        try:
+            # 1. Create a local instance first.
+            balancer = ModelLoadBalancer(
+                config_path=final_config_path, 
+                logger=logger
+            )
+            # 2. Attempt to load its configuration.
+            balancer.load_config()
+            # 3. Only assign to the global variable on full success.
+            model_balancer = balancer
+        except Exception as e:
+            # If any step fails, the global model_balancer remains None,
+            # allowing for another initialization attempt later.
+            # Re-raise the exception to notify the caller of the failure.
+            raise RuntimeError(f"Failed to initialize and configure ModelLoadBalancer from {final_config_path}: {e}") from e
+
+def get_model_balancer() -> ModelLoadBalancer:
+    if model_balancer is None:
+        raise RuntimeError("ModelLoadBalancer not initialized. Please call init_load_balancer() first.")
+    return model_balancer

crewplus/services/model_load_balancer.py

@@ -0,0 +1,264 @@
+import json
+import random
+import logging
+import threading
+from typing import Dict, List, Optional, Union
+from collections import defaultdict
+from langchain_openai import ChatOpenAI, AzureOpenAIEmbeddings
+from .gemini_chat_model import GeminiChatModel
+from .azure_chat_model import TracedAzureChatOpenAI
+
+
+class ModelLoadBalancer:
+    def __init__(self,
+                 config_path: Optional[str] = "config/models_config.json",
+                 config_data: Optional[Dict] = None,
+                 logger: Optional[logging.Logger] = None):
+        """
+        Initializes the ModelLoadBalancer.
+
+        Args:
+            config_path: Path to the JSON configuration file.
+            config_data: A dictionary containing the model configuration.
+            logger: An optional logger instance. If not provided, a default one is created.
+
+        Raises:
+            ValueError: If neither config_path nor config_data is provided.
+        """
+        if not config_path and not config_data:
+            raise ValueError("Either 'config_path' or 'config_data' must be provided.")
+
+        self.config_path = config_path
+        self.config_data = config_data
+        self.logger = logger or logging.getLogger(__name__)
+        self.models_config: List[Dict] = []
+        self.thread_local = threading.local()
+        self._initialize_state()
+        self._config_loaded = False  # Flag to check if config is loaded
+
+    def load_config(self):
+        """Load and validate model configurations from a file path or a dictionary."""
+        self.logger.debug("Model balancer: loading configuration.")
+        try:
+            config = None
+            if self.config_data:
+                config = self.config_data
+            elif self.config_path:
+                with open(self.config_path, 'r') as f:
+                    config = json.load(f)
+            else:
+                # This case is handled in __init__, but as a safeguard:
+                raise RuntimeError("No configuration source provided (path or data).")
+
+            # Validate config
+            if 'models' not in config or not isinstance(config['models'], list):
+                raise ValueError("Configuration must contain a 'models' list.")
+
+            for model in config.get('models', []):
+                if 'provider' not in model or 'type' not in model or 'id' not in model:
+                    self.logger.error("Model config must contain 'id', 'provider', and 'type' fields.")
+                    raise ValueError("Model config must contain 'id', 'provider', and 'type' fields.")
+
+            self.models_config = config['models']
+
+            self._config_loaded = True
+            self.logger.debug("Model balancer: configuration loaded successfully.")
+        except (FileNotFoundError, json.JSONDecodeError, ValueError) as e:
+            self._config_loaded = False
+            self.logger.error(f"Failed to load model configuration: {e}", exc_info=True)
+            raise RuntimeError(f"Failed to load model configuration: {e}")
+
+    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None, with_metadata: bool = False, selection_strategy: str = 'random', disable_streaming: bool = False):
+        """
+        Get a model instance.
+
+        Can fetch a model in two ways:
+        1. By its specific `deployment_name`.
+        2. By `provider` and `model_type`, which will select a model using a specified strategy.
+
+        Args:
+            provider: The model provider (e.g., 'azure-openai', 'google-genai').
+            model_type: The type of model (e.g., 'inference', 'embedding', 'embedding-large').
+            deployment_name: The unique name for the model deployment.
+            with_metadata: If True, returns a tuple of (model, deployment_name).
+            selection_strategy: The selection strategy ('random', 'round_robin', or 'least_used'). Defaults to 'random'.
+            disable_streaming: If True, get a model instance with streaming disabled.
+
+        Returns:
+            An instantiated language model object, or a tuple if with_metadata is True.
+
+        Raises:
+            RuntimeError: If the model configuration has not been loaded.
+            ValueError: If the requested model cannot be found or if parameters are insufficient.
+        """
+        if not self._config_loaded:
+            self.logger.error("Model configuration not loaded")
+            raise RuntimeError("Model configuration not loaded")
+
+        if deployment_name:
+            for model_config in self.models_config:
+                if model_config.get('deployment_name') == deployment_name:
+                    model = self._get_or_create_model(model_config, disable_streaming)
+                    if with_metadata:
+                        return model, deployment_name
+                    return model
+                
+            self.logger.error(f"No model found for deployment name: {deployment_name}")
+            raise ValueError(f"No model found for deployment name: {deployment_name}")
+
+        if provider and model_type:
+            candidates = [model for model in self.models_config if model.get('provider') == provider and model.get('type') == model_type]
+            if not candidates:
+                self.logger.error(f"No models found for provider '{provider}' and type '{model_type}'")
+                raise ValueError(f"No models found for provider '{provider}' and type '{model_type}'")
+
+            if selection_strategy == 'random':
+                selected_model_config = self._random_selection(candidates)
+            elif selection_strategy == 'round_robin':
+                selected_model_config = self._round_robin_selection(candidates)
+            elif selection_strategy == 'least_used':
+                selected_model_config = self._least_used_selection(candidates)
+            else:
+                self.logger.warning(f"Unsupported selection strategy: '{selection_strategy}'. Defaulting to 'random'.")
+                selected_model_config = self._random_selection(candidates)
+                
+            model = self._get_or_create_model(selected_model_config, disable_streaming)
+            if with_metadata:
+                return model, selected_model_config.get('deployment_name')
+            return model
+
+        raise ValueError("Either 'deployment_name' or both 'provider' and 'model_type' must be provided.")
+
+    def _get_thread_local_models_cache(self) -> Dict:
+        """Gets the model cache for the current thread, creating it if it doesn't exist."""
+        if not hasattr(self.thread_local, 'models_cache'):
+            self.thread_local.models_cache = {}
+        return self.thread_local.models_cache
+
+    def _get_or_create_model(self, model_config: Dict, disable_streaming: bool = False):
+        """
+        Gets a model instance from the thread-local cache. If it doesn't exist,
+        it instantiates, caches, and returns it.
+        """
+        model_id = model_config['id']
+        cache_key = f"{model_id}"
+        if disable_streaming:
+            cache_key += "-non-streaming"
+            
+        models_cache = self._get_thread_local_models_cache()
+
+        if cache_key not in models_cache:
+            self.logger.debug(f"Creating new model instance for id {cache_key} in thread {threading.get_ident()}")
+            models_cache[cache_key] = self._instantiate_model(model_config, disable_streaming)
+        
+        return models_cache[cache_key]
+
+    def _instantiate_model(self, model_config: Dict, disable_streaming: bool = False):
+        """Instantiate and return an LLM object based on the model configuration"""
+        provider = model_config['provider']
+        self.logger.debug(f"Model balancer: instantiating {provider} -- {model_config.get('deployment_name')}")
+
+        if provider == 'azure-openai':
+            kwargs = {
+                'azure_deployment': model_config['deployment_name'],
+                'openai_api_version': model_config['api_version'],
+                'azure_endpoint': model_config['api_base'],
+                'openai_api_key': model_config['api_key']
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            
+            # The 'disable_streaming' parameter takes precedence
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
+            elif model_config.get('deployment_name') == 'o1-mini':
+                kwargs['disable_streaming'] = True
+                
+            return TracedAzureChatOpenAI(**kwargs)
+        elif provider == 'openai':
+            kwargs = {
+                'openai_api_key': model_config['api_key']
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            return ChatOpenAI(**kwargs)
+        elif provider == 'azure-openai-embeddings':
+            try:
+                emb_kwargs = dict(
+                    model=model_config['model_name'],
+                    azure_deployment=model_config['deployment_name'],
+                    openai_api_version=model_config['api_version'],
+                    api_key=model_config['api_key'],
+                    azure_endpoint=model_config['api_base'],
+                    chunk_size=16, request_timeout=60, max_retries=2
+                )
+                if 'dimensions' in model_config:
+                    emb_kwargs['dimensions'] = model_config['dimensions']
+                return AzureOpenAIEmbeddings(**emb_kwargs)
+            except Exception as e:
+                self.logger.error(f"Failed to instantiate AzureOpenAIEmbeddings: {e}")
+                return None
+        elif provider == 'google-genai':
+            kwargs = {
+                'google_api_key': model_config['api_key'],
+                'model_name': model_config['deployment_name']  # Map deployment_name to model_name
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            if 'max_tokens' in model_config:
+                kwargs['max_tokens'] = model_config['max_tokens']
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
+            return GeminiChatModel(**kwargs)
+        elif provider == 'vertex-ai':
+            deployment_name = model_config['deployment_name']
+            
+            # Handle the 'model_name@location' format for deployment_name
+            model_name_for_gemini = deployment_name.split('@')[0] if '@' in deployment_name else deployment_name
+
+            kwargs = {
+                'use_vertex_ai': True,
+                'model_name': model_name_for_gemini,
+                'project_id': model_config['project_id'],
+                'location': model_config['location'],
+            }
+            if 'service_account_file' in model_config:
+                kwargs['service_account_file'] = model_config['service_account_file']
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            if 'max_tokens' in model_config:
+                kwargs['max_tokens'] = model_config['max_tokens']
+            if disable_streaming:
+                kwargs['disable_streaming'] = True
+            return GeminiChatModel(**kwargs)
+        else:
+            self.logger.error(f"Unsupported provider: {provider}")
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    def _initialize_state(self):
+        self.active_models = []
+        self.usage_counter = defaultdict(int)
+        self.current_indices = {}
+
+    def _random_selection(self, candidates: list) -> Dict:
+        """Selects a model randomly from a list of candidates."""
+        model = random.choice(candidates)
+        self.usage_counter[model['id']] += 1
+        return model
+
+    def _round_robin_selection(self, candidates: list) -> Dict:
+        if id(candidates) not in self.current_indices:
+            self.current_indices[id(candidates)] = 0
+        idx = self.current_indices[id(candidates)]
+        model = candidates[idx]
+        self.current_indices[id(candidates)] = (idx + 1) % len(candidates)
+        self.usage_counter[model['id']] += 1
+
+        return model
+
+    def _least_used_selection(self, candidates: list) -> Dict:
+        min_usage = min(self.usage_counter[m['id']] for m in candidates)
+        least_used = [m for m in candidates if self.usage_counter[m['id']] == min_usage]
+        model = random.choice(least_used)
+        self.usage_counter[model['id']] += 1
+        return model

crewplus/services/schemas/feedback.py

@@ -0,0 +1,61 @@
+from typing import Optional, Union, Dict, Any
+from uuid import UUID
+from pydantic import BaseModel, Field, ConfigDict
+import datetime
+
+class FeedbackConfig(BaseModel):
+    api_url: Optional[str] = None
+    api_key: Optional[str] = None
+
+class FeedbackIn(BaseModel):
+    run_id: Union[UUID, str]
+    key: str
+    score: Optional[Union[float, int, bool]] = None
+    value: Optional[Union[float, int, bool, str, Dict]] = None
+    correction: Optional[Dict] = None
+    comment: Optional[str] = None
+    feedback_group_id: Optional[Union[UUID, str]] = None
+    config: Optional[FeedbackConfig] = None
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "run_id": "5323d917-f337-42c2-8437-22e6fa623930",
+                "key": "accuracy",
+                "score": 0.95,
+                "value": "yes",
+                "correction": {"expected": "correct answer"},
+                "comment": "The model performed well on this task.",
+                "feedback_group_id": "123e4567-e89b-12d3-a456-426614174000",
+                "config": {
+                    "api_url": "https://api.smith.langchain.com",
+                    "api_key": "your_api_key_here"
+                }
+            }
+        }
+    )
+
+class FeedbackUpdate(BaseModel):
+    score: Optional[Union[float, int, bool]] = None
+    value: Optional[Union[float, int, bool, str, Dict]] = None
+    correction: Optional[Dict] = None
+    comment: Optional[str] = None
+    config: Optional[FeedbackConfig] = None
+
+
+class FeedbackOut(BaseModel):
+    """
+    Represents a retrieved feedback item, including its unique ID and group context.
+    This corresponds to a Langfuse Score object enriched with its parent's metadata.
+    """
+    feedback_id: str = Field(..., description="The unique identifier for the feedback item (maps to Langfuse score ID).")
+    run_id: str = Field(..., description="The ID of the run (trace) this feedback belongs to.")
+    key: str = Field(..., description="The name of the feedback key.")
+    score: Optional[Union[float, int, bool, str]] = Field(None, description="The numerical or string score of the feedback.")
+    value: Optional[Any] = Field(None, description="The original value of the feedback, which could be a dict or other type.")
+    comment: Optional[str] = Field(None, description="The feedback comment.")
+    correction: Optional[Dict] = Field(None, description="Correction data from the parent observation.")
+    feedback_group_id: Optional[str] = Field(None, description="The ID of the group this feedback belongs to.")
+    created_at: datetime.datetime = Field(..., description="The timestamp when the feedback was created.")
+
+    model_config = ConfigDict(from_attributes=True)

crewplus/services/tracing_manager.py

@@ -0,0 +1,182 @@
+# File: crewplus/services/tracing_manager.py
+
+from typing import Any, Optional, List, Protocol, Dict
+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
+    from ..callbacks.async_langfuse_handler import AsyncLangfuseCallbackHandler
+    from ..utils.tracing_util import get_langfuse_handler, get_async_langfuse_handler
+    LANGFUSE_AVAILABLE = True
+except ImportError:
+    LANGFUSE_AVAILABLE = False
+    LangfuseCallbackHandler = None
+    AsyncLangfuseCallbackHandler = None
+    get_langfuse_handler = None
+    get_async_langfuse_handler = 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._sync_handlers: List[Any] = []
+        self._async_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._sync_handlers = []
+        self._async_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."""
+        self.context.logger.debug("Attempting to initialize Langfuse handlers.")
+        if not LANGFUSE_AVAILABLE:
+            if self.context.enable_tracing is True:
+                self.context.logger.warning("Langfuse is not installed; tracing will be disabled. Install with: pip install langfuse")
+            else:
+                self.context.logger.debug("Langfuse is not installed, skipping handler initialization.")
+            return
+        
+        # Determine if Langfuse should be enabled via an explicit flag or
+        # by detecting its environment variables.
+        enable_langfuse = self.context.enable_tracing
+        if enable_langfuse is None: # Auto-detect if not explicitly set
+            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:
+                # Create both sync and async handlers. We'll pick one at runtime.
+                sync_handler = get_langfuse_handler()
+                self._sync_handlers.append(sync_handler)
+
+                if AsyncLangfuseCallbackHandler:
+                    async_handler = get_async_langfuse_handler()
+                    self._async_handlers.append(async_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}", exc_info=True)
+        else:
+            self.context.logger.info("Langfuse is not enabled, skipping handler initialization.")
+    
+    def add_callbacks_to_config(self, config: Optional[dict], handlers: List[Any]) -> dict:
+        """A generic helper to add a list of handlers to a config object."""
+        if config is None:
+            config = {}
+        
+        self.context.logger.debug(f"Adding callbacks to config. Have {len(handlers)} handlers to add.")
+
+        if not handlers or config.get("metadata", {}).get("tracing_disabled"):
+            self.context.logger.debug("No handlers to add or tracing is disabled for this run.")
+            return config
+        
+        callbacks = config.get("callbacks")
+        
+        if hasattr(callbacks, 'add_handler') and hasattr(callbacks, 'handlers'):
+            # This block is for CallbackManager instances
+            self.context.logger.debug(f"Config has a CallbackManager with {len(callbacks.handlers)} existing handlers.")
+            for handler in handlers:
+                if not any(isinstance(cb, type(handler)) for cb in callbacks.handlers):
+                    callbacks.add_handler(handler, inherit=True)
+            self.context.logger.debug(f"CallbackManager now has {len(callbacks.handlers)} handlers.")
+            return config
+        
+        # This block is for simple lists of callbacks
+        current_callbacks = callbacks or []
+        self.context.logger.debug(f"Config has a list with {len(current_callbacks)} existing callbacks.")
+        new_callbacks = list(current_callbacks)
+        
+        for handler in handlers:
+            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
+
+    def add_sync_callbacks_to_config(self, config: Optional[dict], handlers: Optional[List[Any]] = None) -> dict:
+        """
+        Adds synchronous tracing handlers to the request configuration.
+
+        Args:
+            config: The configuration dictionary to which callbacks will be added.
+            handlers: An optional list of handlers to add. If not provided,
+                      the manager's default synchronous handlers are used.
+        """
+        handlers_to_add = self._sync_handlers if handlers is None else handlers
+        return self.add_callbacks_to_config(config, handlers_to_add)
+
+    def add_async_callbacks_to_config(self, config: Optional[dict], handlers: Optional[List[Any]] = None) -> dict:
+        """
+        Adds asynchronous tracing handlers to the request configuration.
+
+        Args:
+            config: The configuration dictionary to which callbacks will be added.
+            handlers: An optional list of handlers to add. If not provided,
+                      the manager's default asynchronous handlers are used.
+        """
+        handlers_to_add = self._async_handlers if handlers is None else handlers
+        return self.add_callbacks_to_config(config, handlers_to_add)

crewplus/utils/__init__.py

@@ -0,0 +1,4 @@
+from .schema_action import Action
+from .schema_document_updater import SchemaDocumentUpdater
+
+__all__ = ["Action", "SchemaDocumentUpdater"]

crewplus/utils/schema_action.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+class Action(Enum):
+    UPSERT = "upsert"  # Update existing fields; if a match is found, it updates, otherwise, it inserts. Does not delete unmatched existing fields.
+    DELETE = "delete"  # Clear data from fields in the schema.
+    UPDATE = "update"  # Update only the matching original fields.
+    INSERT = "insert"  # Insert data, clearing the original fields before inserting new values.