api模型检测
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

55 lines
2.0 KiB

  1. import os
  2. import openai
  3. OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
  4. OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://lancerouter.ai/v1")
  5. MODEL_ID = os.environ.get("OPENAI_MODEL", "google/gemini-3.1-pro-preview")
  6. def list_models():
  7. client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
  8. models = client.models.list()
  9. model_ids = sorted([m.id for m in models.data])
  10. print(f"共找到 {len(model_ids)} 个模型:")
  11. for model_id in model_ids:
  12. print(f" - {model_id}")
  13. def test_model(stream: bool = False, thinking: bool = True, prompt: str = "Hello, how are you?"):
  14. client = openai.OpenAI(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL)
  15. extra_body = {"thinking": {"type": "enabled", "budget_tokens": 5000}} if thinking else {}
  16. if stream:
  17. response = client.chat.completions.create(
  18. model=MODEL_ID,
  19. messages=[{"role": "user", "content": prompt}],
  20. stream=True,
  21. extra_body=extra_body,
  22. )
  23. for chunk in response:
  24. if not chunk.choices:
  25. continue
  26. delta = chunk.choices[0].delta
  27. # 输出思考内容(thinking block)
  28. if hasattr(delta, "thinking") and delta.thinking:
  29. print(delta.thinking, end="", flush=True)
  30. elif delta.content:
  31. print(delta.content, end="", flush=True)
  32. else:
  33. response = client.chat.completions.create(
  34. model=MODEL_ID,
  35. messages=[{"role": "user", "content": prompt}],
  36. extra_body=extra_body,
  37. )
  38. message = response.choices[0].message
  39. # 输出思考内容(thinking block)
  40. if hasattr(message, "thinking") and message.thinking:
  41. print(f"<thinking>\n{message.thinking}\n</thinking>\n")
  42. print(message.content)
  43. if __name__ == "__main__":
  44. # list_models()
  45. test_model(stream=True, thinking=False, prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。")