AIRouter/examples/demo_usage.py

#!/usr/bin/env python3
"""
Edge AI Gateway 使用示范

本文件演示其他应用如何通过 OpenAI 兼容 API 调用 Edge AI Gateway。
支持两种方式：
1. 使用 openai SDK（推荐）
2. 使用 requests 库直接调用

运行前请确保：
- Edge AI Gateway 已启动 (默认 http://localhost:8080)
- 已配置有效的 API Key
"""

# ============================================================
# 方式一：使用 openai SDK（推荐）
# ============================================================
# pip install openai

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8080/v1",
    api_key="test-key",  # 替换为你的 API Key
)

# --- 非流式对话 ---
print("=== 非流式对话 ===")
response = client.chat.completions.create(
    model="general-chat",  # 逻辑模型名，网关自动路由到实际模型
    messages=[
        {"role": "system", "content": "你是一个简洁的助手"},
        {"role": "user", "content": "你好，介绍一下你自己"},
    ],
    stream=False,
)
print(f"回复: {response.choices[0].message.content}")
print(f"Token: 输入={response.usage.prompt_tokens}, 输出={response.usage.completion_tokens}")
print()

# --- 流式对话 (SSE) ---
print("=== 流式对话 ===")
stream = client.chat.completions.create(
    model="fast-chat",
    messages=[
        {"role": "user", "content": "写一个 Python 的 hello world"},
    ],
    stream=True,
)
for chunk in stream:
    delta = chunk.choices[0].delta.content
    if delta:
        print(delta, end="", flush=True)
print("\n")

# --- 带优先级的请求 (P0 最高) ---
print("=== 高优先级请求 ===")
response = client.chat.completions.create(
    model="general-chat",
    messages=[{"role": "user", "content": "1+1=?"}],
    extra_headers={"X-Priority": "P0"},  # 通过 header 传递优先级
)
print(f"回复: {response.choices[0].message.content}")
print()

# --- 查看可用模型 ---
print("=== 可用模型 ===")
models = client.models.list()
for m in models.data:
    print(f"  - {m.id}")
print()


# ============================================================
# 方式二：使用 requests 库直接调用
# ============================================================
# pip install requests

import requests

BASE_URL = "http://localhost:8080"
API_KEY = "test-key"
HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json",
}

# --- 非流式 ---
print("=== requests: 非流式 ===")
resp = requests.post(
    f"{BASE_URL}/v1/chat/completions",
    headers=HEADERS,
    json={
        "model": "general-chat",
        "messages": [{"role": "user", "content": "你好"}],
        "stream": False,
    },
)
data = resp.json()
print(f"状态: {data.get('status')}")
print(f"回复: {data['choices'][0]['message']['content']}")
print(f"模型: {data.get('actual_model')}")
print()

# --- 流式 (SSE) ---
print("=== requests: 流式 ===")
resp = requests.post(
    f"{BASE_URL}/v1/chat/completions",
    headers=HEADERS,
    json={
        "model": "fast-chat",
        "messages": [{"role": "user", "content": "讲个笑话"}],
        "stream": True,
    },
    stream=True,
)
for line in resp.iter_lines():
    if line:
        line = line.decode("utf-8")
        if line.startswith("data: ") and line != "data: [DONE]":
            import json
            chunk = json.loads(line[6:])
            delta = chunk["choices"][0]["delta"].get("content", "")
            if delta:
                print(delta, end="", flush=True)
print("\n")

# --- 会话管理 ---
print("=== 会话管理 ===")
# 创建会话
resp = requests.post(
    f"{BASE_URL}/v1/sessions",
    headers=HEADERS,
    json={"application_id": "my-app", "user_id": "user123"},
)
session = resp.json()
session_id = session["session_id"]
print(f"创建会话: {session_id}")

# 查询会话
resp = requests.get(
    f"{BASE_URL}/v1/sessions/{session_id}",
    headers=HEADERS,
)
print(f"会话信息: {resp.json()}")

# 删除会话
resp = requests.delete(
    f"{BASE_URL}/v1/sessions/{session_id}",
    headers=HEADERS,
)
print(f"删除会话: HTTP {resp.status_code}")
print()

# --- 错误处理 ---
print("=== 错误处理 ===")
# 无认证
resp = requests.get(f"{BASE_URL}/v1/models")
print(f"无认证: {resp.status_code} - {resp.json()['error']['code']}")

# 无效 key
resp = requests.get(
    f"{BASE_URL}/v1/models",
    headers={"Authorization": "Bearer invalid-key"},
)
print(f"无效key: {resp.status_code} - {resp.json()['error']['code']}")

# 缺少 model
resp = requests.post(
    f"{BASE_URL}/v1/chat/completions",
    headers=HEADERS,
    json={"messages": [{"role": "user", "content": "hi"}]},
)
print(f"缺少model: {resp.status_code} - {resp.json()['error']['code']}")

# 未知 model
resp = requests.post(
    f"{BASE_URL}/v1/chat/completions",
    headers=HEADERS,
    json={"model": "nonexistent", "messages": [{"role": "user", "content": "hi"}]},
)
print(f"未知model: {resp.status_code} - {resp.json()['error']['code']}")

print("\n✅ 示范完成！")