llm_client.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. import os
  2. from openai import OpenAI
  3. from dotenv import load_dotenv
  4. from typing import List, Dict
  5. # 加载 .env 文件中的环境变量
  6. load_dotenv()
  7. class HelloAgentsLLM:
  8. """
  9. 为本书 "Hello Agents" 定制的LLM客户端。
  10. 它用于调用任何兼容OpenAI接口的服务,并默认使用流式响应。
  11. """
  12. def __init__(self, model: str = None, apiKey: str = None, baseUrl: str = None, timeout: int = None):
  13. """
  14. 初始化客户端。优先使用传入参数,如果未提供,则从环境变量加载。
  15. """
  16. self.model = model or os.getenv("LLM_MODEL_ID")
  17. apiKey = apiKey or os.getenv("LLM_API_KEY")
  18. baseUrl = baseUrl or os.getenv("LLM_BASE_URL")
  19. timeout = timeout or int(os.getenv("LLM_TIMEOUT", 60))
  20. if not all([self.model, apiKey, baseUrl]):
  21. raise ValueError("模型ID、API密钥和服务地址必须被提供或在.env文件中定义。")
  22. self.client = OpenAI(api_key=apiKey, base_url=baseUrl, timeout=timeout)
  23. def think(self, messages: List[Dict[str, str]], temperature: float = 0) -> str:
  24. """
  25. 调用大语言模型进行思考,并返回其响应。
  26. """
  27. print(f"🧠 正在调用 {self.model} 模型...")
  28. try:
  29. response = self.client.chat.completions.create(
  30. model=self.model,
  31. messages=messages,
  32. temperature=temperature,
  33. stream=True,
  34. )
  35. # 处理流式响应
  36. print("✅ 大语言模型响应成功:")
  37. collected_content = []
  38. for chunk in response:
  39. content = chunk.choices[0].delta.content or ""
  40. print(content, end="", flush=True)
  41. collected_content.append(content)
  42. print() # 在流式输出结束后换行
  43. return "".join(collected_content)
  44. except Exception as e:
  45. print(f"❌ 调用LLM API时发生错误: {e}")
  46. return None
  47. # --- 客户端使用示例 ---
  48. if __name__ == '__main__':
  49. try:
  50. llmClient = HelloAgentsLLM()
  51. exampleMessages = [
  52. {"role": "system", "content": "You are a helpful assistant that writes Python code."},
  53. {"role": "user", "content": "写一个快速排序算法"}
  54. ]
  55. print("--- 调用LLM ---")
  56. responseText = llmClient.think(exampleMessages)
  57. if responseText:
  58. print("\n\n--- 完整模型响应 ---")
  59. print(responseText)
  60. except ValueError as e:
  61. print(e)