llm_client.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  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. if not chunk.choices:
  40. continue
  41. content = chunk.choices[0].delta.content or ""
  42. print(content, end="", flush=True)
  43. collected_content.append(content)
  44. print() # 在流式输出结束后换行
  45. return "".join(collected_content)
  46. except Exception as e:
  47. print(f"❌ 调用LLM API时发生错误: {e}")
  48. return None
  49. # --- 客户端使用示例 ---
  50. if __name__ == '__main__':
  51. try:
  52. llmClient = HelloAgentsLLM()
  53. exampleMessages = [
  54. {"role": "system", "content": "You are a helpful assistant that writes Python code."},
  55. {"role": "user", "content": "写一个快速排序算法"}
  56. ]
  57. print("--- 调用LLM ---")
  58. responseText = llmClient.think(exampleMessages)
  59. if responseText:
  60. print("\n\n--- 完整模型响应 ---")
  61. print(responseText)
  62. except ValueError as e:
  63. print(e)