scitex 2.0.0__py2.py3-none-any.whl

scitex/ai/genai/auth_manager.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2025-05-31 10:00:00"
+# Author: ywatanabe
+# File: ./src/scitex/ai/genai/auth_manager.py
+
+"""
+Handles API key management and validation for AI providers.
+
+This module provides secure handling of API keys including:
+- Environment variable retrieval
+- Key validation
+- Masked key display for security
+"""
+
+import os
+from typing import Optional, Dict, Any
+
+
+class AuthManager:
+    """Manages API key authentication for AI providers.
+
+    Example
+    -------
+    >>> auth = AuthManager("sk-abc123...", "OpenAI")
+    >>> auth.validate_key()
+    True
+    >>> auth.get_masked_key()
+    'sk-a****3...'
+
+    Parameters
+    ----------
+    api_key : str
+        The API key for authentication
+    provider : str
+        The provider name (e.g., "OpenAI", "Anthropic")
+    """
+
+    # Mapping of providers to environment variable names
+    ENV_VAR_MAPPING: Dict[str, str] = {
+        "OpenAI": "OPENAI_API_KEY",
+        "Anthropic": "ANTHROPIC_API_KEY",
+        "Google": "GOOGLE_API_KEY",
+        "Groq": "GROQ_API_KEY",
+        "DeepSeek": "DEEPSEEK_API_KEY",
+        "Perplexity": "PERPLEXITY_API_KEY",
+        "Llama": "LLAMA_API_KEY",
+    }
+
+    def __init__(self, api_key: Optional[str], provider: str):
+        """Initialize AuthManager with API key and provider.
+
+        Parameters
+        ----------
+        api_key : Optional[str]
+            The API key. If None, will attempt to get from environment
+        provider : str
+            The provider name
+        """
+        # Normalize provider name to lowercase
+        self.provider = provider.lower()
+
+        # Check if provider is known
+        if self.provider not in [p.lower() for p in self.ENV_VAR_MAPPING.keys()]:
+            raise ValueError(f"Unknown provider: {provider}")
+
+        # Store private api key
+        self._api_key = api_key or self.get_key_from_env(provider)
+
+        if not self._api_key:
+            raise ValueError(f"No API key provided for {provider}")
+
+        # Public property for backward compatibility
+        self.api_key = self._api_key
+
+    def validate(self) -> bool:
+        """Validate the API key format.
+
+        Returns
+        -------
+        bool
+            True if key appears valid, False otherwise
+
+        Raises
+        ------
+        ValueError
+            If API key is missing or invalid
+        """
+        if not self._api_key:
+            raise ValueError("No API key configured")
+
+        # Basic validation - ensure key is not empty and has reasonable length
+        if len(self._api_key) < 10:
+            raise ValueError(
+                f"API key for {self.provider} appears to be invalid (too short)"
+            )
+
+        return True
+
+    def get_masked_key(self) -> str:
+        """Get a masked version of the API key for display.
+
+        Returns
+        -------
+        str
+            Masked API key showing only first few and last few characters
+        """
+        if not self._api_key:
+            return "No API key"
+
+        if len(self._api_key) <= 8:
+            return "*****"
+
+        # Match test expectation: show first 3 and last 4 characters with "..." in between
+        return f"{self._api_key[:3]}...{self._api_key[-4:]}"
+
+    def get_client_config(self) -> Dict[str, Any]:
+        """Get client configuration for the provider.
+
+        Returns
+        -------
+        Dict[str, any]
+            Configuration dictionary for the provider client
+        """
+        config = {"api_key": self._api_key}
+
+        # Provider-specific configurations
+        if self.provider == "openai":
+            # Check for optional organization
+            org = os.getenv("OPENAI_ORGANIZATION")
+            if org:
+                config["organization"] = org
+        elif self.provider == "anthropic":
+            config["max_retries"] = 3
+        elif self.provider == "google":
+            # Google only needs API key
+            pass
+
+        return config
+
+    @classmethod
+    def get_key_from_env(cls, provider: str) -> Optional[str]:
+        """Get API key from environment variable.
+
+        Parameters
+        ----------
+        provider : str
+            The provider name
+
+        Returns
+        -------
+        Optional[str]
+            The API key if found in environment, None otherwise
+        """
+        # Find the env var name case-insensitively
+        for p, env_var in cls.ENV_VAR_MAPPING.items():
+            if p.lower() == provider.lower():
+                return os.getenv(env_var)
+
+        # Try generic pattern if not found
+        env_var = f"{provider.upper()}_API_KEY"
+        return os.getenv(env_var)
+
+    def get_api_key(self, provider: str, api_key: Optional[str] = None) -> str:
+        """Get API key for a provider.
+
+        This method is used by provider_base.py for compatibility.
+
+        Parameters
+        ----------
+        provider : str
+            The provider name
+        api_key : Optional[str]
+            Explicit API key, if None will use environment
+
+        Returns
+        -------
+        str
+            The API key
+
+        Raises
+        ------
+        ValueError
+            If no API key is found
+        """
+        if api_key:
+            return api_key
+
+        key = self.get_key_from_env(provider)
+        if not key:
+            raise ValueError(f"No API key provided for {provider}")
+
+        return key
+
+    def __repr__(self) -> str:
+        """String representation of AuthManager."""
+        return f"AuthManager(provider={self.provider}, key={self.get_masked_key()})"
+
+
+# EOF

