dataknobs-bots 0.2.4__py3-none-any.whl

dataknobs_bots/middleware/cost.py

@@ -0,0 +1,421 @@
+"""Cost tracking middleware for monitoring LLM usage."""
+
+import json
+import logging
+from typing import Any
+
+from dataknobs_bots.bot.context import BotContext
+
+from .base import Middleware
+
+logger = logging.getLogger(__name__)
+
+
+class CostTrackingMiddleware(Middleware):
+    """Middleware for tracking LLM API costs and usage.
+
+    Monitors token usage across different providers (Ollama, OpenAI, Anthropic, etc.)
+    to help optimize costs and track budgets.
+
+    Attributes:
+        track_tokens: Whether to track token usage
+        cost_rates: Token cost rates per provider/model
+        usage_stats: Accumulated usage statistics by client_id
+
+    Example:
+        ```python
+        # Create middleware with default rates
+        middleware = CostTrackingMiddleware()
+
+        # Or with custom rates
+        middleware = CostTrackingMiddleware(
+            cost_rates={
+                "openai": {
+                    "gpt-4o": {"input": 0.0025, "output": 0.01},
+                },
+            }
+        )
+
+        # Get stats
+        stats = middleware.get_client_stats("my-client")
+        total = middleware.get_total_cost()
+
+        # Export to JSON
+        json_data = middleware.export_stats_json()
+        ```
+    """
+
+    # Default cost rates (USD per 1K tokens) - Updated Dec 2024
+    DEFAULT_RATES: dict[str, Any] = {
+        "ollama": {"input": 0.0, "output": 0.0},  # Free (infrastructure cost only)
+        "openai": {
+            "gpt-4o": {"input": 0.0025, "output": 0.01},
+            "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
+            "gpt-4-turbo": {"input": 0.01, "output": 0.03},
+            "gpt-4": {"input": 0.03, "output": 0.06},
+            "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
+            "o1": {"input": 0.015, "output": 0.06},
+            "o1-mini": {"input": 0.003, "output": 0.012},
+        },
+        "anthropic": {
+            "claude-3-5-sonnet": {"input": 0.003, "output": 0.015},
+            "claude-3-5-haiku": {"input": 0.0008, "output": 0.004},
+            "claude-3-opus": {"input": 0.015, "output": 0.075},
+            "claude-3-sonnet": {"input": 0.003, "output": 0.015},
+            "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
+        },
+        "google": {
+            "gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
+            "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
+            "gemini-2.0-flash": {"input": 0.0001, "output": 0.0004},
+        },
+    }
+
+    def __init__(
+        self,
+        track_tokens: bool = True,
+        cost_rates: dict[str, Any] | None = None,
+    ):
+        """Initialize cost tracking middleware.
+
+        Args:
+            track_tokens: Enable token tracking
+            cost_rates: Optional custom cost rates (merged with defaults)
+        """
+        self.track_tokens = track_tokens
+        # Merge custom rates with defaults
+        self.cost_rates = self.DEFAULT_RATES.copy()
+        if cost_rates:
+            for provider, rates in cost_rates.items():
+                if provider in self.cost_rates:
+                    if isinstance(rates, dict) and isinstance(
+                        self.cost_rates[provider], dict
+                    ):
+                        self.cost_rates[provider].update(rates)
+                    else:
+                        self.cost_rates[provider] = rates
+                else:
+                    self.cost_rates[provider] = rates
+
+        self._usage_stats: dict[str, dict[str, Any]] = {}
+        self._logger = logging.getLogger(f"{__name__}.CostTracker")
+
+    async def before_message(self, message: str, context: BotContext) -> None:
+        """Track message before processing (mainly for logging).
+
+        Args:
+            message: User's input message
+            context: Bot context
+        """
+        # Estimate input tokens (rough approximation: ~4 chars per token)
+        estimated_tokens = len(message) // 4
+        self._logger.debug(f"Estimated input tokens: {estimated_tokens}")
+
+    async def after_message(
+        self, response: str, context: BotContext, **kwargs: Any
+    ) -> None:
+        """Track costs after bot response.
+
+        Args:
+            response: Bot's generated response
+            context: Bot context
+            **kwargs: Should contain 'tokens_used', 'provider', 'model' if available
+        """
+        if not self.track_tokens:
+            return
+
+        client_id = context.client_id
+
+        # Extract provider and model info
+        provider = kwargs.get("provider", "unknown")
+        model = kwargs.get("model", "unknown")
+
+        # Get token counts
+        tokens_used = kwargs.get("tokens_used", {})
+        if isinstance(tokens_used, int):
+            # If single number, assume it's total and estimate split
+            input_tokens = len(context.session_metadata.get("last_message", "")) // 4
+            output_tokens = tokens_used - input_tokens
+        else:
+            input_tokens = int(
+                tokens_used.get(
+                    "input",
+                    tokens_used.get(
+                        "prompt_tokens",
+                        len(context.session_metadata.get("last_message", "")) // 4,
+                    ),
+                )
+            )
+            output_tokens = int(
+                tokens_used.get(
+                    "output",
+                    tokens_used.get("completion_tokens", len(response) // 4),
+                )
+            )
+
+        # Calculate cost
+        cost = self._calculate_cost(provider, model, input_tokens, output_tokens)
+
+        # Update stats
+        if client_id not in self._usage_stats:
+            self._usage_stats[client_id] = {
+                "client_id": client_id,
+                "total_requests": 0,
+                "total_input_tokens": 0,
+                "total_output_tokens": 0,
+                "total_cost_usd": 0.0,
+                "by_provider": {},
+            }
+
+        stats = self._usage_stats[client_id]
+        stats["total_requests"] += 1
+        stats["total_input_tokens"] += input_tokens
+        stats["total_output_tokens"] += output_tokens
+        stats["total_cost_usd"] += cost
+
+        # Track by provider
+        if provider not in stats["by_provider"]:
+            stats["by_provider"][provider] = {
+                "requests": 0,
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "cost_usd": 0.0,
+                "by_model": {},
+            }
+
+        provider_stats = stats["by_provider"][provider]
+        provider_stats["requests"] += 1
+        provider_stats["input_tokens"] += input_tokens
+        provider_stats["output_tokens"] += output_tokens
+        provider_stats["cost_usd"] += cost
+
+        # Track by model within provider
+        if model not in provider_stats["by_model"]:
+            provider_stats["by_model"][model] = {
+                "requests": 0,
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "cost_usd": 0.0,
+            }
+
+        model_stats = provider_stats["by_model"][model]
+        model_stats["requests"] += 1
+        model_stats["input_tokens"] += input_tokens
+        model_stats["output_tokens"] += output_tokens
+        model_stats["cost_usd"] += cost
+
+        self._logger.info(
+            f"Client {client_id}: {provider}/{model} - "
+            f"{input_tokens} in + {output_tokens} out tokens, "
+            f"cost: ${cost:.6f}, total: ${stats['total_cost_usd']:.6f}"
+        )
+
+    async def post_stream(
+        self, message: str, response: str, context: BotContext
+    ) -> None:
+        """Track costs after streaming completes.
+
+        For streaming responses, token counts are estimated from text length
+        since exact counts may not be available until the stream completes.
+
+        Args:
+            message: Original user message
+            response: Complete accumulated response from streaming
+            context: Bot context
+        """
+        if not self.track_tokens:
+            return
+
+        client_id = context.client_id
+
+        # For streaming, we estimate tokens from text length (~4 chars per token)
+        input_tokens = len(message) // 4
+        output_tokens = len(response) // 4
+
+        # Get provider/model from context metadata if available
+        provider = context.session_metadata.get("provider", "unknown")
+        model = context.session_metadata.get("model", "unknown")
+
+        # Calculate cost
+        cost = self._calculate_cost(provider, model, input_tokens, output_tokens)
+
+        # Update stats
+        if client_id not in self._usage_stats:
+            self._usage_stats[client_id] = {
+                "total_requests": 0,
+                "total_input_tokens": 0,
+                "total_output_tokens": 0,
+                "total_cost_usd": 0.0,
+                "by_provider": {},
+            }
+
+        stats = self._usage_stats[client_id]
+        stats["total_requests"] += 1
+        stats["total_input_tokens"] += input_tokens
+        stats["total_output_tokens"] += output_tokens
+        stats["total_cost_usd"] += cost
+
+        # Track by provider
+        if provider not in stats["by_provider"]:
+            stats["by_provider"][provider] = {
+                "requests": 0,
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "cost_usd": 0.0,
+            }
+
+        provider_stats = stats["by_provider"][provider]
+        provider_stats["requests"] += 1
+        provider_stats["input_tokens"] += input_tokens
+        provider_stats["output_tokens"] += output_tokens
+        provider_stats["cost_usd"] += cost
+
+        self._logger.info(
+            f"Stream complete - Client {client_id}: {provider}/{model} - "
+            f"~{input_tokens} in + ~{output_tokens} out tokens (estimated), "
+            f"cost: ${cost:.6f}, total: ${stats['total_cost_usd']:.6f}"
+        )
+
+    async def on_error(
+        self, error: Exception, message: str, context: BotContext
+    ) -> None:
+        """Log errors but don't track costs for failed requests.
+
+        Args:
+            error: The exception that occurred
+            message: User message that caused the error
+            context: Bot context
+        """
+        self._logger.warning(
+            f"Error during request for client {context.client_id}: {error}"
+        )
+
+    def _calculate_cost(
+        self, provider: str, model: str, input_tokens: int, output_tokens: int
+    ) -> float:
+        """Calculate cost for token usage.
+
+        Args:
+            provider: LLM provider name
+            model: Model name
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+
+        Returns:
+            Cost in USD
+        """
+        # Get rates for provider/model
+        if provider in self.cost_rates:
+            provider_rates = self.cost_rates[provider]
+
+            if isinstance(provider_rates, dict):
+                # Check if model-specific rates exist
+                if model in provider_rates:
+                    rates = provider_rates[model]
+                elif "input" in provider_rates:
+                    # Use generic rates for provider (e.g., ollama)
+                    rates = provider_rates
+                else:
+                    # Try partial model name match
+                    for model_key in provider_rates:
+                        if model_key in model or model in model_key:
+                            rates = provider_rates[model_key]
+                            break
+                    else:
+                        return 0.0
+            else:
+                return 0.0
+
+            # Calculate cost (rates are per 1K tokens)
+            input_cost = (input_tokens / 1000) * float(rates.get("input", 0.0))
+            output_cost = (output_tokens / 1000) * float(rates.get("output", 0.0))
+            return float(input_cost + output_cost)
+
+        return 0.0
+
+    def get_client_stats(self, client_id: str) -> dict[str, Any] | None:
+        """Get usage statistics for a client.
+
+        Args:
+            client_id: Client identifier
+
+        Returns:
+            Usage statistics or None if not found
+        """
+        return self._usage_stats.get(client_id)
+
+    def get_all_stats(self) -> dict[str, dict[str, Any]]:
+        """Get all usage statistics.
+
+        Returns:
+            Dictionary mapping client_id to statistics
+        """
+        return self._usage_stats.copy()
+
+    def get_total_cost(self) -> float:
+        """Get total cost across all clients.
+
+        Returns:
+            Total cost in USD
+        """
+        return float(
+            sum(stats["total_cost_usd"] for stats in self._usage_stats.values())
+        )
+
+    def get_total_tokens(self) -> dict[str, int]:
+        """Get total tokens across all clients.
+
+        Returns:
+            Dictionary with 'input', 'output', and 'total' token counts
+        """
+        input_tokens = sum(
+            stats["total_input_tokens"] for stats in self._usage_stats.values()
+        )
+        output_tokens = sum(
+            stats["total_output_tokens"] for stats in self._usage_stats.values()
+        )
+        return {
+            "input": input_tokens,
+            "output": output_tokens,
+            "total": input_tokens + output_tokens,
+        }
+
+    def clear_stats(self, client_id: str | None = None) -> None:
+        """Clear usage statistics.
+
+        Args:
+            client_id: If provided, clear only this client. Otherwise clear all.
+        """
+        if client_id:
+            if client_id in self._usage_stats:
+                del self._usage_stats[client_id]
+        else:
+            self._usage_stats.clear()
+
+    def export_stats_json(self, indent: int = 2) -> str:
+        """Export all statistics as JSON.
+
+        Args:
+            indent: JSON indentation level
+
+        Returns:
+            JSON string of all statistics
+        """
+        return json.dumps(self._usage_stats, indent=indent)
+
+    def export_stats_csv(self) -> str:
+        """Export statistics as CSV (one row per client).
+
+        Returns:
+            CSV string with headers
+        """
+        lines = [
+            "client_id,total_requests,total_input_tokens,total_output_tokens,total_cost_usd"
+        ]
+        for client_id, stats in self._usage_stats.items():
+            lines.append(
+                f"{client_id},{stats['total_requests']},"
+                f"{stats['total_input_tokens']},{stats['total_output_tokens']},"
+                f"{stats['total_cost_usd']:.6f}"
+            )
+        return "\n".join(lines)

dataknobs_bots/middleware/logging.py

@@ -0,0 +1,184 @@
+"""Logging middleware for conversation tracking."""
+
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from dataknobs_bots.bot.context import BotContext
+
+from .base import Middleware
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingMiddleware(Middleware):
+    """Middleware for tracking conversation interactions.
+
+    Logs all user messages and bot responses with context
+    for monitoring, debugging, and analytics.
+
+    Attributes:
+        log_level: Logging level to use (default: INFO)
+        include_metadata: Whether to include full context metadata
+        json_format: Whether to output logs in JSON format
+
+    Example:
+        ```python
+        # Basic usage
+        middleware = LoggingMiddleware()
+
+        # With JSON format for log aggregation
+        middleware = LoggingMiddleware(
+            log_level="INFO",
+            include_metadata=True,
+            json_format=True
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        log_level: str = "INFO",
+        include_metadata: bool = True,
+        json_format: bool = False,
+    ):
+        """Initialize logging middleware.
+
+        Args:
+            log_level: Logging level (DEBUG, INFO, WARNING, ERROR)
+            include_metadata: Whether to log full context metadata
+            json_format: Whether to output in JSON format
+        """
+        self.log_level = log_level
+        self.include_metadata = include_metadata
+        self.json_format = json_format
+        self._logger = logging.getLogger(f"{__name__}.ConversationLogger")
+        self._logger.setLevel(getattr(logging, log_level.upper()))
+
+    async def before_message(self, message: str, context: BotContext) -> None:
+        """Called before processing user message.
+
+        Args:
+            message: User's input message
+            context: Bot context with conversation and user info
+        """
+        log_data = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event": "user_message",
+            "client_id": context.client_id,
+            "user_id": context.user_id,
+            "conversation_id": context.conversation_id,
+            "message_length": len(message),
+        }
+
+        if self.include_metadata:
+            log_data["session_metadata"] = context.session_metadata
+            log_data["request_metadata"] = context.request_metadata
+
+        if self.json_format:
+            self._logger.info(json.dumps(log_data))
+        else:
+            self._logger.info(f"User message: {log_data}")
+
+        # Log content at DEBUG level (first 200 chars)
+        self._logger.debug(f"Message content: {message[:200]}...")
+
+    async def after_message(
+        self, response: str, context: BotContext, **kwargs: Any
+    ) -> None:
+        """Called after generating bot response.
+
+        Args:
+            response: Bot's generated response
+            context: Bot context
+            **kwargs: Additional data (e.g., tokens_used, response_time_ms)
+        """
+        log_data = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event": "bot_response",
+            "client_id": context.client_id,
+            "user_id": context.user_id,
+            "conversation_id": context.conversation_id,
+            "response_length": len(response),
+        }
+
+        # Add optional metrics
+        if "tokens_used" in kwargs:
+            log_data["tokens_used"] = kwargs["tokens_used"]
+        if "response_time_ms" in kwargs:
+            log_data["response_time_ms"] = kwargs["response_time_ms"]
+        if "provider" in kwargs:
+            log_data["provider"] = kwargs["provider"]
+        if "model" in kwargs:
+            log_data["model"] = kwargs["model"]
+
+        if self.include_metadata:
+            log_data["session_metadata"] = context.session_metadata
+            log_data["request_metadata"] = context.request_metadata
+
+        if self.json_format:
+            self._logger.info(json.dumps(log_data))
+        else:
+            self._logger.info(f"Bot response: {log_data}")
+
+        # Log content at DEBUG level (first 200 chars)
+        self._logger.debug(f"Response content: {response[:200]}...")
+
+    async def post_stream(
+        self, message: str, response: str, context: BotContext
+    ) -> None:
+        """Called after streaming response completes.
+
+        Args:
+            message: Original user message
+            response: Complete accumulated response from streaming
+            context: Bot context
+        """
+        log_data = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event": "stream_complete",
+            "client_id": context.client_id,
+            "user_id": context.user_id,
+            "conversation_id": context.conversation_id,
+            "message_length": len(message),
+            "response_length": len(response),
+        }
+
+        if self.include_metadata:
+            log_data["session_metadata"] = context.session_metadata
+            log_data["request_metadata"] = context.request_metadata
+
+        if self.json_format:
+            self._logger.info(json.dumps(log_data))
+        else:
+            self._logger.info(f"Stream complete: {log_data}")
+
+        # Log content at DEBUG level (first 200 chars each)
+        self._logger.debug(f"Streamed message: {message[:200]}...")
+        self._logger.debug(f"Streamed response: {response[:200]}...")
+
+    async def on_error(
+        self, error: Exception, message: str, context: BotContext
+    ) -> None:
+        """Called when an error occurs during message processing.
+
+        Args:
+            error: The exception that occurred
+            message: User message that caused the error
+            context: Bot context
+        """
+        log_data = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "event": "error",
+            "client_id": context.client_id,
+            "user_id": context.user_id,
+            "conversation_id": context.conversation_id,
+            "error_type": type(error).__name__,
+            "error_message": str(error),
+        }
+
+        if self.json_format:
+            self._logger.error(json.dumps(log_data), exc_info=error)
+        else:
+            self._logger.error(f"Error processing message: {log_data}", exc_info=error)

dataknobs_bots/reasoning/__init__.py

@@ -0,0 +1,65 @@
+"""Reasoning strategies for DynaBot."""
+
+from typing import Any
+
+from .base import ReasoningStrategy
+from .react import ReActReasoning
+from .simple import SimpleReasoning
+
+__all__ = [
+    "ReasoningStrategy",
+    "SimpleReasoning",
+    "ReActReasoning",
+    "create_reasoning_from_config",
+]
+
+
+def create_reasoning_from_config(config: dict[str, Any]) -> ReasoningStrategy:
+    """Create reasoning strategy from configuration.
+
+    Args:
+        config: Reasoning configuration with:
+            - strategy: Strategy type ('simple', 'react')
+            - max_iterations: For ReAct, max reasoning loops (default: 5)
+            - verbose: Enable debug logging for reasoning steps (default: False)
+            - store_trace: Store reasoning trace in conversation metadata (default: False)
+
+    Returns:
+        Configured reasoning strategy instance
+
+    Raises:
+        ValueError: If strategy type is not supported
+
+    Example:
+        ```python
+        # Simple reasoning
+        config = {"strategy": "simple"}
+        strategy = create_reasoning_from_config(config)
+
+        # ReAct reasoning with trace storage
+        config = {
+            "strategy": "react",
+            "max_iterations": 5,
+            "verbose": True,
+            "store_trace": True
+        }
+        strategy = create_reasoning_from_config(config)
+        ```
+    """
+    strategy_type = config.get("strategy", "simple").lower()
+
+    if strategy_type == "simple":
+        return SimpleReasoning()
+
+    elif strategy_type == "react":
+        return ReActReasoning(
+            max_iterations=config.get("max_iterations", 5),
+            verbose=config.get("verbose", False),
+            store_trace=config.get("store_trace", False),
+        )
+
+    else:
+        raise ValueError(
+            f"Unknown reasoning strategy: {strategy_type}. "
+            f"Available strategies: simple, react"
+        )