- Reset master to upstream/main (16,697 commits) - Overlay 2,271 local-only files (skills, tools, workspace, configs, apps) - Restore IDENTITY.md and USER.md templates - Build verified, gateway running, Discord working Co-Authored-By: Claude Opus 4.6 <[email protected]>
122 lines
3.6 KiB
Python
122 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""nvidia_api.py
|
|
Simple CLI wrapper for NVIDIA Inference API (OpenAI-compatible).
|
|
|
|
Usage:
|
|
python3 nvidia_api.py health
|
|
python3 nvidia_api.py chat --model <model> --prompt "..."
|
|
python3 nvidia_api.py bench --model <model> --requests 3
|
|
|
|
The script looks for credentials in ~/.openclaw/credentials/nvidia.key or env NVIDIA_API_KEY.
|
|
"""
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import argparse
|
|
from typing import Optional
|
|
|
|
import requests
|
|
|
|
BASE = os.environ.get('NVIDIA_BASE_URL','https://integrate.api.nvidia.com/v1')
|
|
KEY = os.environ.get('NVIDIA_API_KEY')
|
|
if not KEY:
|
|
key_path = os.path.expanduser('~/.openclaw/credentials/nvidia.key')
|
|
if os.path.exists(key_path):
|
|
with open(key_path,'r') as f:
|
|
KEY = f.read().strip()
|
|
|
|
HEADERS = {
|
|
'Authorization': f'Bearer {KEY}' if KEY else '',
|
|
'Content-Type': 'application/json'
|
|
}
|
|
|
|
|
|
def health():
|
|
url = f"{BASE}/models"
|
|
r = requests.get(url, headers=HEADERS, timeout=20)
|
|
try:
|
|
r.raise_for_status()
|
|
except Exception as e:
|
|
print(json.dumps({'ok': False, 'error': str(e), 'status_code': r.status_code, 'body': r.text}))
|
|
sys.exit(2)
|
|
data = r.json()
|
|
print(json.dumps({'ok': True, 'models_count': len(data.get('data', data) if isinstance(data, dict) else data)}))
|
|
return data
|
|
|
|
|
|
def chat(model: str, prompt: str, temperature: float=0.2, max_tokens: int=512):
|
|
url = f"{BASE}/chat/completions"
|
|
payload = {
|
|
'model': model,
|
|
'messages': [{'role':'user','content': prompt}],
|
|
'temperature': temperature,
|
|
'max_tokens': max_tokens,
|
|
}
|
|
r = requests.post(url, headers=HEADERS, json=payload, timeout=60)
|
|
try:
|
|
r.raise_for_status()
|
|
except Exception as e:
|
|
print(json.dumps({'ok': False, 'error': str(e), 'status_code': r.status_code, 'body': r.text}))
|
|
sys.exit(2)
|
|
data = r.json()
|
|
# Print concise output
|
|
choice = None
|
|
try:
|
|
choice = data['choices'][0]['message']
|
|
except Exception:
|
|
choice = data
|
|
print(json.dumps({'ok': True, 'response': choice}))
|
|
return data
|
|
|
|
|
|
def bench(model: str, prompt: str = 'Hello', requests_n: int =3):
|
|
results = []
|
|
for i in range(requests_n):
|
|
t0 = time.time()
|
|
try:
|
|
data = chat(model, prompt)
|
|
ok = True
|
|
except SystemExit:
|
|
ok = False
|
|
data = None
|
|
dt = time.time() - t0
|
|
results.append({'index': i, 'ok': ok, 'latency_s': dt})
|
|
time.sleep(0.5)
|
|
summary = {
|
|
'model': model,
|
|
'requests': requests_n,
|
|
'avg_latency_s': sum(r['latency_s'] for r in results)/len(results) if results else None,
|
|
'results': results
|
|
}
|
|
print(json.dumps({'ok': True, 'summary': summary}))
|
|
return summary
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
sub = p.add_subparsers(dest='cmd')
|
|
sub.add_parser('health')
|
|
pc = sub.add_parser('chat')
|
|
pc.add_argument('--model', required=True)
|
|
pc.add_argument('--prompt', required=True)
|
|
pc.add_argument('--temperature', type=float, default=0.2)
|
|
pc.add_argument('--max_tokens', type=int, default=512)
|
|
pb = sub.add_parser('bench')
|
|
pb.add_argument('--model', required=True)
|
|
pb.add_argument('--prompt', default='Hello')
|
|
pb.add_argument('--requests', type=int, default=3)
|
|
|
|
args = p.parse_args()
|
|
if args.cmd == 'health':
|
|
health()
|
|
elif args.cmd == 'chat':
|
|
chat(args.model, args.prompt, temperature=args.temperature, max_tokens=args.max_tokens)
|
|
elif args.cmd == 'bench':
|
|
bench(args.model, args.prompt, requests_n=args.requests)
|
|
else:
|
|
p.print_help()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|