scitex/ai/genai/base_genai.py

@@ -0,0 +1,336 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-05-03 11:55:54 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_repo/src/scitex/ai/_gen_ai/_BaseGenAI.py
+# ----------------------------------------
+import os
+
+__FILE__ = "./src/scitex/ai/_gen_ai/_BaseGenAI.py"
+__DIR__ = os.path.dirname(__FILE__)
+# ----------------------------------------
+
+import base64
+import sys
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generator, List, Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ...io._load import load
+from .calc_cost import calc_cost
+from .format_output_func import format_output_func
+from .params import MODELS
+
+
+class BaseGenAI(ABC):
+    def __init__(
+        self,
+        system_setting: str = "",
+        model: str = "",
+        api_key: str = "",
+        stream: bool = False,
+        seed: Optional[int] = None,
+        n_keep: int = 1,
+        temperature: float = 1.0,
+        provider: str = "",
+        chat_history: Optional[List[Dict[str, str]]] = None,
+        max_tokens: int = 4_096,
+    ) -> None:
+        self.provider = provider
+        self.system_setting = system_setting
+        self.model = model
+        self.api_key = api_key
+        self.stream = stream
+        self.seed = seed
+        self.n_keep = n_keep
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self._error_messages: List[str] = []
+
+        self.reset(system_setting)
+        self.history = chat_history if chat_history else []
+
+        try:
+            self.verify_model()
+            self.client = self._init_client()
+        except Exception as error:
+            print(error)
+            self._error_messages.append(f"\nError:\n{str(error)}")
+
+    @classmethod
+    def list_models(cls, provider: Optional[str] = None) -> List[str]:
+        """List available models for the provider. If provider is None, list all models."""
+        if provider:
+            indi = [
+                provider.lower() in api_key_env.lower()
+                for api_key_env in MODELS["api_key_env"]
+            ]
+            models = MODELS[indi].name.tolist()
+            providers = MODELS[indi].provider.tolist()
+
+        else:
+            indi = np.arange(len(MODELS))
+            models = MODELS.name.tolist()
+            providers = MODELS.provider.tolist()
+
+        for provider, model in zip(providers, models):
+            print(f"- {provider} - {model}")
+
+        return models
+
+    def gen_error(
+        self, return_stream: bool
+    ) -> tuple[bool, Optional[Union[str, Generator]]]:
+        error_exists = bool(self._error_messages)
+        if not error_exists:
+            return False, None
+
+        error_msgs = self._error_messages
+        self._error_messages = []
+
+        if not self.stream:
+            return True, "".join(error_msgs)
+
+        stream_obj = self._to_stream(error_msgs)
+        return True, (
+            self._yield_stream(stream_obj) if not return_stream else stream_obj
+        )
+
+    def __call__(
+        self,
+        prompt: Optional[str] = None,
+        prompt_file: Optional[str] = None,
+        images: List[Any] = None,
+        format_output: bool = False,
+        return_stream: bool = False,
+    ) -> Union[str, Generator]:
+
+        # ----------------------------------------
+        # Handles Prompt and Prompt File
+        if (not prompt) and (not prompt_file):
+            print("Please input prompt\n")
+            return
+
+        if prompt_file:
+            file_content = load(prompt_file)
+            # Escape special characters
+            escaped_content = [repr(line)[1:-1] for line in file_content]
+            prompt = (
+                str(prompt).strip() + "\n\n" + str("\n".join(escaped_content)).strip()
+            )
+
+        # if prompt_file:
+        #     prompt = (
+        #         str(prompt).strip()
+        #         + "\n\n"
+        #         + str("\n".join(load(prompt_file))).strip()
+        #     )
+
+        if prompt.strip() == "":
+            print("Please input prompt\n")
+            return
+        # ----------------------------------------
+
+        self.update_history("user", prompt or "", images=images)
+
+        error_flag, error_obj = self.gen_error(return_stream)
+        if error_flag:
+            return error_obj
+
+        try:
+            if not self.stream:
+                return self._call_static(format_output)
+
+            if return_stream:
+                self.stream, orig_stream = return_stream, self.stream
+                stream_obj = self._call_stream(format_output)
+                self.stream = orig_stream
+                return stream_obj
+
+            return self._yield_stream(self._call_stream(format_output))
+
+        except Exception as error:
+            self._error_messages.append(f"\nError:\n{str(error)}")
+            error_flag, error_obj = self.gen_error(return_stream)
+            if error_flag:
+                return error_obj
+
+    def _yield_stream(self, stream_obj: Generator) -> str:
+        accumulated = []
+        for chunk in stream_obj:
+            if chunk:
+                sys.stdout.write(chunk)
+                sys.stdout.flush()
+                accumulated.append(chunk)
+        result = "".join(accumulated)
+        self.update_history("assistant", result)
+        return result
+
+    def _call_static(self, format_output: bool = True) -> str:
+        out_text = self._api_call_static()
+        out_text = format_output_func(out_text) if format_output else out_text
+        self.update_history("assistant", out_text)
+        return out_text
+
+    def _call_stream(self, format_output: Optional[bool] = None) -> Generator:
+        return self._api_call_stream()
+
+    @abstractmethod
+    def _init_client(self) -> Any:
+        """Returns client"""
+        pass
+
+    def _api_format_history(self, history):
+        """Returns chat_history by handling differences in API expectations"""
+        return history
+
+    @abstractmethod
+    def _api_call_static(self) -> str:
+        """Returns out_text by handling differences in API expectations"""
+        pass
+
+    @abstractmethod
+    def _api_call_stream(self) -> Generator:
+        """Returns stream by handling differences in API expectations"""
+        pass
+
+    def _get_available_models(self) -> List[str]:
+        indi = [
+            self.provider.lower() in api_key_env.lower()
+            for api_key_env in MODELS["api_key_env"]
+        ]
+        return MODELS[indi].name.tolist()
+
+    @property
+    def available_models(self) -> List[str]:
+        return self._get_available_models()
+
+    def reset(self, system_setting: str = "") -> None:
+        self.history = []
+        if system_setting:
+            self.history.append({"role": "system", "content": system_setting})
+
+    def _ensure_alternative_history(
+        self, history: List[Dict[str, str]]
+    ) -> List[Dict[str, str]]:
+        if len(history) < 2:
+            return history
+
+        if history[-1]["role"] == history[-2]["role"]:
+            last_content = history.pop()["content"]
+            history[-1]["content"] += f"\n\n{last_content}"
+            return self._ensure_alternative_history(history)
+
+        return history
+
+    @staticmethod
+    def _ensure_start_from_user(history: List[Dict[str, str]]) -> List[Dict[str, str]]:
+        if history and history[0]["role"] != "user":
+            history.pop(0)
+        return history
+
+    @staticmethod
+    def _ensure_base64_encoding(image, max_size=512):
+        import io
+
+        from PIL import Image
+
+        def resize_image(img):
+            # Calculate new dimensions while maintaining aspect ratio
+            ratio = max_size / max(img.size)
+            if ratio < 1:
+                new_size = tuple(int(dim * ratio) for dim in img.size)
+                img = img.resize(new_size, Image.Resampling.LANCZOS)
+            return img
+
+        if isinstance(image, str):
+            try:
+                # Try to open and resize as file path
+                img = Image.open(image)
+                img = resize_image(img)
+                buffer = io.BytesIO()
+                img.save(buffer, format="JPEG")
+                return base64.b64encode(buffer.getvalue()).decode("utf-8")
+            except:
+                # If fails, assume it's already base64 string
+                return image
+        elif isinstance(image, bytes):
+            # Convert bytes to image, resize, then back to base64
+            img = Image.open(io.BytesIO(image))
+            img = resize_image(img)
+            buffer = io.BytesIO()
+            img.save(buffer, format="JPEG")
+            return base64.b64encode(buffer.getvalue()).decode("utf-8")
+        else:
+            raise ValueError("Unsupported image format")
+
+    def update_history(self, role: str, content: str, images=None) -> None:
+        if images is not None:
+            content = [
+                {"type": "text", "text": content},
+                *[
+                    {
+                        "type": "_image",
+                        "_image": self._ensure_base64_encoding(image),
+                    }
+                    for image in images
+                ],
+            ]
+
+        self.history.append({"role": role, "content": content})
+
+        if len(self.history) > self.n_keep:
+            self.history = self.history[-self.n_keep :]
+
+        self.history = self._ensure_alternative_history(self.history)
+        self.history = self._ensure_start_from_user(self.history)
+        self.history = self._api_format_history(self.history)
+
+    def verify_model(self) -> None:
+        if self.model not in self.available_models:
+            message = (
+                f"Specified model {self.model} is not supported for the API Key ({self.masked_api_key}). "
+                f"Available models for {str(self)} are as follows:\n{self.available_models}"
+            )
+            raise ValueError(message)
+
+    @property
+    def masked_api_key(self) -> str:
+        return f"{self.api_key[:4]}****{self.api_key[-4:]}"
+
+    def _add_masked_api_key(self, text: str) -> str:
+        return text + f"\n(API Key: {self.masked_api_key}"
+
+    @property
+    def cost(self) -> float:
+        return calc_cost(self.model, self.input_tokens, self.output_tokens)
+
+    @staticmethod
+    def _to_stream(string: Union[str, List[str]]) -> Generator[str, None, None]:
+        """Converts string or list of strings to generator for streaming."""
+        chunks = string if isinstance(string, list) else [string]
+        for chunk in chunks:
+            if chunk:
+                yield chunk
+
+
+def main() -> None:
+    pass
+
+
+if __name__ == "__main__":
+    import scitex
+
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt, verbose=False)
+    main()
+    scitex.gen.close(CONFIG, verbose=False, notify=False)
+
+"""
+python src/scitex/ai/_gen_ai/_BaseGenAI.py
+python -m src.scitex.ai._gen_ai._BaseGenAI
+"""
+
+# EOF