embed-client 1.0.1.1__py3-none-any.whl → 3.1.0.0__py3-none-any.whl

embed_client/async_client.py

@@ -5,6 +5,12 @@ Async client for Embedding Service API (OpenAPI 3.0.2)
 - English docstrings and examples
 - Ready for PyPi
 - Supports new API format with body, embedding, and chunks
+- Supports all authentication methods (API Key, JWT, Basic Auth, Certificate)
+- Integrates with mcp_security_framework
+- Supports all security modes (HTTP, HTTPS, mTLS)
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
 """
 
 from typing import Any, Dict, List, Optional, Union
@@ -13,6 +19,12 @@ import asyncio
 import os
 import json
 import logging
+from pathlib import Path
+
+# Import authentication, configuration, and SSL systems
+from .auth import ClientAuthManager, create_auth_manager
+from .config import ClientConfig
+from .ssl_manager import ClientSSLManager, create_ssl_manager
 
 class EmbeddingServiceError(Exception):
     """Base exception for EmbeddingServiceAsyncClient."""
@@ -48,19 +60,62 @@ class EmbeddingServiceAsyncClient:
     
     Supports both old and new API formats:
     - Old format: {"result": {"success": true, "data": {"embeddings": [...]}}}
-    - New format: {"result": {"success": true, "data": [{"body": "text", "embedding": [...], "chunks": [...]}]}}
+    - New format: {"result": {"success": true, "data": {"embeddings": [...], "results": [{"body": "text", "embedding": [...], "tokens": [...], "bm25_tokens": [...]}]}}}
+    
+    Supports all authentication methods and security modes:
+    - API Key authentication
+    - JWT token authentication
+    - Basic authentication
+    - Certificate authentication (mTLS)
+    - HTTP, HTTPS, and mTLS security modes
     
     Args:
-        base_url (str): Base URL of the embedding service (e.g., "http://localhost").
-        port (int): Port of the embedding service (e.g., 8001).
+        base_url (str, optional): Base URL of the embedding service (e.g., "http://localhost").
+        port (int, optional): Port of the embedding service (e.g., 8001).
         timeout (float): Request timeout in seconds (default: 30).
+        config (ClientConfig, optional): Configuration object with authentication and SSL settings.
+        config_dict (dict, optional): Configuration dictionary with authentication and SSL settings.
+        auth_manager (ClientAuthManager, optional): Authentication manager instance.
+        
     Raises:
         EmbeddingServiceConfigError: If base_url or port is invalid.
     """
-    def __init__(self, base_url: Optional[str] = None, port: Optional[int] = None, timeout: float = 30.0):
-        # Validate and set base_url
+    def __init__(self, 
+                 base_url: Optional[str] = None, 
+                 port: Optional[int] = None, 
+                 timeout: float = 30.0,
+                 config: Optional[ClientConfig] = None,
+                 config_dict: Optional[Dict[str, Any]] = None,
+                 auth_manager: Optional[ClientAuthManager] = None):
+        # Initialize configuration
+        self.config = config
+        self.config_dict = config_dict
+        self.auth_manager = auth_manager
+        
+        # If config is provided, use it to set base_url and port
+        if config:
+            self.base_url = config.get("server.host", base_url or os.getenv("EMBEDDING_SERVICE_BASE_URL", "http://localhost"))
+            self.port = config.get("server.port", port or int(os.getenv("EMBEDDING_SERVICE_PORT", "8001")))
+            self.timeout = config.get("client.timeout", timeout)
+        elif config_dict:
+            self.base_url = config_dict.get("server", {}).get("host", base_url or os.getenv("EMBEDDING_SERVICE_BASE_URL", "http://localhost"))
+            self.port = config_dict.get("server", {}).get("port", port or int(os.getenv("EMBEDDING_SERVICE_PORT", "8001")))
+            self.timeout = config_dict.get("client", {}).get("timeout", timeout)
+        else:
+            # Use provided parameters or environment variables
+            try:
+                self.base_url = base_url or os.getenv("EMBEDDING_SERVICE_BASE_URL", "http://localhost")
+            except (TypeError, AttributeError) as e:
+                raise EmbeddingServiceConfigError(f"Invalid base_url configuration: {e}") from e
+            
+            try:
+                self.port = port or int(os.getenv("EMBEDDING_SERVICE_PORT", "8001"))
+            except (ValueError, TypeError) as e:
+                raise EmbeddingServiceConfigError(f"Invalid port configuration: {e}") from e
+            self.timeout = timeout
+        
+        # Validate base_url
         try:
