diff --git a/nomyo/SecureCompletionClient.py b/nomyo/SecureCompletionClient.py index 8f2be54..27d3863 100644 --- a/nomyo/SecureCompletionClient.py +++ b/nomyo/SecureCompletionClient.py @@ -616,7 +616,7 @@ class SecureCompletionClient: return response - async def send_secure_request(self, payload: Dict[str, Any], payload_id: str, api_key: Optional[str] = None) -> Dict[str, Any]: + async def send_secure_request(self, payload: Dict[str, Any], payload_id: str, api_key: Optional[str] = None, security_tier: Optional[str] = None) -> Dict[str, Any]: """ Send a secure chat completion request to the router. @@ -624,6 +624,12 @@ class SecureCompletionClient: payload: Chat completion request payload payload_id: Unique identifier for this request api_key: Optional API key for bearer authentication + security_tier: Optional security tier for routing ("standard", "high", or "maximum"). + Controls hardware preference: + - "standard": general secure inference + - "high": sensitive business data + - "maximum": maximum isolation (PHI, classified data) + If not specified, server uses default based on model name mapping. Returns: Decrypted response from the LLM @@ -634,9 +640,19 @@ class SecureCompletionClient: APIError: For other HTTP errors APIConnectionError: If connection fails SecurityError: If encryption/decryption fails + ValueError: If security_tier is invalid """ logger.info("Sending secure chat completion request...") + # Validate security tier if provided + if security_tier is not None: + valid_tiers = ["standard", "high", "maximum"] + if security_tier not in valid_tiers: + raise ValueError( + f"Invalid security_tier: '{security_tier}'. " + f"Must be one of: {', '.join(valid_tiers)}" + ) + # Step 1: Encrypt the payload encrypted_payload = await self.encrypt_payload(payload) @@ -651,6 +667,10 @@ class SecureCompletionClient: if api_key: headers["Authorization"] = f"Bearer {api_key}" + # Add Security-Tier header if security_tier is provided + if security_tier: + headers["X-Security-Tier"] = security_tier + # Step 3: Send request to router url = f"{self.router_url}/v1/chat/secure_completion" logger.debug("Target URL: %s", url) diff --git a/nomyo/nomyo.py b/nomyo/nomyo.py index bfd02d0..59a61d3 100644 --- a/nomyo/nomyo.py +++ b/nomyo/nomyo.py @@ -120,6 +120,12 @@ class SecureChatCompletion: - logit_bias: Dict[str, float] - user: str - base_url: str (alternative to initializing with router_url) + - security_tier: str ("standard", "high", or "maximum") + Controls hardware routing and security level: + * "standard": general secure inference + * "high": sensitive business data + * "maximum": maximum isolation (PHI, classified data) + If not specified, server uses default based on model name mapping. Returns: A dictionary containing the chat completion response with the following structure: @@ -153,6 +159,9 @@ class SecureChatCompletion: """ # Extract base_url if provided (OpenAI compatibility) base_url = kwargs.pop("base_url", None) + + # Extract security_tier if provided + security_tier = kwargs.pop("security_tier", None) # Use the instance's client unless base_url is explicitly overridden if base_url is not None: @@ -179,8 +188,8 @@ class SecureChatCompletion: # Use instance's api_key if not overridden in kwargs request_api_key = kwargs.pop("api_key", instance.api_key) - # Send secure request - response = await instance.client.send_secure_request(payload, payload_id, request_api_key) + # Send secure request with security tier + response = await instance.client.send_secure_request(payload, payload_id, request_api_key, security_tier) return response