diff --git a/.gitignore b/.gitignore index 9bb840a..05b9cf4 100644 --- a/.gitignore +++ b/.gitignore @@ -103,3 +103,5 @@ client_keys/ *.p7c *.p7r *.p7s + +test/ \ No newline at end of file diff --git a/nomyo/SecureCompletionClient.py b/nomyo/SecureCompletionClient.py index cdbd8fb..3d1e896 100644 --- a/nomyo/SecureCompletionClient.py +++ b/nomyo/SecureCompletionClient.py @@ -438,13 +438,14 @@ class SecureCompletionClient: return response - async def send_secure_request(self, payload: Dict[str, Any], payload_id: str) -> Dict[str, Any]: + async def send_secure_request(self, payload: Dict[str, Any], payload_id: str, api_key: Optional[str] = None) -> Dict[str, Any]: """ Send a secure chat completion request to the router. Args: payload: Chat completion request payload payload_id: Unique identifier for this request + api_key: Optional API key for bearer authentication Returns: Decrypted response from the LLM @@ -461,6 +462,10 @@ class SecureCompletionClient: "Content-Type": "application/octet-stream" } + # Add Authorization header if api_key is provided + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + # 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 4af1f64..1bc7dec 100644 --- a/nomyo/nomyo.py +++ b/nomyo/nomyo.py @@ -1,5 +1,5 @@ import uuid -from typing import Dict, Any, List +from typing import Dict, Any, List, Optional from .SecureCompletionClient import SecureCompletionClient class SecureChatCompletion: @@ -37,7 +37,7 @@ class SecureChatCompletion: ``` """ - def __init__(self, base_url: str = "https://api.nomyo.ai:12434", allow_http: bool = False): + def __init__(self, base_url: str = "https://api.nomyo.ai:12434", allow_http: bool = False, api_key: Optional[str] = None): """ Initialize the secure chat completion client. @@ -45,10 +45,13 @@ class SecureChatCompletion: base_url: Base URL of the NOMYO Router (must use HTTPS for production) This parameter is named 'base_url' for OpenAI compatibility. allow_http: Allow HTTP connections (ONLY for local development, never in production) + api_key: Optional API key for bearer authentication. If provided, it will be + used for all requests made with this client. """ self.client = SecureCompletionClient(router_url=base_url, allow_http=allow_http) self._keys_initialized = False + self.api_key = api_key async def _ensure_keys(self): """Ensure keys are loaded or generated.""" @@ -139,10 +142,13 @@ class SecureChatCompletion: } # Generate a unique payload ID - payload_id = f"openai-compat-{uuid.uuid4()}" + payload_id = f"{uuid.uuid4()}" + + # 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) + response = await instance.client.send_secure_request(payload, payload_id, request_api_key) return response diff --git a/test.py b/test.py index 5bd2456..040af4a 100644 --- a/test.py +++ b/test.py @@ -9,7 +9,7 @@ the same interface as OpenAI's ChatCompletion.create() method. import asyncio from nomyo import SecureChatCompletion -client = SecureChatCompletion(base_url="http://localhost:12434", allow_http=True) +client = SecureChatCompletion(api_key="test-api-key-12345") async def test_basic_chat(): """Test basic chat completion with OpenAI-style API."""