-            self.base_url = base_url or os.getenv("EMBEDDING_SERVICE_BASE_URL", "http://localhost")
             if not self.base_url:
                 raise EmbeddingServiceConfigError("base_url must be provided.")
             if not isinstance(self.base_url, str):
@@ -72,10 +127,8 @@ class EmbeddingServiceAsyncClient:
         except (TypeError, AttributeError) as e:
             raise EmbeddingServiceConfigError(f"Invalid base_url configuration: {e}") from e
         
-        # Validate and set port
+        # Validate port
         try:
-            port_env = os.getenv("EMBEDDING_SERVICE_PORT", "8001")
-            self.port = port if port is not None else int(port_env)
             if self.port is None:
                 raise EmbeddingServiceConfigError("port must be provided.")
             if not isinstance(self.port, int) or self.port <= 0 or self.port > 65535:
@@ -85,12 +138,23 @@ class EmbeddingServiceAsyncClient:
         
         # Validate timeout
         try:
-            self.timeout = float(timeout)
+            self.timeout = float(self.timeout)
             if self.timeout <= 0:
                 raise EmbeddingServiceConfigError("timeout must be positive.")
         except (ValueError, TypeError) as e:
             raise EmbeddingServiceConfigError(f"Invalid timeout configuration: {e}") from e
         
+        # Initialize authentication manager if not provided
+        if not self.auth_manager and (self.config or self.config_dict):
+            config_data = self.config_dict if self.config_dict else self.config.get_all()
+            self.auth_manager = create_auth_manager(config_data)
+        
+        # Initialize SSL manager
+        self.ssl_manager = None
+        if self.config or self.config_dict:
+            config_data = self.config_dict if self.config_dict else self.config.get_all()
+            self.ssl_manager = create_ssl_manager(config_data)
+        
         self._session: Optional[aiohttp.ClientSession] = None
 
     def _make_url(self, path: str, base_url: Optional[str] = None, port: Optional[int] = None) -> str:
@@ -172,13 +236,33 @@ class EmbeddingServiceAsyncClient:
             result: API response dictionary
             
         Returns:
-            List of dictionaries with 'body', 'embedding', and 'chunks' fields
+            List of dictionaries with 'body', 'embedding', 'tokens', and 'bm25_tokens' fields
             
         Raises:
             ValueError: If data cannot be extracted or is in old format
         """
         if "result" in result and isinstance(result["result"], dict):
             res = result["result"]
+            if "data" in res and isinstance(res["data"], dict) and "results" in res["data"]:
+                # New format: result.data.results[]
+                results = res["data"]["results"]
+                if isinstance(results, list):
+                    # Validate that all items have required fields
+                    for i, item in enumerate(results):
+                        if not isinstance(item, dict):
+                            raise ValueError(f"Item {i} is not a dictionary: {item}")
+                        if "body" not in item:
+                            raise ValueError(f"Item {i} missing 'body' field: {item}")
+                        if "embedding" not in item:
+                            raise ValueError(f"Item {i} missing 'embedding' field: {item}")
+                        if "tokens" not in item:
+                            raise ValueError(f"Item {i} missing 'tokens' field: {item}")
+                        if "bm25_tokens" not in item:
+                            raise ValueError(f"Item {i} missing 'bm25_tokens' field: {item}")
+                    
+                    return results
+            
+            # Legacy support for old format: result.data[]
             if "data" in res and isinstance(res["data"], list):
                 # Validate that all items have required fields
                 for i, item in enumerate(res["data"]):
@@ -188,8 +272,9 @@ class EmbeddingServiceAsyncClient:
                         raise ValueError(f"Item {i} missing 'body' field: {item}")
                     if "embedding" not in item:
                         raise ValueError(f"Item {i} missing 'embedding' field: {item}")
-                    if "chunks" not in item:
-                        raise ValueError(f"Item {i} missing 'chunks' field: {item}")
+                    # Old format had 'chunks' instead of 'tokens'
+                    if "chunks" not in item and "tokens" not in item:
+                        raise ValueError(f"Item {i} missing 'chunks' or 'tokens' field: {item}")
                 
                 return res["data"]
         
@@ -214,24 +299,72 @@ class EmbeddingServiceAsyncClient:
     def extract_chunks(self, result: Dict[str, Any]) -> List[List[str]]:
         """
         Extract text chunks from API response (new format only).
+        Note: This method now extracts 'tokens' instead of 'chunks' for compatibility.
         
         Args:
             result: API response dictionary
             
         Returns:
