scitex 2.0.0__py2.py3-none-any.whl

scitex/ai/genai/cost_tracker.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2025-05-31 10:15:00"
+# Author: ywatanabe
+# File: ./src/scitex/ai/genai/cost_tracker.py
+
+"""
+Tracks token usage and costs for AI API calls.
+
+This module provides cost tracking functionality including:
+- Token counting for input/output
+- Cost calculation based on model pricing
+- Usage statistics aggregation
+"""
+
+from typing import Dict, Optional, List, Any
+from dataclasses import dataclass, field
+
+from .calc_cost import calc_cost
+
+
+@dataclass
+class TokenUsage:
+    """Token usage statistics.
+
+    Attributes
+    ----------
+    input_tokens : int
+        Total input tokens used
+    output_tokens : int
+        Total output tokens used
+    total_tokens : int
+        Total tokens (input + output)
+    """
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+    @property
+    def total_tokens(self) -> int:
+        """Get total token count."""
+        return self.input_tokens + self.output_tokens
+
+
+class CostTracker:
+    """Tracks token usage and associated costs.
+
+    Example
+    -------
+    >>> tracker = CostTracker("openai", "gpt-4")
+    >>> tracker.update(input_tokens=100, output_tokens=50)
+    >>> cost = tracker.calculate_cost()
+    >>> print(f"Total cost: ${cost:.4f}")
+    Total cost: $0.0045
+
+    Parameters
+    ----------
+    provider : str
+        The provider name (e.g., "openai", "anthropic")
+    model : str
+        The model name for cost calculation
+    """
+
+    def __init__(self, provider: str, model: str):
+        """Initialize cost tracker.
+
+        Parameters
+        ----------
+        provider : str
+            The provider name
+        model : str
+            The model name
+        """
+        self.provider = provider
+        self.model = model
+
+        # Total token counters
+        self.input_tokens = 0
+        self.output_tokens = 0
+
+        # Session token counters
+        self.session_input_tokens = 0
+        self.session_output_tokens = 0
+
+        # Token history for detailed tracking
+        self._history: List[TokenUsage] = []
+
+        # Request counter
+        self.request_count = 0
+
+    # Properties for backward compatibility
+    @property
+    def total_prompt_tokens(self) -> int:
+        """Get total prompt (input) tokens."""
+        return self.input_tokens
+
+    @property
+    def total_completion_tokens(self) -> int:
+        """Get total completion (output) tokens."""
+        return self.output_tokens
+
+    @property
+    def total_cost(self) -> float:
+        """Get total cost."""
+        return self.calculate_cost()
+
+    def update(self, input_tokens: int, output_tokens: int) -> None:
+        """Update token counts.
+
+        Parameters
+        ----------
+        input_tokens : int
+            Number of input tokens to add
+        output_tokens : int
+            Number of output tokens to add
+        """
+        # Update total counters
+        self.input_tokens += input_tokens
+        self.output_tokens += output_tokens
+
+        # Update session counters
+        self.session_input_tokens += input_tokens
+        self.session_output_tokens += output_tokens
+
+        # Record in history
+        usage = TokenUsage(input_tokens=input_tokens, output_tokens=output_tokens)
+        self._history.append(usage)
+
+        # Increment request count
+        self.request_count += 1
+
+    def track_usage(self, model: str, usage: TokenUsage) -> None:
+        """Track usage (for backward compatibility).
+
+        Parameters
+        ----------
+        model : str
+            Model name (ignored, uses instance model)
+        usage : TokenUsage
+            Token usage to track
+        """
+        self.update(usage.input_tokens, usage.output_tokens)
+
+    def calculate_cost(self) -> float:
+        """Calculate total cost based on token usage.
+
+        Returns
+        -------
+        float
+            Total cost in USD
+        """
+        return calc_cost(
+            model=self.model,
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+        )
+
+    def calculate_session_cost(self) -> float:
+        """Calculate session cost based on session token usage.
+
+        Returns
+        -------
+        float
+            Session cost in USD
+        """
+        return calc_cost(
+            model=self.model,
+            input_tokens=self.session_input_tokens,
+            output_tokens=self.session_output_tokens,
+        )
+
+    def reset_session(self) -> None:
+        """Reset session counters while keeping totals."""
+        self.session_input_tokens = 0
+        self.session_output_tokens = 0
+
+    def reset(self) -> None:
+        """Reset all token usage and cost tracking."""
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self.session_input_tokens = 0
+        self.session_output_tokens = 0
+        self._history = []
+
+    def get_usage_stats(self) -> Dict[str, Any]:
+        """Get comprehensive usage statistics.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Dictionary containing usage statistics
+        """
+        return {
+            "total_input_tokens": self.input_tokens,
+            "total_output_tokens": self.output_tokens,
+            "total_tokens": self.input_tokens + self.output_tokens,
+            "session_input_tokens": self.session_input_tokens,
+            "session_output_tokens": self.session_output_tokens,
+            "session_total_tokens": self.session_input_tokens
+            + self.session_output_tokens,
+        }
+
+    def get_summary(self) -> str:
+        """Get a cost summary (alias for format_cost_summary).
+
+        Returns
+        -------
+        str
+            Formatted cost summary
+        """
+        return self.format_cost_summary()
+
+    def format_cost_summary(self) -> str:
+        """Format a human-readable cost summary.
+
+        Returns
+        -------
+        str
+            Formatted cost summary
+        """
+        total_cost = self.calculate_cost()
+        session_cost = self.calculate_session_cost()
+
+        return (
+            f"Cost Summary for {self.provider} - {self.model}\n"
+            f"{'=' * 50}\n"
+            f"Total Usage:\n"
+            f"  Input tokens:  {self.input_tokens:,}\n"
+            f"  Output tokens: {self.output_tokens:,}\n"
+            f"  Total tokens:  {self.input_tokens + self.output_tokens:,}\n"
+            f"  Total cost:    ${total_cost:.4f}\n"
+            f"\n"
+            f"Session Usage:\n"
+            f"  Input tokens:  {self.session_input_tokens:,}\n"
+            f"  Output tokens: {self.session_output_tokens:,}\n"
+            f"  Total tokens:  {self.session_input_tokens + self.session_output_tokens:,}\n"
+            f"  Session cost:  ${session_cost:.4f}\n"
+        )
+
+    def get_detailed_report(self) -> Dict[str, Any]:
+        """Get a detailed cost and usage report.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Comprehensive report including usage, costs, and statistics
+        """
+        return {
+            "provider": self.provider,
+            "model": self.model,
+            "total_cost": self.calculate_cost(),
+            "session_cost": self.calculate_session_cost(),
+            "usage_stats": self.get_usage_stats(),
+            "num_api_calls": len(self._history),
+            "average_per_call": {
+                "input_tokens": (
+                    self.input_tokens / len(self._history) if self._history else 0
+                ),
+                "output_tokens": (
+                    self.output_tokens / len(self._history) if self._history else 0
+                ),
+            },
+        }
+
+    def __repr__(self) -> str:
+        """String representation of CostTracker."""
+        cost = self.calculate_cost()
+        return (
+            f"CostTracker(provider={self.provider}, "
+            f"model={self.model}, "
+            f"total_tokens={self.input_tokens + self.output_tokens}, "
+            f"cost=${cost:.4f})"
+        )
+
+
+# EOF

