Hello World!
This commit is contained in:
commit
5c27f51941
9 changed files with 1384 additions and 0 deletions
201
LICENSE
Normal file
201
LICENSE
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright 2025 OpenAI
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
3
MANIFEST.in
Normal file
3
MANIFEST.in
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
include README.md
|
||||
include LICENSE
|
||||
recursive-include nomyo *
|
||||
341
README.md
Normal file
341
README.md
Normal file
|
|
@ -0,0 +1,341 @@
|
|||
# NOMYO Secure Python Chat Client
|
||||
|
||||
**OpenAI-compatible secure chat client with end-to-end encryption with NOMYO Inference Endpoints**
|
||||
|
||||
🔒 **All prompts and responses are automatically encrypted and decrypted**
|
||||
🔑 **Uses hybrid encryption (AES-256-GCM + RSA-OAEP with 4096-bit keys)**
|
||||
🔄 **Drop-in replacement for OpenAI's ChatCompletion API**
|
||||
|
||||
## 🚀 Quick Start
|
||||
|
||||
### 1. Install dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 2. Use the client (same API as OpenAI)
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
async def main():
|
||||
# Initialize client (defaults to http://api.nomyo.ai:12434)
|
||||
client = SecureChatCompletion(base_url="http://api.nomyo.ai:12434")
|
||||
|
||||
# Simple chat completion
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello! How are you today?"}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response['choices'][0]['message']['content'])
|
||||
|
||||
# Run the async function
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### 3. Run tests
|
||||
|
||||
```bash
|
||||
python3 test.py
|
||||
```
|
||||
|
||||
## 🔐 Security Features
|
||||
|
||||
### Hybrid Encryption
|
||||
|
||||
- **Payload encryption**: AES-256-GCM (authenticated encryption)
|
||||
- **Key exchange**: RSA-OAEP with SHA-256
|
||||
- **Key size**: 4096-bit RSA keys
|
||||
- **All communication**: End-to-end encrypted
|
||||
|
||||
### Key Management
|
||||
|
||||
- Automatic key generation and management
|
||||
- Keys stored with restricted permissions (600 for private key)
|
||||
- Optional password protection for private keys
|
||||
- Key persistence across sessions
|
||||
|
||||
## 🔄 OpenAI Compatibility
|
||||
|
||||
The `SecureChatCompletion` class provides **exact API compatibility** with OpenAI's `ChatCompletion.create()` method.
|
||||
|
||||
### Supported Parameters
|
||||
|
||||
All standard OpenAI parameters are supported:
|
||||
|
||||
- `model`: Model identifier
|
||||
- `messages`: List of message objects
|
||||
- `temperature`: Sampling temperature (0-2)
|
||||
- `max_tokens`: Maximum tokens to generate
|
||||
- `top_p`: Nucleus sampling
|
||||
- `frequency_penalty`: Frequency penalty
|
||||
- `presence_penalty`: Presence penalty
|
||||
- `stop`: Stop sequences
|
||||
- `n`: Number of completions
|
||||
- `stream`: Streaming (not yet implemented)
|
||||
- `tools`: Tool definitions
|
||||
- `tool_choice`: Tool selection strategy
|
||||
- `user`: User identifier
|
||||
- And more...
|
||||
|
||||
### Response Format
|
||||
|
||||
Responses follow the OpenAI format exactly, with an additional `_metadata` field for debugging and security information:
|
||||
|
||||
```python
|
||||
{
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! I'm doing well, thank you for asking.",
|
||||
"tool_calls": [...] # if tools were used
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 30
|
||||
},
|
||||
"_metadata": {
|
||||
"payload_id": "openai-compat-abc123", # Unique identifier for this request
|
||||
"processed_at": 1765250382, # Timestamp when server processed the request
|
||||
"is_encrypted": True, # Indicates this response was decrypted
|
||||
"encryption_algorithm": "hybrid-aes256-rsa4096", # Encryption method used
|
||||
"response_status": "success" # Status of the decryption/processing
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The `_metadata` field contains security-related information about the encrypted communication and is automatically added to all responses.
|
||||
|
||||
## 🛠️ Usage Examples
|
||||
|
||||
### Basic Chat
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
async def main():
|
||||
client = SecureChatCompletion(base_url="http://api.nomyo.ai:12434")
|
||||
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response['choices'][0]['message']['content'])
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### With Tools
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
async def main():
|
||||
client = SecureChatCompletion(base_url="http://api.nomyo.ai:12434")
|
||||
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "What's the weather in Paris?"}
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {"type": "string"}
|
||||
},
|
||||
"required": ["location"]
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response['choices'][0]['message']['content'])
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Using acreate() Alias
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
async def main():
|
||||
client = SecureChatCompletion(base_url="http://api.nomyo.ai:12434")
|
||||
|
||||
response = await client.acreate(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello!"}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response['choices'][0]['message']['content'])
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
## 📦 Dependencies
|
||||
|
||||
See `requirements.txt` for the complete list:
|
||||
|
||||
- `cryptography`: Cryptographic primitives (RSA, AES, etc.)
|
||||
- `httpx`: Async HTTP client
|
||||
- `anyio`: Async compatibility layer
|
||||
|
||||
## 🔧 Configuration
|
||||
|
||||
### Custom Base URL
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
async def main():
|
||||
client = SecureChatCompletion(base_url="http://NOMYO-Pro-Router:12434")
|
||||
# ... rest of your code
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Key Management
|
||||
|
||||
Keys are automatically generated on first use and stored in `client_keys/` directory.
|
||||
|
||||
#### Generate Keys Manually
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo.SecureCompletionClient import SecureCompletionClient
|
||||
|
||||
async def main():
|
||||
client = SecureCompletionClient()
|
||||
await client.generate_keys(save_to_file=True, password="your-password")
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
#### Load Existing Keys
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from nomyo.SecureCompletionClient import SecureCompletionClient
|
||||
|
||||
async def main():
|
||||
client = SecureCompletionClient()
|
||||
await client.load_keys("client_keys/private_key.pem", "client_keys/public_key.pem", password="your-password")
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
## 🧪 Testing
|
||||
|
||||
Run the comprehensive test suite:
|
||||
|
||||
```bash
|
||||
python3 test.py
|
||||
```
|
||||
|
||||
Tests verify:
|
||||
|
||||
- ✅ OpenAI API compatibility
|
||||
- ✅ Basic chat completion
|
||||
- ✅ Tool usage
|
||||
- ✅ All OpenAI parameters
|
||||
- ✅ Async methods
|
||||
- ✅ Error handling
|
||||
|
||||
## 📚 API Reference
|
||||
|
||||
### SecureChatCompletion
|
||||
|
||||
#### Constructor
|
||||
|
||||
```python
|
||||
SecureChatCompletion(base_url: str = "http://api.nomyo.ai:12434")
|
||||
```
|
||||
|
||||
#### Methods
|
||||
|
||||
- `create(model, messages, **kwargs)`: Create a chat completion
|
||||
- `acreate(model, messages, **kwargs)`: Async alias for create()
|
||||
|
||||
### SecureCompletionClient
|
||||
|
||||
#### Constructor
|
||||
|
||||
```python
|
||||
SecureCompletionClient(router_url: str = "http://api.nomyo.ai:12434")
|
||||
```
|
||||
|
||||
#### Methods
|
||||
|
||||
- `generate_keys(save_to_file=False, key_dir="client_keys", password=None)`: Generate RSA key pair
|
||||
- `load_keys(private_key_path, public_key_path=None, password=None)`: Load keys from files
|
||||
- `fetch_server_public_key()`: Fetch server's public key
|
||||
- `encrypt_payload(payload)`: Encrypt a payload
|
||||
- `decrypt_response(encrypted_response, payload_id)`: Decrypt a response
|
||||
- `send_secure_request(payload, payload_id)`: Send encrypted request and receive decrypted response
|
||||
|
||||
## 📝 Notes
|
||||
|
||||
### Security Best Practices
|
||||
|
||||
- Always use password protection for private keys in production
|
||||
- Keep private keys secure (permissions set to 600)
|
||||
- Never share your private key
|
||||
- Verify server's public key fingerprint before first use
|
||||
|
||||
### Performance
|
||||
|
||||
- Key generation takes ~1-2 seconds (one-time operation)
|
||||
- Encryption/decryption adds minimal overhead (~10-20ms per request)
|
||||
|
||||
### Compatibility
|
||||
|
||||
- Works with any OpenAI-compatible code
|
||||
- No changes needed to existing OpenAI client code
|
||||
- Simply replace `openai.ChatCompletion.create()` with `SecureChatCompletion.create()`
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! Please open issues or pull requests on the project repository.
|
||||
|
||||
## 📄 License
|
||||
|
||||
See LICENSE file for licensing information.
|
||||
|
||||
## 📞 Support
|
||||
|
||||
For questions or issues, please refer to the project documentation or open an issue.
|
||||
382
nomyo/SecureCompletionClient.py
Normal file
382
nomyo/SecureCompletionClient.py
Normal file
|
|
@ -0,0 +1,382 @@
|
|||
import json, base64, urllib.parse, httpx, os
|
||||
from typing import Dict, Any, Optional
|
||||
from cryptography.hazmat.primitives import serialization, hashes
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa, padding
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
||||
|
||||
|
||||
class SecureCompletionClient:
|
||||
"""
|
||||
Client for the /v1/chat/secure_completion endpoint.
|
||||
|
||||
Handles:
|
||||
- Key generation and management
|
||||
- Hybrid encryption/decryption
|
||||
- API communication
|
||||
- Response parsing
|
||||
"""
|
||||
|
||||
def __init__(self, router_url: str = "http://api.nomyo.ai:12434"):
|
||||
"""
|
||||
Initialize the secure completion client.
|
||||
|
||||
Args:
|
||||
router_url: Base URL of the NOMYO Router (e.g., "http://api.nomyo.ai:12434")
|
||||
"""
|
||||
self.router_url = router_url.rstrip('/')
|
||||
self.private_key = None
|
||||
self.public_key_pem = None
|
||||
self.key_size = 4096 # RSA key size
|
||||
|
||||
async def generate_keys(self, save_to_file: bool = False, key_dir: str = "client_keys", password: Optional[str] = None) -> None:
|
||||
"""
|
||||
Generate RSA key pair for secure communication.
|
||||
|
||||
Args:
|
||||
save_to_file: Whether to save keys to files
|
||||
key_dir: Directory to save keys (if save_to_file is True)
|
||||
password: Optional password to encrypt private key (recommended for production)
|
||||
"""
|
||||
print("🔑 Generating RSA key pair...")
|
||||
|
||||
# Generate private key
|
||||
self.private_key = rsa.generate_private_key(
|
||||
public_exponent=65537,
|
||||
key_size=self.key_size,
|
||||
backend=default_backend()
|
||||
)
|
||||
|
||||
# Get public key
|
||||
public_key = self.private_key.public_key()
|
||||
|
||||
# Serialize public key to PEM format
|
||||
self.public_key_pem = public_key.public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
||||
).decode('utf-8')
|
||||
|
||||
print(f" ✓ Generated {self.key_size}-bit RSA key pair")
|
||||
|
||||
if save_to_file:
|
||||
os.makedirs(key_dir, exist_ok=True)
|
||||
|
||||
# Save private key
|
||||
if password:
|
||||
# Encrypt private key with user-provided password
|
||||
private_pem = self.private_key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.BestAvailableEncryption(password.encode('utf-8'))
|
||||
)
|
||||
print(f" ✓ Private key encrypted with password")
|
||||
else:
|
||||
# Save unencrypted for convenience (not recommended for production)
|
||||
private_pem = self.private_key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption()
|
||||
)
|
||||
print(f" ⚠️ Private key saved UNENCRYPTED (not recommended for production)")
|
||||
|
||||
# Write private key with restricted permissions (readable only by owner)
|
||||
private_key_path = os.path.join(key_dir, "private_key.pem")
|
||||
with open(private_key_path, "wb") as f:
|
||||
f.write(private_pem)
|
||||
try:
|
||||
os.chmod(private_key_path, 0o600) # Only owner can read/write
|
||||
print(f" ✓ Private key permissions set to 600 (owner-only access)")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ Could not set private key permissions: {e}")
|
||||
|
||||
# Save public key (always unencrypted, but with restricted permissions)
|
||||
public_key_path = os.path.join(key_dir, "public_key.pem")
|
||||
with open(public_key_path, "w") as f:
|
||||
f.write(self.public_key_pem)
|
||||
try:
|
||||
os.chmod(public_key_path, 0o644) # Owner read/write, group/others read
|
||||
print(f" ✓ Public key permissions set to 644")
|
||||
except Exception as e:
|
||||
print(f" ⚠️ Could not set public key permissions: {e}")
|
||||
|
||||
print(f" ✓ Keys saved to {key_dir}/")
|
||||
|
||||
async def load_keys(self, private_key_path: str, public_key_path: Optional[str] = None, password: Optional[str] = None) -> None:
|
||||
"""
|
||||
Load RSA keys from files.
|
||||
|
||||
Args:
|
||||
private_key_path: Path to private key file
|
||||
public_key_path: Path to public key file (optional, derived from private key if not provided)
|
||||
password: Optional password for encrypted private key
|
||||
"""
|
||||
print(f"🔑 Loading keys from files...")
|
||||
|
||||
# Load private key
|
||||
with open(private_key_path, "rb") as f:
|
||||
private_pem = f.read()
|
||||
|
||||
# Try different password options
|
||||
password_options = []
|
||||
if password:
|
||||
password_options.append(password.encode('utf-8'))
|
||||
password_options.append(None) # Try without password
|
||||
|
||||
last_error = None
|
||||
for pwd in password_options:
|
||||
try:
|
||||
self.private_key = serialization.load_pem_private_key(
|
||||
private_pem,
|
||||
password=pwd,
|
||||
backend=default_backend()
|
||||
)
|
||||
print(f" ✓ Private key loaded {'with password' if pwd else 'without password'}")
|
||||
break
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
continue
|
||||
else:
|
||||
raise ValueError(f"Failed to load private key. Tried all password options. Error: {last_error}")
|
||||
|
||||
# Get public key
|
||||
public_key = self.private_key.public_key()
|
||||
|
||||
# Load public key from file if provided, otherwise derive from private key
|
||||
if public_key_path:
|
||||
with open(public_key_path, "r") as f:
|
||||
self.public_key_pem = f.read().strip()
|
||||
else:
|
||||
self.public_key_pem = public_key.public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
||||
).decode('utf-8')
|
||||
|
||||
print(" ✓ Keys loaded successfully")
|
||||
|
||||
async def fetch_server_public_key(self) -> str:
|
||||
"""
|
||||
Fetch the server's public key from the /pki/public_key endpoint.
|
||||
|
||||
Returns:
|
||||
Server's public key as PEM string
|
||||
"""
|
||||
print("🔑 Fetching server's public key...")
|
||||
|
||||
url = f"{self.router_url}/pki/public_key"
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.get(url)
|
||||
|
||||
if response.status_code == 200:
|
||||
server_public_key = response.text
|
||||
print(" ✓ Server's public key fetched successfully")
|
||||
return server_public_key
|
||||
else:
|
||||
raise ValueError(f"Failed to fetch server's public key: HTTP {response.status_code}")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to fetch server's public key: {e}")
|
||||
|
||||
async def encrypt_payload(self, payload: Dict[str, Any]) -> bytes:
|
||||
"""
|
||||
Encrypt a payload using hybrid encryption (AES-256-GCM + RSA-OAEP).
|
||||
|
||||
Args:
|
||||
payload: Dictionary containing the chat completion request
|
||||
|
||||
Returns:
|
||||
Encrypted payload as bytes
|
||||
|
||||
Raises:
|
||||
Exception: If encryption fails
|
||||
"""
|
||||
print("🔒 Encrypting payload...")
|
||||
|
||||
try:
|
||||
# Serialize payload to JSON
|
||||
payload_json = json.dumps(payload).encode('utf-8')
|
||||
print(f" Payload size: {len(payload_json)} bytes")
|
||||
|
||||
# Generate random AES key
|
||||
aes_key = os.urandom(32) # 256-bit key
|
||||
|
||||
# Encrypt payload with AES-GCM using Cipher API (matching server implementation)
|
||||
nonce = os.urandom(12) # 96-bit nonce for GCM
|
||||
cipher = Cipher(
|
||||
algorithms.AES(aes_key),
|
||||
modes.GCM(nonce),
|
||||
backend=default_backend()
|
||||
)
|
||||
encryptor = cipher.encryptor()
|
||||
ciphertext = encryptor.update(payload_json) + encryptor.finalize()
|
||||
tag = encryptor.tag
|
||||
|
||||
# Fetch server's public key for encrypting the AES key
|
||||
server_public_key_pem = await self.fetch_server_public_key()
|
||||
|
||||
# Encrypt AES key with server's RSA-OAEP
|
||||
server_public_key = serialization.load_pem_public_key(
|
||||
server_public_key_pem.encode('utf-8'),
|
||||
backend=default_backend()
|
||||
)
|
||||
encrypted_aes_key = server_public_key.encrypt(
|
||||
aes_key,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None
|
||||
)
|
||||
)
|
||||
|
||||
# Create encrypted package
|
||||
encrypted_package = {
|
||||
"version": "1.0",
|
||||
"algorithm": "hybrid-aes256-rsa4096",
|
||||
"encrypted_payload": {
|
||||
"ciphertext": base64.b64encode(ciphertext).decode('utf-8'),
|
||||
"nonce": base64.b64encode(nonce).decode('utf-8'),
|
||||
"tag": base64.b64encode(tag).decode('utf-8')
|
||||
},
|
||||
"encrypted_aes_key": base64.b64encode(encrypted_aes_key).decode('utf-8'),
|
||||
"key_algorithm": "RSA-OAEP-SHA256",
|
||||
"payload_algorithm": "AES-256-GCM"
|
||||
}
|
||||
|
||||
# Serialize package to JSON and return as bytes
|
||||
package_json = json.dumps(encrypted_package).encode('utf-8')
|
||||
print(f" ✓ Encrypted package size: {len(package_json)} bytes")
|
||||
|
||||
return package_json
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"Encryption failed: {str(e)}")
|
||||
|
||||
async def decrypt_response(self, encrypted_response: bytes, payload_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Decrypt a response from the secure endpoint.
|
||||
|
||||
Args:
|
||||
encrypted_response: Encrypted response bytes
|
||||
payload_id: Payload ID for metadata verification
|
||||
|
||||
Returns:
|
||||
Decrypted response dictionary
|
||||
"""
|
||||
print("🔓 Decrypting response...")
|
||||
|
||||
# Parse encrypted package
|
||||
try:
|
||||
package = json.loads(encrypted_response.decode('utf-8'))
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid encrypted package format: {e}")
|
||||
|
||||
# Validate package structure
|
||||
required_fields = ["version", "algorithm", "encrypted_payload", "encrypted_aes_key"]
|
||||
for field in required_fields:
|
||||
if field not in package:
|
||||
raise ValueError(f"Missing required field in encrypted package: {field}")
|
||||
|
||||
# Decrypt AES key with private key
|
||||
encrypted_aes_key = base64.b64decode(package["encrypted_aes_key"])
|
||||
aes_key = self.private_key.decrypt(
|
||||
encrypted_aes_key,
|
||||
padding.OAEP(
|
||||
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
||||
algorithm=hashes.SHA256(),
|
||||
label=None
|
||||
)
|
||||
)
|
||||
|
||||
# Decrypt payload with AES-GCM using Cipher API (matching server implementation)
|
||||
ciphertext = base64.b64decode(package["encrypted_payload"]["ciphertext"])
|
||||
nonce = base64.b64decode(package["encrypted_payload"]["nonce"])
|
||||
tag = base64.b64decode(package["encrypted_payload"]["tag"])
|
||||
|
||||
cipher = Cipher(
|
||||
algorithms.AES(aes_key),
|
||||
modes.GCM(nonce, tag),
|
||||
backend=default_backend()
|
||||
)
|
||||
decryptor = cipher.decryptor()
|
||||
plaintext = decryptor.update(ciphertext) + decryptor.finalize()
|
||||
|
||||
# Parse decrypted response
|
||||
response = json.loads(plaintext.decode('utf-8'))
|
||||
|
||||
# Add metadata for debugging
|
||||
if "_metadata" not in response:
|
||||
response["_metadata"] = {}
|
||||
response["_metadata"].update({
|
||||
"payload_id": payload_id,
|
||||
"processed_at": package.get("processed_at"),
|
||||
"is_encrypted": True,
|
||||
"encryption_algorithm": package["algorithm"]
|
||||
})
|
||||
|
||||
print(f" ✓ Response decrypted successfully")
|
||||
print(f" Response size: {len(plaintext)} bytes")
|
||||
|
||||
return response
|
||||
|
||||
async def send_secure_request(self, payload: Dict[str, Any], payload_id: str) -> 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
|
||||
|
||||
Returns:
|
||||
Decrypted response from the LLM
|
||||
"""
|
||||
print("\n📤 Sending secure chat completion request...")
|
||||
|
||||
# Step 1: Encrypt the payload
|
||||
encrypted_payload = await self.encrypt_payload(payload)
|
||||
|
||||
# Step 2: Prepare headers
|
||||
headers = {
|
||||
"X-Payload-ID": payload_id,
|
||||
"X-Public-Key": urllib.parse.quote(self.public_key_pem),
|
||||
"Content-Type": "application/octet-stream"
|
||||
}
|
||||
|
||||
# Step 3: Send request to router
|
||||
url = f"{self.router_url}/v1/chat/secure_completion"
|
||||
print(f" Target URL: {url}")
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
content=encrypted_payload
|
||||
)
|
||||
|
||||
print(f" HTTP Status: {response.status_code}")
|
||||
|
||||
if response.status_code == 200:
|
||||
# Step 4: Decrypt the response
|
||||
encrypted_response = response.content
|
||||
decrypted_response = await self.decrypt_response(encrypted_response, payload_id)
|
||||
return decrypted_response
|
||||
|
||||
elif response.status_code == 400:
|
||||
error = response.json()
|
||||
raise ValueError(f"Bad request: {error.get('detail', 'Unknown error')}")
|
||||
|
||||
elif response.status_code == 404:
|
||||
error = response.json()
|
||||
raise ValueError(f"Endpoint not found: {error.get('detail', 'Secure inference not enabled')}")
|
||||
|
||||
elif response.status_code == 500:
|
||||
error = response.json()
|
||||
raise ValueError(f"Server error: {error.get('detail', 'Internal server error')}")
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unexpected status code: {response.status_code}")
|
||||
|
||||
except httpx.NetworkError as e:
|
||||
raise ConnectionError(f"Failed to connect to router: {e}")
|
||||
except Exception as e:
|
||||
raise Exception(f"Request failed: {e}")
|
||||
12
nomyo/__init__.py
Normal file
12
nomyo/__init__.py
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
"""
|
||||
NOMYO Secure Python Chat Client
|
||||
|
||||
OpenAI-compatible secure chat client with end-to-end encryption.
|
||||
"""
|
||||
|
||||
from .nomyo import SecureChatCompletion
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__author__ = "NOMYO AI"
|
||||
__license__ = "Apache-2.0"
|
||||
__all__ = ["SecureChatCompletion"]
|
||||
161
nomyo/nomyo.py
Normal file
161
nomyo/nomyo.py
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
import uuid
|
||||
from typing import Dict, Any, List
|
||||
from .SecureCompletionClient import SecureCompletionClient
|
||||
|
||||
class SecureChatCompletion:
|
||||
"""
|
||||
OpenAI-compatible secure chat completion client.
|
||||
|
||||
This class provides the same interface as OpenAI's ChatCompletion.create()
|
||||
method, but automatically encrypts all requests and decrypts all responses
|
||||
for secure communication with the NOMYO Router's /v1/chat/secure_completion
|
||||
endpoint.
|
||||
|
||||
Usage:
|
||||
```python
|
||||
# Create a client instance
|
||||
client = SecureChatCompletion(base_url="http://api.nomyo.ai:12434")
|
||||
|
||||
# Simple chat completion
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "What is the capital of France?"}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
# With tools
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "What's the weather in Paris?"}
|
||||
],
|
||||
tools=[...],
|
||||
temperature=0.7
|
||||
)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str = "http://api.nomyo.ai:12434"):
|
||||
"""
|
||||
Initialize the secure chat completion client.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the NOMYO Router (e.g., "http://api.nomyo.ai:12434")
|
||||
This parameter is named 'base_url' for OpenAI compatibility.
|
||||
"""
|
||||
|
||||
self.client = SecureCompletionClient(router_url=base_url)
|
||||
self._keys_initialized = False
|
||||
|
||||
async def _ensure_keys(self):
|
||||
"""Ensure keys are loaded or generated."""
|
||||
if not self._keys_initialized:
|
||||
# Try to load existing keys
|
||||
try:
|
||||
await self.client.load_keys("client_keys/private_key.pem", "client_keys/public_key.pem")
|
||||
self._keys_initialized = True
|
||||
except Exception:
|
||||
# Generate new keys if loading fails
|
||||
await self.client.generate_keys()
|
||||
self._keys_initialized = True
|
||||
|
||||
async def create(self, model: str, messages: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
Creates a new chat completion for the provided messages and parameters.
|
||||
|
||||
This method provides the same interface as OpenAI's ChatCompletion.create()
|
||||
but automatically handles encryption and decryption for secure communication.
|
||||
|
||||
Args:
|
||||
model: The model to use for the chat completion.
|
||||
messages: A list of message objects. Each message has a role ("system",
|
||||
"user", or "assistant") and content.
|
||||
**kwargs: Additional parameters that can be passed to the API.
|
||||
Supported parameters include:
|
||||
- temperature: float (0-2)
|
||||
- max_tokens: int
|
||||
- tools: List of tool definitions
|
||||
- tool_choice: str ("auto", "none", or specific tool name)
|
||||
- stop: Union[str, List[str]]
|
||||
- presence_penalty: float
|
||||
- frequency_penalty: float
|
||||
- logit_bias: Dict[str, float]
|
||||
- user: str
|
||||
- base_url: str (alternative to initializing with router_url)
|
||||
|
||||
Returns:
|
||||
A dictionary containing the chat completion response with the following structure:
|
||||
{
|
||||
"id": str,
|
||||
"object": "chat.completion",
|
||||
"created": int,
|
||||
"model": str,
|
||||
"choices": [
|
||||
{
|
||||
"index": int,
|
||||
"message": {
|
||||
"role": str,
|
||||
"content": str,
|
||||
"tool_calls": List[Dict] # if tools were used
|
||||
},
|
||||
"finish_reason": str
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": int,
|
||||
"completion_tokens": int,
|
||||
"total_tokens": int
|
||||
}
|
||||
}
|
||||
|
||||
Raises:
|
||||
ValueError: If required parameters are missing or invalid.
|
||||
ConnectionError: If the connection to the router fails.
|
||||
Exception: For other errors during the request.
|
||||
"""
|
||||
# Extract base_url if provided (OpenAI compatibility)
|
||||
base_url = kwargs.pop("base_url", None)
|
||||
|
||||
# Use the instance's client unless base_url is explicitly overridden
|
||||
if base_url is not None:
|
||||
# Create a temporary client with overridden base_url
|
||||
temp_client = type(self)(base_url=base_url)
|
||||
instance = temp_client
|
||||
else:
|
||||
# Use the instance's existing client
|
||||
instance = self
|
||||
|
||||
# Ensure keys are available
|
||||
await instance._ensure_keys()
|
||||
|
||||
# Prepare payload in OpenAI format
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
**kwargs
|
||||
}
|
||||
|
||||
# Generate a unique payload ID
|
||||
payload_id = f"openai-compat-{uuid.uuid4()}"
|
||||
|
||||
# Send secure request
|
||||
response = await instance.client.send_secure_request(payload, payload_id)
|
||||
|
||||
return response
|
||||
|
||||
async def acreate(self, model: str, messages: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
Async alias for create() method.
|
||||
|
||||
This provides the same functionality as create() but with an explicit
|
||||
async name, following OpenAI's naming conventions.
|
||||
|
||||
Args:
|
||||
Same as create() method.
|
||||
|
||||
Returns:
|
||||
Same as create() method.
|
||||
"""
|
||||
return await self.create(model, messages, **kwargs)
|
||||
52
pyproject.toml
Normal file
52
pyproject.toml
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
[build-system]
|
||||
requires = ["hatchling>=1.0.0", "wheel"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "nomyo"
|
||||
version = "0.1.0"
|
||||
description = "OpenAI-compatible secure chat client with end-to-end encryption for NOMYO Inference Endpoints"
|
||||
authors = [
|
||||
{name = "NOMYO.AI", email = "ichi@nomyo.ai"},
|
||||
]
|
||||
readme = "README.md"
|
||||
license = {text = "Apache-2.0"}
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Information Technology",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.8",
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Topic :: Security :: Cryptography",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
"Topic :: Communications :: Chat",
|
||||
"Operating System :: OS Independent",
|
||||
]
|
||||
requires-python = ">=3.8"
|
||||
dependencies = [
|
||||
"anyio==4.12.0",
|
||||
"certifi==2025.11.12",
|
||||
"cffi==2.0.0",
|
||||
"cryptography==46.0.3",
|
||||
"exceptiongroup==1.3.1",
|
||||
"h11==0.16.0",
|
||||
"httpcore==1.0.9",
|
||||
"httpx==0.28.1",
|
||||
"idna==3.11",
|
||||
"pycparser==2.23",
|
||||
"typing_extensions==4.15.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://nomyo.ai"
|
||||
Documentation = "https://nomyo.ai/nomyo-docs"
|
||||
Repository = "https://github.com/nomyo-ai/nomyo"
|
||||
Issues = "https://github.com/nomyo-ai/nomyo/issues"
|
||||
|
||||
[tool.setuptools]
|
||||
packages = ["nomyo"]
|
||||
11
requirements.txt
Normal file
11
requirements.txt
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
anyio==4.12.0
|
||||
certifi==2025.11.12
|
||||
cffi==2.0.0
|
||||
cryptography==46.0.3
|
||||
exceptiongroup==1.3.1
|
||||
h11==0.16.0
|
||||
httpcore==1.0.9
|
||||
httpx==0.28.1
|
||||
idna==3.11
|
||||
pycparser==2.23
|
||||
typing_extensions==4.15.0
|
||||
221
test.py
Normal file
221
test.py
Normal file
|
|
@ -0,0 +1,221 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Test script to verify OpenAI compatibility of SecureChatCompletion.
|
||||
|
||||
This script demonstrates that the SecureChatCompletion class provides
|
||||
the same interface as OpenAI's ChatCompletion.create() method.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from nomyo import SecureChatCompletion
|
||||
|
||||
client = SecureChatCompletion(base_url="http://localhost:12434")
|
||||
|
||||
async def test_basic_chat():
|
||||
"""Test basic chat completion with OpenAI-style API."""
|
||||
print("=" * 70)
|
||||
print("TEST 1: Basic Chat Completion (OpenAI-style API)")
|
||||
print("=" * 70)
|
||||
|
||||
# This is how you would use OpenAI's client:
|
||||
# response = await openai.ChatCompletion.create(
|
||||
# model="gpt-3.5-turbo",
|
||||
# messages=[...],
|
||||
# temperature=0.7
|
||||
# )
|
||||
|
||||
# Now with SecureChatCompletion (same API!):
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"}
|
||||
],
|
||||
temperature=0.7,
|
||||
)
|
||||
|
||||
# Verify response structure matches OpenAI format
|
||||
assert "choices" in response, "Response missing 'choices' field"
|
||||
assert len(response["choices"]) > 0, "No choices in response"
|
||||
assert "message" in response["choices"][0], "Choice missing 'message' field"
|
||||
assert "content" in response["choices"][0]["message"], "Message missing 'content' field"
|
||||
assert "finish_reason" in response["choices"][0], "Choice missing 'finish_reason' field"
|
||||
|
||||
print("✅ Response structure matches OpenAI format")
|
||||
print(f"✅ Model: {response.get('model')}")
|
||||
print(f"✅ Finish Reason: {response['choices'][0].get('finish_reason')}")
|
||||
print(f"✅ Content: {response['choices'][0]['message']['content']}...")
|
||||
return True
|
||||
|
||||
async def test_chat_with_tools():
|
||||
"""Test chat completion with tools (OpenAI-style API)."""
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 2: Chat with Tools (OpenAI-style API)")
|
||||
print("=" * 70)
|
||||
|
||||
# This is how you would use OpenAI's client with tools:
|
||||
# response = await openai.ChatCompletion.create(
|
||||
# model="gpt-3.5-turbo",
|
||||
# messages=[...],
|
||||
# tools=[...],
|
||||
# temperature=0.7
|
||||
# )
|
||||
|
||||
# Now with SecureChatCompletion (same API!):
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant with tools."},
|
||||
{"role": "user", "content": "What's the weather in Paris?"}
|
||||
],
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather information for a location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "City name"
|
||||
}
|
||||
},
|
||||
"required": ["location"]
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
temperature=0.7,
|
||||
max_tokens=2000
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
assert "choices" in response, "Response missing 'choices' field"
|
||||
assert "message" in response["choices"][0], "Choice missing 'message' field"
|
||||
|
||||
print("✅ Response structure matches OpenAI format")
|
||||
print(f"✅ Model: {response.get('model')}")
|
||||
print(f"✅ Content: {response['choices'][0]['message']['content']}...")
|
||||
|
||||
# Check for tool calls if present
|
||||
if 'tool_calls' in response['choices'][0]['message']:
|
||||
print("✅ Tool calls detected in response")
|
||||
for tool_call in response['choices'][0]['message']['tool_calls']:
|
||||
print(f" - Function: {tool_call['function']['name']}")
|
||||
return True
|
||||
|
||||
async def test_all_openai_parameters():
|
||||
"""Test that all common OpenAI parameters are supported."""
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 3: All OpenAI Parameters Support")
|
||||
print("=" * 70)
|
||||
|
||||
# Test with various OpenAI parameters
|
||||
response = await client.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello!"}
|
||||
],
|
||||
temperature=0.7,
|
||||
max_tokens=100,
|
||||
top_p=0.9,
|
||||
frequency_penalty=0.0,
|
||||
presence_penalty=0.0,
|
||||
stop=None,
|
||||
n=1,
|
||||
stream=False,
|
||||
user="test_user"
|
||||
)
|
||||
|
||||
print("✅ All OpenAI parameters accepted")
|
||||
print(f"✅ Response received successfully")
|
||||
return True
|
||||
|
||||
async def test_async_alias():
|
||||
"""Test the acreate async alias method."""
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 4: Async Alias (acreate)")
|
||||
print("=" * 70)
|
||||
|
||||
# Test using the acreate alias on the client instance
|
||||
response = await client.acreate(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "user", "content": "Test message"}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print("✅ acreate() method works correctly")
|
||||
print(f"✅ Response received: {response['choices'][0]['message']['content']}...")
|
||||
return True
|
||||
|
||||
async def test_error_handling():
|
||||
"""Test error handling."""
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 5: Error Handling")
|
||||
print("=" * 70)
|
||||
|
||||
try:
|
||||
# This should fail gracefully
|
||||
response = await client.create(
|
||||
model="nonexistent-model",
|
||||
messages=[
|
||||
{"role": "user", "content": "Test"}
|
||||
]
|
||||
)
|
||||
print("⚠️ Expected error did not occur")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"✅ Error handled correctly: {type(e).__name__}")
|
||||
return True
|
||||
|
||||
async def main():
|
||||
"""Run all compatibility tests."""
|
||||
print("=" * 70)
|
||||
print("SECURE CHAT CLIENT - OpenAI Compatibility Tests")
|
||||
print("=" * 70)
|
||||
print("\nTesting that SecureChatCompletion provides the same API as")
|
||||
print("openai.ChatCompletion.create() with end-to-end encryption...\n")
|
||||
|
||||
tests = [
|
||||
test_basic_chat,
|
||||
test_chat_with_tools,
|
||||
test_all_openai_parameters,
|
||||
test_async_alias,
|
||||
test_error_handling,
|
||||
]
|
||||
|
||||
results = []
|
||||
for test in tests:
|
||||
try:
|
||||
result = await test()
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
print(f"\n❌ Test failed with exception: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
results.append(False)
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST SUMMARY")
|
||||
print("=" * 70)
|
||||
passed = sum(results)
|
||||
total = len(results)
|
||||
print(f"Passed: {passed}/{total}")
|
||||
|
||||
if passed == total:
|
||||
print("\n🎉 ALL TESTS PASSED!")
|
||||
print("\nThe SecureChatCompletion class is fully compatible with")
|
||||
print("OpenAI's ChatCompletion.create() API while providing")
|
||||
print("end-to-end encryption for secure communication.")
|
||||
else:
|
||||
print(f"\n⚠️ {total - passed} test(s) failed")
|
||||
|
||||
return passed == total
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = asyncio.run(main())
|
||||
exit(0 if success else 1)
|
||||
Loading…
Add table
Add a link
Reference in a new issue