| import os
|
|
|
| from openai import OpenAI
|
|
|
|
|
| def openAIChatLLM(model_name=None, api_key=None, base_url=None):
|
| """
|
| model_name 取值
|
| - deepseek-chat
|
| """
|
| api_key = os.environ.get("OPENAI_API_KEY", api_key)
|
| base_url = os.environ.get("OPENAI_BASE_URL", base_url)
|
|
|
| client = OpenAI(api_key=api_key, base_url=base_url)
|
|
|
| def chatLLM(
|
| messages: list,
|
| temperature=None,
|
| top_p=None,
|
| max_tokens=None,
|
| stream=False,
|
| model=model_name,
|
| ) -> dict:
|
| if not stream:
|
| response = client.chat.completions.create(
|
|
|
| model=model,
|
| messages=messages,
|
| temperature=temperature,
|
| top_p=top_p,
|
| max_tokens=max_tokens,
|
| )
|
| return {
|
| "content": response.choices[0].message.content,
|
| "total_tokens": response.usage.total_tokens,
|
| }
|
| else:
|
| responses = client.chat.completions.create(
|
|
|
| model=model,
|
| messages=messages,
|
| temperature=temperature,
|
| top_p=top_p,
|
| max_tokens=max_tokens,
|
| stream=True,
|
| )
|
|
|
| def respGenerator():
|
| content = ""
|
| for response in responses:
|
| delta = response.choices[0].delta.content
|
|
|
|
|
| if delta is not None:
|
| content += delta
|
|
|
|
|
|
|
|
|
|
|
| total_tokens = None
|
|
|
| yield {
|
| "content": content,
|
| "total_tokens": total_tokens,
|
| }
|
|
|
| return respGenerator()
|
|
|
| return chatLLM
|
|
|
