|
- import os
-
- import openai
-
- OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
- OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://lancerouter.ai/v1")
- MODEL_ID = os.environ.get("OPENAI_MODEL", "google/gemini-3.1-pro-preview")
-
- def list_models():
- client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
- models = client.models.list()
- model_ids = sorted([m.id for m in models.data])
- print(f"共找到 {len(model_ids)} 个模型:")
- for model_id in model_ids:
- print(f" - {model_id}")
-
-
- def test_model(stream: bool = False, thinking: bool = True, prompt: str = "Hello, how are you?"):
- client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
-
- extra_body = {"thinking": {"type": "enabled", "budget_tokens": 5000}} if thinking else {}
-
- if stream:
- response = client.chat.completions.create(
- model=MODEL_ID,
- messages=[{"role": "user", "content": prompt}],
- stream=True,
- extra_body=extra_body,
- )
- for chunk in response:
- if not chunk.choices:
- continue
- delta = chunk.choices[0].delta
- # 输出思考内容(thinking block)
- if hasattr(delta, "thinking") and delta.thinking:
- print(delta.thinking, end="", flush=True)
- elif delta.content:
- print(delta.content, end="", flush=True)
- else:
- response = client.chat.completions.create(
- model=MODEL_ID,
- messages=[{"role": "user", "content": prompt}],
- extra_body=extra_body,
- )
- message = response.choices[0].message
- # 输出思考内容(thinking block)
- if hasattr(message, "thinking") and message.thinking:
- print(f"<thinking>\n{message.thinking}\n</thinking>\n")
- print(message.content)
-
-
- if __name__ == "__main__":
- # list_models()
- test_model(stream=True, thinking=False, prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。")
|