api模型检测
Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

78 строки
2.5 KiB

  1. from google import genai
  2. from google.genai import types
  3. import os
  4. GEMINI_API_KEY = os.environ["GEMINI_API_KEY"]
  5. GEMINI_BASE_URL = os.environ.get("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta")
  6. MODEL_ID = os.environ.get("GEMINI_MODEL", "gemini-3.1-pro-preview")
  7. def _make_client() -> genai.Client:
  8. return genai.Client(
  9. api_key=GEMINI_API_KEY,
  10. http_options=types.HttpOptions(base_url=GEMINI_BASE_URL),
  11. )
  12. def list_models():
  13. client = _make_client()
  14. models = client.models.list()
  15. model_ids = sorted([m.name for m in models])
  16. print(f"共找到 {len(model_ids)} 个模型:")
  17. for model_id in model_ids:
  18. print(f" - {model_id}")
  19. def test_model(
  20. stream: bool = False,
  21. thinking: bool = True,
  22. prompt: str = "Hello, how are you?",
  23. ):
  24. client = _make_client()
  25. thinking_config = (
  26. types.ThinkingConfig(include_thoughts=True, thinking_budget=5000)
  27. if thinking
  28. else types.ThinkingConfig(include_thoughts=False)
  29. )
  30. config = types.GenerateContentConfig(thinking_config=thinking_config)
  31. if stream:
  32. in_thinking = False
  33. for chunk in client.models.generate_content_stream(
  34. model=MODEL_ID, contents=prompt, config=config
  35. ):
  36. if not chunk.candidates:
  37. continue
  38. for part in chunk.candidates[0].content.parts or []:
  39. if part.thought:
  40. if not in_thinking:
  41. print("<thinking>", flush=True)
  42. in_thinking = True
  43. print(part.text, end="", flush=True)
  44. else:
  45. if in_thinking:
  46. print("\n</thinking>\n", flush=True)
  47. in_thinking = False
  48. print(part.text or "", end="", flush=True)
  49. if in_thinking:
  50. print("\n</thinking>", flush=True)
  51. print()
  52. else:
  53. response = client.models.generate_content(
  54. model=MODEL_ID, contents=prompt, config=config
  55. )
  56. for part in response.candidates[0].content.parts:
  57. if part.thought:
  58. print(f"<thinking>\n{part.text}\n</thinking>\n")
  59. else:
  60. print(part.text or "", end="")
  61. print()
  62. if __name__ == "__main__":
  63. test_model(
  64. stream=True,
  65. thinking=False,
  66. prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。",
  67. )