aiecs 1.0.0__py3-none-any.whl

aiecs/llm/custom_callbacks.py

@@ -0,0 +1,228 @@
+from typing import Any, List, Optional
+import logging
+
+# Import the base callback handler from utils
+from ..utils.base_callback import CustomAsyncCallbackHandler
+# Import LLM types for internal use only
+from .base_client import LLMMessage, LLMResponse
+# Import token usage repository
+from ..utils.token_usage_repository import token_usage_repo
+
+logger = logging.getLogger(__name__)
+
+
+class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Concrete token recording callback handler.
+    Responsible for recording token usage after LLM calls by delegating to the repository.
+    """
+
+    def __init__(self, user_id: str, cycle_start_date: Optional[str] = None):
+        if not user_id:
+            raise ValueError("user_id must be provided for RedisTokenCallbackHandler")
+        self.user_id = user_id
+        self.cycle_start_date = cycle_start_date
+        self.start_time = None
+        self.messages = None
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Triggered when LLM call starts"""
+        import time
+        self.start_time = time.time()
+        self.messages = messages
+
+        logger.info(f"[Callback] LLM call started for user '{self.user_id}' with {len(messages)} messages")
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Triggered when LLM call ends successfully"""
+        try:
+            # Record call duration
+            if self.start_time:
+                import time
+                call_duration = time.time() - self.start_time
+                logger.info(f"[Callback] LLM call completed for user '{self.user_id}' in {call_duration:.2f}s")
+
+            # Extract token usage from response dictionary
+            tokens_used = response.get("tokens_used")
+
+            if tokens_used and tokens_used > 0:
+                # Delegate recording work to repository
+                await token_usage_repo.increment_total_usage(
+                    self.user_id,
+                    tokens_used,
+                    self.cycle_start_date
+                )
+
+                logger.info(f"[Callback] Recorded {tokens_used} tokens for user '{self.user_id}'")
+            else:
+                logger.warning(f"[Callback] No token usage data available for user '{self.user_id}'")
+
+        except Exception as e:
+            logger.error(f"[Callback] Failed to record token usage for user '{self.user_id}': {e}")
+            # Don't re-raise exception to avoid affecting main LLM call flow
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Triggered when LLM call encounters an error"""
+        if self.start_time:
+            import time
+            call_duration = time.time() - self.start_time
+            logger.error(f"[Callback] LLM call failed for user '{self.user_id}' after {call_duration:.2f}s: {error}")
+        else:
+            logger.error(f"[Callback] LLM call failed for user '{self.user_id}': {error}")
+
+
+class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Detailed token recording callback handler.
+    Records separate prompt and completion token usage in addition to total usage.
+    """
+
+    def __init__(self, user_id: str, cycle_start_date: Optional[str] = None):
+        if not user_id:
+            raise ValueError("user_id must be provided for DetailedRedisTokenCallbackHandler")
+        self.user_id = user_id
+        self.cycle_start_date = cycle_start_date
+        self.start_time = None
+        self.messages = None
+        self.prompt_tokens = 0
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Triggered when LLM call starts"""
+        import time
+        self.start_time = time.time()
+        self.messages = messages
+
+        # Estimate input token count
+        self.prompt_tokens = self._estimate_prompt_tokens(messages)
+
+        logger.info(f"[DetailedCallback] LLM call started for user '{self.user_id}' with estimated {self.prompt_tokens} prompt tokens")
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Triggered when LLM call ends successfully"""
+        try:
+            # Record call duration
+            if self.start_time:
+                import time
+                call_duration = time.time() - self.start_time
+                logger.info(f"[DetailedCallback] LLM call completed for user '{self.user_id}' in {call_duration:.2f}s")
+
+            # Extract detailed token information from response
+            prompt_tokens, completion_tokens = self._extract_detailed_tokens(response)
+
+            # Ensure we have valid integers (not None)
+            prompt_tokens = prompt_tokens or 0
+            completion_tokens = completion_tokens or 0
+
+            if prompt_tokens > 0 or completion_tokens > 0:
+                # Use detailed token recording method
+                await token_usage_repo.increment_detailed_usage(
+                    self.user_id,
+                    prompt_tokens,
+                    completion_tokens,
+                    self.cycle_start_date
+                )
+
+                logger.info(f"[DetailedCallback] Recorded detailed tokens for user '{self.user_id}': prompt={prompt_tokens}, completion={completion_tokens}")
+            else:
+                logger.warning(f"[DetailedCallback] No detailed token usage data available for user '{self.user_id}'")
+
+        except Exception as e:
+            logger.error(f"[DetailedCallback] Failed to record detailed token usage for user '{self.user_id}': {e}")
+            # Don't re-raise exception to avoid affecting main LLM call flow
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Triggered when LLM call encounters an error"""
+        if self.start_time:
+            import time
+            call_duration = time.time() - self.start_time
+            logger.error(f"[DetailedCallback] LLM call failed for user '{self.user_id}' after {call_duration:.2f}s: {error}")
+        else:
+            logger.error(f"[DetailedCallback] LLM call failed for user '{self.user_id}': {error}")
+
+    def _estimate_prompt_tokens(self, messages: List[dict]) -> int:
+        """Estimate token count for input messages"""
+        total_chars = sum(len(msg.get('content', '')) for msg in messages)
+        # Rough estimation: 4 characters ≈ 1 token
+        return total_chars // 4
+
+    def _extract_detailed_tokens(self, response: dict) -> tuple[int, int]:
+        """
+        Extract detailed token information from response dictionary
+
+        Returns:
+            tuple: (prompt_tokens, completion_tokens)
+        """
+        # If response has detailed token information, use it first
+        prompt_tokens = response.get('prompt_tokens') or 0
+        completion_tokens = response.get('completion_tokens') or 0
+
+        if prompt_tokens > 0 and completion_tokens > 0:
+            return prompt_tokens, completion_tokens
+
+        # If only total token count is available, try to allocate
+        tokens_used = response.get('tokens_used') or 0
+        if tokens_used > 0:
+            # Use previously estimated prompt tokens
+            prompt_tokens = self.prompt_tokens
+            completion_tokens = max(0, tokens_used - prompt_tokens)
+            return prompt_tokens, completion_tokens
+
+        # If no token information, try to estimate from response content
+        content = response.get('content', '')
+        if content:
+            completion_tokens = len(content) // 4
+            prompt_tokens = self.prompt_tokens
+            return prompt_tokens, completion_tokens
+
+        return 0, 0
+
+
+class CompositeCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Composite callback handler that can execute multiple callback handlers simultaneously
+    """
+
+    def __init__(self, handlers: List[CustomAsyncCallbackHandler]):
+        self.handlers = handlers or []
+
+    def add_handler(self, handler: CustomAsyncCallbackHandler):
+        """Add a callback handler"""
+        self.handlers.append(handler)
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Execute start callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_start(messages, **kwargs)
+            except Exception as e:
+                logger.error(f"Error in callback handler {type(handler).__name__}.on_llm_start: {e}")
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Execute end callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_end(response, **kwargs)
+            except Exception as e:
+                logger.error(f"Error in callback handler {type(handler).__name__}.on_llm_end: {e}")
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Execute error callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_error(error, **kwargs)
+            except Exception as e:
+                logger.error(f"Error in callback handler {type(handler).__name__}.on_llm_error: {e}")
+
+
+# Convenience functions for creating common callback handlers
+def create_token_callback(user_id: str, cycle_start_date: Optional[str] = None) -> RedisTokenCallbackHandler:
+    """Create a basic token recording callback handler"""
+    return RedisTokenCallbackHandler(user_id, cycle_start_date)
+
+def create_detailed_token_callback(user_id: str, cycle_start_date: Optional[str] = None) -> DetailedRedisTokenCallbackHandler:
+    """Create a detailed token recording callback handler"""
+    return DetailedRedisTokenCallbackHandler(user_id, cycle_start_date)
+
+def create_composite_callback(*handlers: CustomAsyncCallbackHandler) -> CompositeCallbackHandler:
+    """Create a composite callback handler"""
+    return CompositeCallbackHandler(list(handlers))

aiecs/llm/openai_client.py

@@ -0,0 +1,125 @@
+import asyncio
+import logging
+from typing import Dict, Any, Optional, List, AsyncGenerator
+from openai import AsyncOpenAI
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+import httpx
+
+from .base_client import BaseLLMClient, LLMMessage, LLMResponse, ProviderNotAvailableError, RateLimitError
+from aiecs.config.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+class OpenAIClient(BaseLLMClient):
+    """OpenAI provider client"""
+
+    def __init__(self):
+        super().__init__("OpenAI")
+        self.settings = get_settings()
+        self._client: Optional[AsyncOpenAI] = None
+
+        # Token cost estimates (USD per 1K tokens)
+        self.token_costs = {
+            "gpt-4": {"input": 0.03, "output": 0.06},
+            "gpt-4-turbo": {"input": 0.01, "output": 0.03},
+            "gpt-3.5-turbo": {"input": 0.0015, "output": 0.002},
+            "gpt-4o": {"input": 0.005, "output": 0.015},
+            "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+        }
+
+    def _get_client(self) -> AsyncOpenAI:
+        """Lazy initialization of OpenAI client"""
+        if not self._client:
+            if not self.settings.openai_api_key:
+                raise ProviderNotAvailableError("OpenAI API key not configured")
+            self._client = AsyncOpenAI(api_key=self.settings.openai_api_key)
+        return self._client
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((httpx.RequestError, RateLimitError))
+    )
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """Generate text using OpenAI API"""
+        client = self._get_client()
+        model = model or "gpt-4-turbo"
+
+        # Convert to OpenAI message format
+        openai_messages = [{"role": msg.role, "content": msg.content} for msg in messages]
+
+        try:
+            response = await client.chat.completions.create(
+                model=model,
+                messages=openai_messages,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs
+            )
+
+            content = response.choices[0].message.content
+            tokens_used = response.usage.total_tokens if response.usage else None
+
+            # Estimate cost
+            input_tokens = response.usage.prompt_tokens if response.usage else 0
+            output_tokens = response.usage.completion_tokens if response.usage else 0
+            cost = self._estimate_cost(model, input_tokens, output_tokens, self.token_costs)
+
+            return LLMResponse(
+                content=content,
+                provider=self.provider_name,
+                model=model,
+                tokens_used=tokens_used,
+                cost_estimate=cost
+            )
+
+        except Exception as e:
+            if "rate_limit" in str(e).lower():
+                raise RateLimitError(f"OpenAI rate limit exceeded: {str(e)}")
+            raise
+
+    async def stream_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> AsyncGenerator[str, None]:
+        """Stream text using OpenAI API"""
+        client = self._get_client()
+        model = model or "gpt-4-turbo"
+
+        openai_messages = [{"role": msg.role, "content": msg.content} for msg in messages]
+
+        try:
+            stream = await client.chat.completions.create(
+                model=model,
+                messages=openai_messages,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                stream=True,
+                **kwargs
+            )
+
+            async for chunk in stream:
+                if chunk.choices[0].delta.content:
+                    yield chunk.choices[0].delta.content
+
+        except Exception as e:
+            if "rate_limit" in str(e).lower():
+                raise RateLimitError(f"OpenAI rate limit exceeded: {str(e)}")
+            raise
+
+    async def close(self):
+        """Clean up resources"""
+        if self._client:
+            await self._client.close()
+            self._client = None

aiecs/llm/vertex_client.py

@@ -0,0 +1,186 @@
+import asyncio
+import logging
+import os
+from typing import Dict, Any, Optional, List, AsyncGenerator
+from vertexai.generative_models import GenerativeModel, HarmCategory, HarmBlockThreshold
+import vertexai
+from google.oauth2 import service_account
+
+from .base_client import BaseLLMClient, LLMMessage, LLMResponse, ProviderNotAvailableError, RateLimitError
+from aiecs.config.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+class VertexAIClient(BaseLLMClient):
+    """Vertex AI provider client"""
+
+    def __init__(self):
+        super().__init__("Vertex")
+        self.settings = get_settings()
+        self._initialized = False
+
+        # Token cost estimates (USD per 1K tokens)
+        self.token_costs = {
+            "gemini-2.5-pro": {"input": 0.00125, "output": 0.00375},
+            "gemini-2.5-flash": {"input": 0.000075, "output": 0.0003},
+        }
+
+    def _init_vertex_ai(self):
+        """Lazy initialization of Vertex AI with proper authentication"""
+        if not self._initialized:
+            if not self.settings.vertex_project_id:
+                raise ProviderNotAvailableError("Vertex AI project ID not configured")
+
+            try:
+                # Set up Google Cloud authentication
+                credentials = None
+
+                # Check if GOOGLE_APPLICATION_CREDENTIALS is configured
+                if self.settings.google_application_credentials:
+                    credentials_path = self.settings.google_application_credentials
+                    if os.path.exists(credentials_path):
+                        # Set the environment variable for Google Cloud SDK
+                        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = credentials_path
+                        self.logger.info(f"Using Google Cloud credentials from: {credentials_path}")
+                    else:
+                        self.logger.warning(f"Google Cloud credentials file not found: {credentials_path}")
+                        raise ProviderNotAvailableError(f"Google Cloud credentials file not found: {credentials_path}")
+                elif 'GOOGLE_APPLICATION_CREDENTIALS' in os.environ:
+                    self.logger.info("Using Google Cloud credentials from environment variable")
+                else:
+                    self.logger.warning("No Google Cloud credentials configured. Using default authentication.")
+
+                # Initialize Vertex AI
+                vertexai.init(
+                    project=self.settings.vertex_project_id,
+                    location=getattr(self.settings, 'vertex_location', 'us-central1')
+                )
+                self._initialized = True
+                self.logger.info(f"Vertex AI initialized for project {self.settings.vertex_project_id}")
+
+            except Exception as e:
+                raise ProviderNotAvailableError(f"Failed to initialize Vertex AI: {str(e)}")
+
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """Generate text using Vertex AI"""
+        self._init_vertex_ai()
+        model_name = model or "gemini-2.5-pro"
+
+        try:
+            # Use the stable Vertex AI API
+            model_instance = GenerativeModel(model_name)
+
+            # Convert messages to Vertex AI format
+            if len(messages) == 1 and messages[0].role == "user":
+                prompt = messages[0].content
+            else:
+                # For multi-turn conversations, combine messages
+                prompt = "\n".join([f"{msg.role}: {msg.content}" for msg in messages])
+
+            response = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: model_instance.generate_content(
+                    prompt,
+                    generation_config={
+                        "temperature": temperature,
+                        "max_output_tokens": max_tokens or 8192,  # Increased to account for thinking tokens
+                        "top_p": 0.95,
+                        "top_k": 40,
+                    },
+                    safety_settings={
+                        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                    }
+                )
+            )
+
+            # Handle response content safely
+            try:
+                content = response.text
+                self.logger.debug(f"Vertex AI response received: {content[:100]}...")
+            except ValueError as ve:
+                # Handle cases where response has no content (safety filters, etc.)
+                self.logger.warning(f"Vertex AI response error: {str(ve)}")
+                self.logger.debug(f"Full response object: {response}")
+
+                # Check if response has candidates but no text
+                if hasattr(response, 'candidates') and response.candidates:
+                    candidate = response.candidates[0]
+                    self.logger.debug(f"Candidate finish_reason: {getattr(candidate, 'finish_reason', 'unknown')}")
+
+                    # If finish_reason is MAX_TOKENS, it might be due to thinking tokens
+                    if hasattr(candidate, 'finish_reason') and candidate.finish_reason == 'MAX_TOKENS':
+                        content = "[Response truncated due to token limit - consider increasing max_tokens for Gemini 2.5 models]"
+                        self.logger.warning("Response truncated due to MAX_TOKENS - Gemini 2.5 uses thinking tokens")
+                    elif "no parts" in str(ve).lower() or "safety filters" in str(ve).lower():
+                        content = "[Response blocked by safety filters or has no content]"
+                        self.logger.warning(f"Vertex AI response blocked or empty: {str(ve)}")
+                    else:
+                        content = f"[Response error: {str(ve)}]"
+                else:
+                    content = f"[Response error: {str(ve)}]"
+
+            # Vertex AI doesn't provide detailed token usage in the response
+            tokens_used = self._count_tokens_estimate(prompt + content)
+            cost = self._estimate_cost(
+                model_name,
+                self._count_tokens_estimate(prompt),
+                self._count_tokens_estimate(content),
+                self.token_costs
+            )
+
+            return LLMResponse(
+                content=content,
+                provider=self.provider_name,
+                model=model_name,
+                tokens_used=tokens_used,
+                cost_estimate=cost
+            )
+
+        except Exception as e:
+            if "quota" in str(e).lower() or "limit" in str(e).lower():
+                raise RateLimitError(f"Vertex AI quota exceeded: {str(e)}")
+            # Handle specific Vertex AI response errors
+            if "cannot get the response text" in str(e).lower() or "safety filters" in str(e).lower():
+                self.logger.warning(f"Vertex AI response issue: {str(e)}")
+                # Return a response indicating the issue
+                return LLMResponse(
+                    content="[Response unavailable due to safety filters or content policy]",
+                    provider=self.provider_name,
+                    model=model_name,
+                    tokens_used=self._count_tokens_estimate(prompt),
+                    cost_estimate=0.0
+                )
+            raise
+
+    async def stream_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> AsyncGenerator[str, None]:
+        """Stream text using Vertex AI (simulated streaming)"""
+        # Vertex AI streaming is more complex, for now fall back to non-streaming
+        response = await self.generate_text(messages, model, temperature, max_tokens, **kwargs)
+
+        # Simulate streaming by yielding words
+        words = response.content.split()
+        for word in words:
+            yield word + " "
+            await asyncio.sleep(0.05)  # Small delay to simulate streaming
+
+    async def close(self):
+        """Clean up resources"""
+        # Vertex AI doesn't require explicit cleanup
+        self._initialized = False