sagecode/api_client.py at master · Hugoean/sagecode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
"""
api_client.py
统一API客户端工厂 — SAGE-v2

支持两种后端（通过环境变量切换）：
  DeepSeek: DEEPSEEK_API_KEY + DEEPSEEK_BASE_URL
  OpenAI:   OPENAI_API_KEY  + OPENAI_PROXY（可选）

优先使用DeepSeek，若无DeepSeek Key则回退到OpenAI。

环境变量示例（写入.env文件）：
  # DeepSeek模式
  DEEPSEEK_API_KEY=sk-xxx
  DEEPSEEK_BASE_URL=https://api.deepseek.com
  GENERATOR_MODEL=deepseek-chat
  REVIEWER_MODEL=deepseek-reasoner

  # OpenAI模式（带代理）
  OPENAI_API_KEY=sk-proj-xxx
  OPENAI_PROXY=http://172.26.224.1:7897
  GENERATOR_MODEL=gpt-4o-mini
  REVIEWER_MODEL=o1-mini
"""

from __future__ import annotations
import os
import httpx
from openai import AsyncOpenAI
from dotenv import load_dotenv

load_dotenv()


def make_client() -> AsyncOpenAI:
    """
    创建AsyncOpenAI客户端。
    优先DeepSeek，回退OpenAI，支持代理。
    """
    deepseek_key = os.getenv("DEEPSEEK_API_KEY", "")
    openai_key   = os.getenv("OPENAI_API_KEY", "")
    proxy        = os.getenv("OPENAI_PROXY", "")

    if deepseek_key:
        return AsyncOpenAI(
            api_key=deepseek_key,
            base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
        )

    # OpenAI fallback（带代理）
    http_client = (
        httpx.AsyncClient(proxy=proxy, timeout=60)
        if proxy
        else httpx.AsyncClient(timeout=60)
    )
    return AsyncOpenAI(api_key=openai_key, http_client=http_client)


def get_generator_model() -> str:
    return os.getenv("GENERATOR_MODEL", "deepseek-chat")


def get_reviewer_model() -> str:
    return os.getenv("REVIEWER_MODEL", "deepseek-reasoner")