-            List of chunk lists for each text
+            List of token lists for each text
             
         Raises:
             ValueError: If chunks cannot be extracted or is in old format
         """
         data = self.extract_embedding_data(result)
-        return [item["chunks"] for item in data]
+        chunks = []
+        for item in data:
+            # New format uses 'tokens', old format used 'chunks'
+            if "tokens" in item:
+                chunks.append(item["tokens"])
+            elif "chunks" in item:
+                chunks.append(item["chunks"])
+            else:
+                raise ValueError(f"Item missing both 'tokens' and 'chunks' fields: {item}")
+        return chunks
+
+    def extract_tokens(self, result: Dict[str, Any]) -> List[List[str]]:
+        """
+        Extract tokens from API response (new format only).
+        
+        Args:
+            result: API response dictionary
+            
+        Returns:
+            List of token lists for each text
+            
+        Raises:
+            ValueError: If tokens cannot be extracted or is in old format
+        """
+        data = self.extract_embedding_data(result)
+        return [item["tokens"] for item in data]
+
+    def extract_bm25_tokens(self, result: Dict[str, Any]) -> List[List[str]]:
+        """
+        Extract BM25 tokens from API response (new format only).
+        
+        Args:
+            result: API response dictionary
+            
+        Returns:
+            List of BM25 token lists for each text
+            
+        Raises:
+            ValueError: If BM25 tokens cannot be extracted or is in old format
+        """
+        data = self.extract_embedding_data(result)
+        return [item["bm25_tokens"] for item in data]
 
     async def __aenter__(self):
         try:
             # Create session with timeout configuration
             timeout = aiohttp.ClientTimeout(total=self.timeout)
-            self._session = aiohttp.ClientSession(timeout=timeout)
+            
+            # Create SSL connector if SSL manager is available
+            connector = None
+            if self.ssl_manager:
+                connector = self.ssl_manager.create_connector()
+            
+            self._session = aiohttp.ClientSession(timeout=timeout, connector=connector)
             return self
         except Exception as e:
             raise EmbeddingServiceError(f"Failed to create HTTP session: {e}") from e
@@ -281,8 +414,9 @@ class EmbeddingServiceAsyncClient:
             dict: Health status and model info.
         """
         url = self._make_url("/health", base_url, port)
+        headers = self.get_auth_headers()
         try:
-            async with self._session.get(url, timeout=self.timeout) as resp:
+            async with self._session.get(url, headers=headers, timeout=self.timeout) as resp:
                 await self._raise_for_status(resp)
                 try:
                     data = await resp.json()
@@ -324,8 +458,9 @@ class EmbeddingServiceAsyncClient:
             dict: OpenAPI schema.
         """
         url = self._make_url("/openapi.json", base_url, port)
+        headers = self.get_auth_headers()
         try:
-            async with self._session.get(url, timeout=self.timeout) as resp:
+            async with self._session.get(url, headers=headers, timeout=self.timeout) as resp:
                 await self._raise_for_status(resp)
                 try:
                     data = await resp.json()
@@ -367,8 +502,9 @@ class EmbeddingServiceAsyncClient:
             dict: List of commands and their descriptions.
         """
         url = self._make_url("/api/commands", base_url, port)
+        headers = self.get_auth_headers()
         try:
-            async with self._session.get(url, timeout=self.timeout) as resp:
+            async with self._session.get(url, headers=headers, timeout=self.timeout) as resp:
                 await self._raise_for_status(resp)
                 try:
                     data = await resp.json()
@@ -470,12 +606,13 @@ class EmbeddingServiceAsyncClient:
 
         logger = logging.getLogger('EmbeddingServiceAsyncClient.cmd')
         url = self._make_url("/cmd", base_url, port)
+        headers = self.get_auth_headers()
         payload = {"command": command}
         if params is not None:
             payload["params"] = params
-        logger.info(f"Sending embedding command: url={url}, payload={payload}")
+        logger.info(f"Sending embedding command: url={url}, payload={payload}, headers={headers}")
         try:
-            async with self._session.post(url, json=payload, timeout=self.timeout) as resp:
+            async with self._session.post(url, json=payload, headers=headers, timeout=self.timeout) as resp:
                 logger.info(f"Embedding service HTTP status: {resp.status}")
                 await self._raise_for_status(resp)
                 try:
@@ -541,4 +678,290 @@ class EmbeddingServiceAsyncClient:
             finally:
                 self._session = None
 