scitex/ai/genai/deepseek.py

@@ -0,0 +1,188 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-08 20:33:49 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/ai/_gen_ai/_DeepSeek.py
+
+"""
+1. Functionality:
+   - Implements DeepSeek Code LLM API interface
+2. Input:
+   - Text prompts for code generation
+3. Output:
+   - Generated code responses (streaming or static)
+4. Prerequisites:
+   - DEEPSEEK_API_KEY environment variable
+   - requests library
+
+DEPRECATED: This module is deprecated. Please use deepseek_provider.py instead.
+The new provider-based architecture offers better modularity and maintainability.
+"""
+
+"""Imports"""
+import json
+import os
+import sys
+from typing import Dict, Generator, List, Optional
+import warnings
+
+import scitex
+import requests
+
+from .base_genai import BaseGenAI
+
+warnings.warn(
+    "deepseek.py is deprecated. Please use deepseek_provider.py instead. "
+    "See PROVIDER_MIGRATION_GUIDE.md for migration instructions.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+"""Warnings"""
+# scitex.pd.ignore_SettingWithCopyWarning()
+# warnings.simplefilter("ignore", UserWarning)
+# with warnings.catch_warnings():
+#     warnings.simplefilter("ignore", UserWarning)
+
+"""Parameters"""
+# from scitex.io import load_configs
+# CONFIG = load_configs()
+
+"""Functions & Classes"""
+"""Imports"""
+from .base_genai import BaseGenAI
+from openai import OpenAI as _OpenAI
+
+"""Functions & Classes"""
+
+
+class DeepSeek(BaseGenAI):
+    def __init__(
+        self,
+        system_setting="",
+        model="deepseek-chat",
+        api_key="",
+        stream=False,
+        seed=None,
+        n_keep=1,
+        temperature=1.0,
+        chat_history=None,
+        max_tokens=4096,
+    ):
+        super().__init__(
+            system_setting=system_setting,
+            model=model,
+            api_key=api_key,
+            stream=stream,
+            n_keep=n_keep,
+            temperature=temperature,
+            provider="DeepSeek",
+            chat_history=chat_history,
+            max_tokens=max_tokens,
+        )
+
+    def _init_client(self):
+        # client = _OpenAI(api_key=self.api_key, base_url="https://api.deepseek.com")
+        client = _OpenAI(api_key=self.api_key, base_url="https://api.deepseek.com/beta")
+
+        return client
+
+    def _api_call_static(self):
+        kwargs = dict(
+            model=self.model,
+            messages=self.history,
+            seed=self.seed,
+            stream=False,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+        )
+
+        output = self.client.chat.completions.create(**kwargs)
+        self.input_tokens += output.usage.prompt_tokens
+        self.output_tokens += output.usage.completion_tokens
+
+        out_text = output.choices[0].message.content
+
+        return out_text
+
+    def _api_call_stream(self):
+        kwargs = dict(
+            model=self.model,
+            messages=self.history,
+            max_tokens=self.max_tokens,
+            n=1,
+            stream=self.stream,
+            seed=self.seed,
+            temperature=self.temperature,
+        )
+
+        stream = self.client.chat.completions.create(**kwargs)
+        buffer = ""
+
+        for chunk in stream:
+            if chunk:
+                try:
+                    self.input_tokens += chunk.usage.prompt_tokens
+                except:
+                    pass
+                try:
+                    self.output_tokens += chunk.usage.completion_tokens
+                except:
+                    pass
+
+                try:
+                    current_text = chunk.choices[0].delta.content
+                    if current_text:
+                        buffer += current_text
+                        if any(char in ".!?\n " for char in current_text):
+                            yield buffer
+                            buffer = ""
+                except Exception as e:
+                    pass
+
+        if buffer:
+            yield buffer
+
+
+if __name__ == "__main__":
+    # -----------------------------------
+    # Initiatialization of scitex format
+    # -----------------------------------
+    import sys
+
+    import matplotlib.pyplot as plt
+
+    # Configurations
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(
+        sys,
+        plt,
+        verbose=False,
+        agg=True,
+        # sdir_suffix="",
+    )
+
+    # # Argument parser
+    # script_mode = scitex.gen.is_script()
+    # import argparse
+    # parser = argparse.ArgumentParser(description='')
+    # parser.add_argument('--var', '-v', type=int, choices=None, default=1, help='(default: %%(default)s)')
+    # parser.add_argument('--flag', '-f', action='store_true', default=False, help='(default: %%(default)s)')
+    # args = parser.parse_args()
+    # scitex.gen.print_block(args, c='yellow')
+
+    # -----------------------------------
+    # Main
+    # -----------------------------------
+    exit_status = main()
+
+    # -----------------------------------
+    # Cleanup scitex format
+    # -----------------------------------
+    scitex.gen.close(
+        CONFIG,
+        verbose=False,
+        notify=False,
+        message="",
+        exit_status=exit_status,
+    )
+
+# EOF

scitex/ai/genai/deepseek_provider.py

@@ -0,0 +1,251 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-25 12:00:00"
+# Author: Yusuke Watanabe (ywatanabe@alumni.u-tokyo.ac.jp)
+# scitex/src/scitex/ai/genai/deepseek_provider.py
+
+"""
+DeepSeek provider implementation for GenAI.
+
+Provides access to DeepSeek's API for code generation and chat.
+Uses OpenAI-compatible API endpoint.
+"""
+
+from typing import List, Dict, Any, Optional, Generator
+import logging
+import os
+
+from .base_provider import BaseProvider, CompletionResponse, Provider
+from .provider_factory import register_provider
+
+logger = logging.getLogger(__name__)
+
+
+class DeepSeekProvider(BaseProvider):
+    """
+    DeepSeek provider implementation.
+
+    Supports DeepSeek Chat and DeepSeek Coder models.
+    """
+
+    SUPPORTED_MODELS = [
+        "deepseek-chat",
+        "deepseek-coder",
+    ]
+
+    DEFAULT_MODEL = "deepseek-chat"
+
+    def __init__(self, config):
+        """Initialize DeepSeek provider."""
+        self.config = config
+        self.api_key = config.api_key or os.getenv("DEEPSEEK_API_KEY")
+        self.model = config.model or self.DEFAULT_MODEL
+        self.kwargs = config.kwargs or {}
+
+        if not self.api_key:
+            raise ValueError(
+                "DEEPSEEK_API_KEY not provided and not found in environment"
+            )
+
+        # Import OpenAI client (DeepSeek uses OpenAI-compatible API)
+        try:
+            from openai import OpenAI as OpenAIClient
+
+            # DeepSeek uses a custom base URL
+            self.client = OpenAIClient(
+                api_key=self.api_key, base_url="https://api.deepseek.com/beta"
+            )
+        except ImportError:
+            raise ImportError(
+                "OpenAI package not installed. Install with: pip install openai"
+            )
+
+    def complete(self, messages: List[Dict[str, Any]], **kwargs) -> CompletionResponse:
+        """
+        Generate completion using DeepSeek API.
+
+        Args:
+            messages: List of message dictionaries
+            **kwargs: Additional parameters (temperature, max_tokens, etc.)
+
+        Returns:
+            CompletionResponse with generated text and usage info
+        """
+        # Validate messages
+        if not self.validate_messages(messages):
+            raise ValueError("Invalid message format")
+
+        # Format messages for DeepSeek (same as OpenAI format)
+        formatted_messages = self.format_messages(messages)
+
+        # Prepare API parameters
+        api_params = {
+            "model": self.model,
+            "messages": formatted_messages,
+        }
+
+        # Add optional parameters
+        for param in ["temperature", "max_tokens", "top_p", "stop", "seed", "n"]:
+            if param in kwargs:
+                api_params[param] = kwargs[param]
+
+        try:
+            # Make API call
+            response = self.client.chat.completions.create(**api_params)
+
+            # Extract content and usage
+            content = response.choices[0].message.content
+            usage = response.usage
+
+            return CompletionResponse(
+                content=content,
+                input_tokens=usage.prompt_tokens,
+                output_tokens=usage.completion_tokens,
+                finish_reason=response.choices[0].finish_reason,
+                provider_response=response,
+            )
+
+        except Exception as e:
+            logger.error(f"DeepSeek API error: {str(e)}")
+            raise
+
+    def format_messages(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Format messages for DeepSeek API.
+
+        DeepSeek uses the same message format as OpenAI.
+        """
+        formatted_messages = []
+
+        for msg in messages:
+            formatted_msg = {"role": msg["role"], "content": msg["content"]}
+            formatted_messages.append(formatted_msg)
+
+        return formatted_messages
+
+    def validate_messages(self, messages: List[Dict[str, Any]]) -> bool:
+        """Validate message format."""
+        if not messages:
+            return False
+
+        for msg in messages:
+            if not isinstance(msg, dict):
+                return False
+            if "role" not in msg or "content" not in msg:
+                return False
+            if msg["role"] not in ["system", "user", "assistant"]:
+                return False
+
+        return True
+
+    def stream(
+        self, messages: List[Dict[str, Any]], **kwargs
+    ) -> Generator[str, None, CompletionResponse]:
+        """Generate a streaming completion.
+
+        Args:
+            messages: List of messages in standard format
+            **kwargs: Additional parameters
+
+        Yields:
+            Text chunks during streaming
+
+        Returns:
+            Final CompletionResponse when complete
+        """
+        # Validate messages
+        if not self.validate_messages(messages):
+            raise ValueError("Invalid message format")
+
+        # Format messages for DeepSeek
+        formatted_messages = self.format_messages(messages)
+
+        # Prepare API parameters
+        api_params = {
+            "model": self.model,
+            "messages": formatted_messages,
+            "stream": True,
+        }
+
+        # Add optional parameters
+        for param in ["temperature", "max_tokens", "top_p", "stop", "seed", "n"]:
+            if param in kwargs:
+                api_params[param] = kwargs[param]
+
+        try:
+            # Make streaming API call
+            stream = self.client.chat.completions.create(**api_params)
+
+            # Track content and usage
+            full_content = ""
+            prompt_tokens = 0
+            completion_tokens = 0
+
+            for chunk in stream:
+                if chunk.choices[0].delta.content:
+                    content = chunk.choices[0].delta.content
+                    full_content += content
+                    yield content
+
+                # Try to get usage info from chunks (if available)
+                if hasattr(chunk, "usage") and chunk.usage:
+                    if hasattr(chunk.usage, "prompt_tokens"):
+                        prompt_tokens = chunk.usage.prompt_tokens
+                    if hasattr(chunk.usage, "completion_tokens"):
+                        completion_tokens = chunk.usage.completion_tokens
+
+            # If no usage info from stream, estimate
+            if prompt_tokens == 0:
+                prompt_tokens = self.count_tokens(str(formatted_messages))
+            if completion_tokens == 0:
+                completion_tokens = self.count_tokens(full_content)
+
+            # Return final response
+            return CompletionResponse(
+                content=full_content,
+                input_tokens=prompt_tokens,
+                output_tokens=completion_tokens,
+                finish_reason="stop",
+            )
+
+        except Exception as e:
+            logger.error(f"DeepSeek streaming error: {str(e)}")
+            raise
+
+    def count_tokens(self, text: str) -> int:
+        """Count tokens in the given text.
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Number of tokens (estimated)
+        """
+        # DeepSeek uses similar tokenization to GPT models
+        # Rough estimate: 1 token ≈ 4 characters or 0.75 words
+        return len(text.split()) * 4 // 3
+
+    @property
+    def supports_images(self) -> bool:
+        """Check if this provider/model supports image inputs."""
+        # DeepSeek currently doesn't support multimodal inputs
+        return False
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Check if this provider/model supports streaming."""
+        return True
+
+    @property
+    def max_context_length(self) -> int:
+        """Get maximum context length for this model."""
+        # DeepSeek models typically support 4K-16K context
+        context_lengths = {
+            "deepseek-chat": 4096,
+            "deepseek-coder": 16384,
+        }
+        return context_lengths.get(self.model, 4096)
+
+
+# Auto-register when module is imported
+register_provider(Provider.DEEPSEEK.value, DeepSeekProvider)