66 lines
3.5 KiB
Python
Raw Normal View History

2025-08-23 19:18:29 +08:00
from openai import OpenAI
from typing import Dict, Any, Optional, List
import os
2025-08-23 19:18:29 +08:00
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,
history: Optional[List[Dict[str, Any]]] = None,
current_dir_content: Optional[List[str]] = None,
last_executed_command: str = "") -> str:
2025-08-23 19:18:29 +08:00
"""根据用户输入生成命令"""
# 根据用户输入是否为空选择不同的提示词
if not user_input:
if not prompt:
prompt = self.config.get('recommendation_prompt', '你现在是一个终端助手,根据上下文自动推荐命令:当用户没有输入时,基于最近执行的命令历史和当前目录内容,智能推荐最可能需要的终端命令(仅当有明确上下文线索时);当用户输入命令需求时,生成对应命令。仅输出纯命令文本,不要任何解释或多余内容!')
else:
if not prompt:
prompt = self.config.get('default_prompt', '你现在是一个终端助手,用户输入想要生成的命令,你来输出一个命令,不要任何多余的文本!')
# 构建系统提示,包含上下文信息
system_prompt = prompt
# 添加历史命令上下文
if history:
history_context = "\n最近执行的命令历史:\n"
for i, entry in enumerate(reversed(history), 1):
history_context += f"{i}. 用户输入: {entry.get('user_input', '')} -> 生成命令: {entry.get('generated_command', '')}\n"
system_prompt += history_context
# 添加当前目录内容上下文
if current_dir_content:
dir_context = "\n当前目录下的文件和文件夹:\n" + "\n".join(current_dir_content)
system_prompt += dir_context
# 当用户输入为空时,使用特殊的提示来触发推荐模式
if not user_input:
user_content = f"根据提供的上下文信息推荐一个最可能需要的终端命令仅当有明确的上下文线索时。如果上下文信息不足以确定一个有用的命令则返回空。请直接返回一个可执行的终端命令不要包含任何解释或其他文本。例如ls -la 或 git status。特别注意不要使用echo命令来列出文件应该使用ls命令。推荐命令时请考虑最近执行的命令历史避免重复推荐相同的命令。最后执行的命令是: {last_executed_command}。如果当前目录有pyproject.toml或setup.py文件可以考虑使用pip list查看已安装的包。"
else:
user_content = user_input
2025-08-23 19:18:29 +08:00
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_content}
2025-08-23 19:18:29 +08:00
]
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)}")