-    # TODO: Add methods for /cmd, /api/commands, etc. 
+    # TODO: Add methods for /cmd, /api/commands, etc.
+    
+    @classmethod
+    def from_config(cls, config: ClientConfig) -> "EmbeddingServiceAsyncClient":
+        """
+        Create client from ClientConfig object.
+        
+        Args:
+            config: ClientConfig object with authentication and SSL settings
+            
+        Returns:
+            EmbeddingServiceAsyncClient instance configured with the provided config
+        """
+        return cls(config=config)
+    
+    @classmethod
+    def from_config_dict(cls, config_dict: Dict[str, Any]) -> "EmbeddingServiceAsyncClient":
+        """
+        Create client from configuration dictionary.
+        
+        Args:
+            config_dict: Configuration dictionary with authentication and SSL settings
+            
+        Returns:
+            EmbeddingServiceAsyncClient instance configured with the provided config
+        """
+        return cls(config_dict=config_dict)
+    
+    @classmethod
+    def from_config_file(cls, config_path: Union[str, Path]) -> "EmbeddingServiceAsyncClient":
+        """
+        Create client from configuration file.
+        
+        Args:
+            config_path: Path to configuration file (JSON or YAML)
+            
+        Returns:
+            EmbeddingServiceAsyncClient instance configured with the provided config
+        """
+        config = ClientConfig.load_config(config_path)
+        return cls(config=config)
+    
+    @classmethod
+    def with_auth(cls, 
+                  base_url: str, 
+                  port: int, 
+                  auth_method: str, 
+                  **kwargs) -> "EmbeddingServiceAsyncClient":
+        """
+        Create client with authentication configuration.
+        
+        Args:
+            base_url: Base URL of the embedding service
+            port: Port of the embedding service
+            auth_method: Authentication method ("api_key", "jwt", "basic", "certificate")
+            **kwargs: Additional authentication parameters
+            
+        Returns:
+            EmbeddingServiceAsyncClient instance with authentication configured
+            
+        Examples:
+            # API Key authentication
+            client = EmbeddingServiceAsyncClient.with_auth(
+                "http://localhost", 8001, "api_key", 
+                api_keys={"user": "api_key_123"}
+            )
+            
+            # JWT authentication
+            client = EmbeddingServiceAsyncClient.with_auth(
+                "http://localhost", 8001, "jwt",
+                secret="secret", username="user", password="pass"
+            )
+            
+            # Basic authentication
+            client = EmbeddingServiceAsyncClient.with_auth(
+                "http://localhost", 8001, "basic",
+                username="user", password="pass"
+            )
+            
+            # Certificate authentication
+            client = EmbeddingServiceAsyncClient.with_auth(
+                "https://localhost", 9443, "certificate",
+                cert_file="certs/client.crt", key_file="keys/client.key"
+            )
+        """
+        # Build configuration dictionary
+        config_dict = {
+            "server": {
+                "host": base_url,
+                "port": port
+            },
+            "client": {
+                "timeout": kwargs.get("timeout", 30.0)
+            },
+            "auth": {
+                "method": auth_method
+            }
+        }
+        
+        # Add authentication configuration based on method
+        if auth_method == "api_key":
+            if "api_keys" in kwargs:
+                config_dict["auth"]["api_keys"] = kwargs["api_keys"]
+            elif "api_key" in kwargs:
+                config_dict["auth"]["api_keys"] = {"user": kwargs["api_key"]}
+            else:
+                raise ValueError("api_keys or api_key parameter required for api_key authentication")
+        
+        elif auth_method == "jwt":
+            required_params = ["secret", "username", "password"]
+            for param in required_params:
+                if param not in kwargs:
+                    raise ValueError(f"{param} parameter required for jwt authentication")
+            
+            config_dict["auth"]["jwt"] = {
+                "secret": kwargs["secret"],
+                "username": kwargs["username"],
+                "password": kwargs["password"],
+                "expiry_hours": kwargs.get("expiry_hours", 24)
+            }
+        
+        elif auth_method == "basic":
+            required_params = ["username", "password"]
+            for param in required_params:
+                if param not in kwargs:
+                    raise ValueError(f"{param} parameter required for basic authentication")
+            
+            config_dict["auth"]["basic"] = {
+                "username": kwargs["username"],
+                "password": kwargs["password"]
+            }
+        
+        elif auth_method == "certificate":
+            required_params = ["cert_file", "key_file"]
+            for param in required_params:
+                if param not in kwargs:
+                    raise ValueError(f"{param} parameter required for certificate authentication")
+            
+            config_dict["auth"]["certificate"] = {
+                "cert_file": kwargs["cert_file"],
+                "key_file": kwargs["key_file"]
+            }
+        
+        else:
+            raise ValueError(f"Unsupported authentication method: {auth_method}")
+        
+        # Add SSL configuration if provided or if using HTTPS
+        ssl_enabled = kwargs.get("ssl_enabled")
+        if ssl_enabled is None:
+            ssl_enabled = base_url.startswith("https://")
+        
+        if ssl_enabled or any(key in kwargs for key in ["ca_cert_file", "cert_file", "key_file", "ssl_enabled"]):
+            config_dict["ssl"] = {
+                "enabled": ssl_enabled,
+                "verify_mode": kwargs.get("verify_mode", "CERT_REQUIRED"),
+                "check_hostname": kwargs.get("check_hostname", True),
+                "check_expiry": kwargs.get("check_expiry", True)
+            }
+            
+            if "ca_cert_file" in kwargs:
+                config_dict["ssl"]["ca_cert_file"] = kwargs["ca_cert_file"]
+            
+            if "cert_file" in kwargs:
+                config_dict["ssl"]["cert_file"] = kwargs["cert_file"]
+            
+            if "key_file" in kwargs:
+                config_dict["ssl"]["key_file"] = kwargs["key_file"]
+        
+        return cls(config_dict=config_dict)
+    
+    def get_auth_headers(self) -> Dict[str, str]:
+        """
+        Get authentication headers for requests.
+        
+        Returns:
+            Dictionary of authentication headers
+        """
+        if not self.auth_manager:
+            return {}
+        
+        auth_method = self.auth_manager.get_auth_method()
+        if auth_method == "none":
+            return {}
+        
+        # Get authentication parameters from config
+        auth_config = self.config_dict.get("auth", {}) if self.config_dict else {}
+        if self.config:
+            auth_config = self.config.get("auth", {})
+        
+        if auth_method == "api_key":
+            api_keys = auth_config.get("api_keys", {})
+            # Use first available API key
+            for user_id, api_key in api_keys.items():
+                return self.auth_manager.get_auth_headers("api_key", api_key=api_key)
+        
+        elif auth_method == "jwt":
+            jwt_config = auth_config.get("jwt", {})
+            username = jwt_config.get("username")
+            password = jwt_config.get("password")
+            if username and password:
+                # Create JWT token
+                token = self.auth_manager.create_jwt_token(username, ["user"])
+                return self.auth_manager.get_auth_headers("jwt", token=token)
+        
+        elif auth_method == "basic":
+            basic_config = auth_config.get("basic", {})
+            username = basic_config.get("username")
+            password = basic_config.get("password")
+            if username and password:
+                return self.auth_manager.get_auth_headers("basic", username=username, password=password)
+        
+        return {}
+    
+    def is_authenticated(self) -> bool:
+        """
+        Check if client is configured for authentication.
+        
+        Returns:
+            True if authentication is configured, False otherwise
+        """
+        return self.auth_manager is not None and self.auth_manager.is_auth_enabled()
+    
+    def get_auth_method(self) -> str:
+        """
+        Get current authentication method.
+        
+        Returns:
+            Authentication method name or "none" if not configured
+        """
+        if not self.auth_manager:
+            return "none"
+        return self.auth_manager.get_auth_method()
+    
+    def is_ssl_enabled(self) -> bool:
+        """
+        Check if SSL/TLS is enabled.
+        
+        Returns:
+            True if SSL/TLS is enabled, False otherwise
+        """
+        if not self.ssl_manager:
+            return False
+        return self.ssl_manager.is_ssl_enabled()
+    
+    def is_mtls_enabled(self) -> bool:
+        """
+        Check if mTLS (mutual TLS) is enabled.
+        
+        Returns:
+            True if mTLS is enabled, False otherwise
+        """
+        if not self.ssl_manager:
+            return False
+        return self.ssl_manager.is_mtls_enabled()
+    
+    def get_ssl_config(self) -> Dict[str, Any]:
+        """
+        Get current SSL configuration.
+        
+        Returns:
+            Dictionary with SSL configuration or empty dict if not configured
+        """
+        if not self.ssl_manager:
+            return {}
+        return self.ssl_manager.get_ssl_config()
+    
+    def validate_ssl_config(self) -> List[str]:
+        """
+        Validate SSL configuration.
+        
+        Returns:
+            List of validation errors
+        """
+        if not self.ssl_manager:
+            return []
+        return self.ssl_manager.validate_ssl_config()
+    
+    def get_supported_ssl_protocols(self) -> List[str]:
+        """
+        Get list of supported SSL/TLS protocols.
+        
+        Returns:
+            List of supported protocol names
+        """
+        if not self.ssl_manager:
+            return []
+        return self.ssl_manager.get_supported_protocols()