36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
from openai import OpenAI
|
||
from typing import Dict, Any, Optional
|
||
|
||
class LLMClient:
|
||
"""LLM客户端，封装OpenAI API调用"""
|
||
|
||
def __init__(self, config: Dict[str, Any]):
|
||
self.config = config
|
||
self.client = OpenAI(
|
||
api_key=config.get('api_key'),
|
||
base_url=config.get('base_url')
|
||
)
|
||
|
||
def generate_command(self, user_input: str, prompt: Optional[str] = None) -> str:
|
||
"""根据用户输入生成命令"""
|
||
if not prompt:
|
||
prompt = self.config.get('default_prompt', '你现在是一个终端助手,用户输入想要生成的命令,你来输出一个命令,不要任何多余的文本!')
|
||
|
||
messages = [
|
||
{"role": "system", "content": prompt},
|
||
{"role": "user", "content": user_input}
|
||
]
|
||
|
||
try:
|
||
response = self.client.chat.completions.create(
|
||
model=self.config.get('model'),
|
||
messages=messages,
|
||
temperature=0.1,
|
||
max_tokens=100
|
||
)
|
||
|
||
command = response.choices[0].message.content.strip()
|
||
return command
|
||
except Exception as e:
|
||
raise Exception(f"LLM调用失败: {str(e)}")
|