支持DeepSeek

2025-12-10 16:24:06 +08:00
parent ff3ec65139
commit a53bf1d6e4
5 changed files with 363 additions and 560 deletions
@@ -1,73 +1,66 @@
-import Anthropic from '@anthropic-ai/sdk';
-import type { Tool, ToolResult, Message, AgentConfig } from '../types/index.js';
+import { createAnthropic } from '@ai-sdk/anthropic';
+import { createDeepSeek } from '@ai-sdk/deepseek';
+import { generateText, streamText, type ModelMessage, type Tool as AITool, type LanguageModel } from 'ai';
+import type { Tool, ToolResult, Message, AgentConfig, ProviderType } from '../types/index.js';
+import { buildZodSchema } from '../types/index.js';
+
+// Provider 工厂函数类型
+type ProviderFactory = (apiKey: string) => (model: string) => LanguageModel;
+
+// Provider 注册表
+const providers: Record<ProviderType, ProviderFactory> = {
+  anthropic: (apiKey) => {
+    const client = createAnthropic({ apiKey });
+    return (model) => client(model);
+  },
+  deepseek: (apiKey) => {
+    const client = createDeepSeek({ apiKey });
+    return (model) => client(model);
+  },
+};

 export class Agent {
-  private client: Anthropic;
+  private getModel: (model: string) => LanguageModel;
  private config: AgentConfig;
  private tools: Map<string, Tool> = new Map();
-  private conversationHistory: Message[] = [];
+  private conversationHistory: ModelMessage[] = [];

  constructor(config: AgentConfig) {
    this.config = config;
-    this.client = new Anthropic({
-      apiKey: config.apiKey,
-    });
+
+    const providerFactory = providers[config.provider];
+    if (!providerFactory) {
+      throw new Error(`不支持的 provider: ${config.provider}`);
+    }
+    this.getModel = providerFactory(config.apiKey);
  }

  // 注册工具
-  registerTool(tool: Tool): void {
-    this.tools.set(tool.name, tool);
+  registerTool(customTool: Tool): void {
+    this.tools.set(customTool.name, customTool);
  }

-  // 获取所有工具定义（用于 Claude API）
-  private getToolDefinitions(): Anthropic.Tool[] {
-    return Array.from(this.tools.values()).map((tool) => ({
-      name: tool.name,
-      description: tool.description,
-      input_schema: {
-        type: 'object' as const,
-        properties: Object.fromEntries(
-          Object.entries(tool.parameters).map(([key, param]) => [
-            key,
-            {
-              type: param.type,
-              description: param.description,
-            },
-          ])
-        ),
-        required: Object.entries(tool.parameters)
-          .filter(([, param]) => param.required)
-          .map(([key]) => key),
-      },
-    }));
-  }
+  // 将自定义工具转换为 Vercel AI SDK 的工具格式
+  private getVercelTools(): Record<string, AITool> {
+    const vercelTools: Record<string, AITool> = {};

-  // 执行工具
-  private async executeTool(
-    toolName: string,
-    input: Record<string, unknown>
-  ): Promise<ToolResult> {
-    const tool = this.tools.get(toolName);
-    if (!tool) {
-      return {
-        success: false,
-        output: '',
-        error: `Tool "${toolName}" not found`,
-      };
+    for (const [name, customTool] of this.tools) {
+      const schema = buildZodSchema(customTool.parameters);
+
+      vercelTools[name] = {
+        description: customTool.description,
+        inputSchema: schema,
+        execute: async (params) => {
+          const result = await customTool.execute(params as Record<string, unknown>);
+          return result;
+        },
+      } as AITool;
    }

-    try {
-      return await tool.execute(input);
-    } catch (error) {
-      return {
-        success: false,
-        output: '',
-        error: error instanceof Error ? error.message : String(error),
-      };
-    }
+    return vercelTools;
  }

-  // 发送消息并处理响应
+  // 发送消息并处理响应（流式）
  async chat(
    userMessage: string,
    onStream?: (text: string) => void
@@ -78,84 +71,52 @@ export class Agent {
      content: userMessage,
    });

-    const messages: Anthropic.MessageParam[] = this.conversationHistory.map(
-      (msg) => ({
-        role: msg.role,
-        content: msg.content,
-      })
-    );
-
+    const vercelTools = this.getVercelTools();
    let fullResponse = '';

-    // 循环处理，直到没有工具调用
-    while (true) {
-      const response = await this.client.messages.create({
-        model: this.config.model,
-        max_tokens: this.config.maxTokens,
+    if (onStream) {
+      // 流式模式
+      const result = streamText({
+        model: this.getModel(this.config.model),
        system: this.config.systemPrompt,
-        tools: this.getToolDefinitions(),
-        messages,
-      });
-
-      // 处理响应内容
-      const textBlocks: string[] = [];
-      const toolUseBlocks: Anthropic.ToolUseBlock[] = [];
-
-      for (const block of response.content) {
-        if (block.type === 'text') {
-          textBlocks.push(block.text);
-          if (onStream) {
-            onStream(block.text);
+        messages: this.conversationHistory,
+        tools: vercelTools,
+        maxOutputTokens: this.config.maxTokens,
+        onChunk: ({ chunk }) => {
+          if (chunk.type === 'tool-call') {
+            onStream(`\n[调用工具: ${chunk.toolName}]\n`);
+          } else if (chunk.type === 'tool-result') {
+            const output = (chunk as { output?: ToolResult }).output;
+            if (output && typeof output === 'object') {
+              if (output.success) {
+                onStream(`[结果: ${output.output}]\n`);
+              } else {
+                onStream(`[错误: ${output.error}]\n`);
+              }
+            }
          }
-        } else if (block.type === 'tool_use') {
-          toolUseBlocks.push(block);
-        }
-      }
-
-      fullResponse += textBlocks.join('');
-
-      // 如果没有工具调用，结束循环
-      if (toolUseBlocks.length === 0) {
-        break;
-      }
-
-      // 添加 assistant 消息
-      messages.push({
-        role: 'assistant',
-        content: response.content,
+        },
      });

-      // 处理工具调用
-      const toolResults: Anthropic.ToolResultBlockParam[] = [];
-
-      for (const toolUse of toolUseBlocks) {
-        if (onStream) {
-          onStream(`\n[调用工具: ${toolUse.name}]\n`);
-        }
-
-        const result = await this.executeTool(
-          toolUse.name,
-          toolUse.input as Record<string, unknown>
-        );
-
-        if (onStream) {
-          onStream(
-            result.success ? `[结果: ${result.output}]\n` : `[错误: ${result.error}]\n`
-          );
-        }
-
-        toolResults.push({
-          type: 'tool_result',
-          tool_use_id: toolUse.id,
-          content: result.success ? result.output : `Error: ${result.error}`,
-        });
+      // 流式输出文本
+      for await (const chunk of result.textStream) {
+        fullResponse += chunk;
+        onStream(chunk);
      }

-      // 添加工具结果
-      messages.push({
-        role: 'user',
-        content: toolResults,
+      // 等待完成
+      await result.response;
+    } else {
+      // 非流式模式
+      const result = await generateText({
+        model: this.getModel(this.config.model),
+        system: this.config.systemPrompt,
+        messages: this.conversationHistory,
+        tools: vercelTools,
+        maxOutputTokens: this.config.maxTokens,
      });
+
+      fullResponse = result.text;
    }

    // 保存助手响应到历史
@@ -174,6 +135,13 @@ export class Agent {

  // 获取对话历史
  getHistory(): Message[] {
-    return [...this.conversationHistory];
+    return this.conversationHistory
+      .filter((msg): msg is ModelMessage & { role: 'user' | 'assistant' } =>
+        msg.role === 'user' || msg.role === 'assistant'
+      )
+      .map((msg) => ({
+        role: msg.role,
+        content: typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content),
+      }));
  }
 }
@@ -1,37 +1,45 @@
+import { z } from 'zod';
+
 // 消息类型
 export interface Message {
  role: 'user' | 'assistant';
  content: string;
 }

-// 工具定义
-export interface Tool {
-  name: string;
-  description: string;
-  parameters: Record<string, ToolParameter>;
-  execute: (params: Record<string, unknown>) => Promise<ToolResult>;
-}
-
+// 工具参数定义
 export interface ToolParameter {
  type: 'string' | 'number' | 'boolean' | 'array' | 'object';
  description: string;
  required?: boolean;
 }

+// 工具执行结果
 export interface ToolResult {
  success: boolean;
  output: string;
  error?: string;
 }

+// 工具定义（兼容 Vercel AI SDK 的 tool 格式）
+export interface Tool {
+  name: string;
+  description: string;
+  parameters: Record<string, ToolParameter>;
+  execute: (params: Record<string, unknown>) => Promise<ToolResult>;
+}
+
 // 工具调用请求
 export interface ToolCall {
  name: string;
  input: Record<string, unknown>;
 }

+// 支持的 Provider 类型
+export type ProviderType = 'anthropic' | 'deepseek';
+
 // Agent 配置
 export interface AgentConfig {
+  provider: ProviderType;
  apiKey: string;
  model: string;
  maxTokens: number;
@@ -43,3 +51,40 @@ export interface ConversationContext {
  messages: Message[];
  workingDirectory: string;
 }
+
+// 将自定义 Tool 转换为 Vercel AI SDK 的 zod schema
+export function buildZodSchema(parameters: Record<string, ToolParameter>): z.ZodObject<Record<string, z.ZodTypeAny>> {
+  const schemaObj: Record<string, z.ZodTypeAny> = {};
+
+  for (const [key, param] of Object.entries(parameters)) {
+    let fieldSchema: z.ZodTypeAny;
+
+    switch (param.type) {
+      case 'string':
+        fieldSchema = z.string().describe(param.description);
+        break;
+      case 'number':
+        fieldSchema = z.number().describe(param.description);
+        break;
+      case 'boolean':
+        fieldSchema = z.boolean().describe(param.description);
+        break;
+      case 'array':
+        fieldSchema = z.array(z.unknown()).describe(param.description);
+        break;
+      case 'object':
+        fieldSchema = z.record(z.string(), z.unknown()).describe(param.description);
+        break;
+      default:
+        fieldSchema = z.unknown().describe(param.description);
+    }
+
+    if (!param.required) {
+      fieldSchema = fieldSchema.optional();
+    }
+
+    schemaObj[key] = fieldSchema;
+  }
+
+  return z.object(schemaObj);
+}
@@ -1,17 +1,25 @@
 import * as fs from 'fs';
 import * as path from 'path';
 import * as os from 'os';
-import type { AgentConfig } from '../types/index.js';
+import type { AgentConfig, ProviderType } from '../types/index.js';

 const CONFIG_DIR = path.join(os.homedir(), '.ai-terminal-assistant');
 const CONFIG_FILE = path.join(CONFIG_DIR, 'config.json');

 interface StoredConfig {
+  provider?: ProviderType;
  apiKey?: string;
+  deepseekApiKey?: string;
  model?: string;
  maxTokens?: number;
 }

+// 默认模型配置
+const DEFAULT_MODELS: Record<ProviderType, string> = {
+  anthropic: 'claude-sonnet-4-20250514',
+  deepseek: 'deepseek-chat',
+};
+
 // 默认系统提示词
 const DEFAULT_SYSTEM_PROMPT = `你是一个运行在终端中的 AI 编程助手。你可以帮助用户：
 - 读取和写入文件
@@ -29,12 +37,14 @@ const DEFAULT_SYSTEM_PROMPT = `你是一个运行在终端中的 AI 编程助手

 // 加载配置
 export function loadConfig(): AgentConfig {
-  // 优先从环境变量获取
-  const apiKey = process.env.ANTHROPIC_API_KEY;
-  const model = process.env.AI_MODEL || 'claude-sonnet-4-20250514';
+  // 从环境变量获取
+  const provider = (process.env.AI_PROVIDER as ProviderType) || 'anthropic';
+  const anthropicApiKey = process.env.ANTHROPIC_API_KEY;
+  const deepseekApiKey = process.env.DEEPSEEK_API_KEY;
+  const model = process.env.AI_MODEL;
  const maxTokens = parseInt(process.env.AI_MAX_TOKENS || '4096', 10);

-  // 如果环境变量没有，尝试从配置文件读取
+  // 从配置文件读取
  let storedConfig: StoredConfig = {};
  if (fs.existsSync(CONFIG_FILE)) {
    try {
@@ -45,18 +55,32 @@ export function loadConfig(): AgentConfig {
    }
  }

-  const finalApiKey = apiKey || storedConfig.apiKey;
+  // 确定最终的 provider
+  const finalProvider = storedConfig.provider || provider;
+
+  // 根据 provider 获取对应的 API Key
+  let finalApiKey: string | undefined;
+  if (finalProvider === 'anthropic') {
+    finalApiKey = anthropicApiKey || storedConfig.apiKey;
+  } else if (finalProvider === 'deepseek') {
+    finalApiKey = deepseekApiKey || storedConfig.deepseekApiKey;
+  }

  if (!finalApiKey) {
-    console.error('❌ 错误: 未设置 ANTHROPIC_API_KEY');
-    console.error('请设置环境变量: export ANTHROPIC_API_KEY=your-api-key');
-    console.error('或运行: ai-assist --init 进行初始化配置');
+    const envVar = finalProvider === 'anthropic' ? 'ANTHROPIC_API_KEY' : 'DEEPSEEK_API_KEY';
+    console.error(`❌ 错误: 未设置 ${envVar}`);
+    console.error(`请设置环境变量: export ${envVar}=your-api-key`);
+    console.error('或运行: ai-assist init 进行初始化配置');
    process.exit(1);
  }

+  // 确定模型
+  const finalModel = model || storedConfig.model || DEFAULT_MODELS[finalProvider];
+
  return {
+    provider: finalProvider,
    apiKey: finalApiKey,
-    model: storedConfig.model || model,
+    model: finalModel,
    maxTokens: storedConfig.maxTokens || maxTokens,
    systemPrompt: DEFAULT_SYSTEM_PROMPT,
  };
@@ -91,24 +115,52 @@ export async function initConfig(): Promise<void> {

  console.log('\n🔧 初始化 AI Terminal Assistant 配置\n');

+  // 选择 provider
+  const { provider } = await inquirer.prompt([
+    {
+      type: 'list',
+      name: 'provider',
+      message: '选择 AI 服务商:',
+      choices: [
+        { name: 'Anthropic (Claude)', value: 'anthropic' },
+        { name: 'DeepSeek', value: 'deepseek' },
+      ],
+      default: 'anthropic',
+    },
+  ]);
+
+  // 根据 provider 显示不同的模型选项
+  const modelChoices =
+    provider === 'anthropic'
+      ? [
+          { name: 'Claude Sonnet 4 (推荐，平衡性能和成本)', value: 'claude-sonnet-4-20250514' },
+          { name: 'Claude Opus 4 (最强，成本较高)', value: 'claude-opus-4-20250514' },
+          { name: 'Claude 3.5 Haiku (快速，成本低)', value: 'claude-3-5-haiku-20241022' },
+        ]
+      : [
+          { name: 'DeepSeek Chat (推荐)', value: 'deepseek-chat' },
+          { name: 'DeepSeek Reasoner (推理增强)', value: 'deepseek-reasoner' },
+        ];
+
+  const apiKeyField = provider === 'anthropic' ? 'apiKey' : 'deepseekApiKey';
+  const apiKeyMessage =
+    provider === 'anthropic'
+      ? '请输入你的 Anthropic API Key:'
+      : '请输入你的 DeepSeek API Key:';
+
  const answers = await inquirer.prompt([
    {
      type: 'password',
-      name: 'apiKey',
-      message: '请输入你的 Anthropic API Key:',
-      validate: (input: string) =>
-        input.length > 0 || 'API Key 不能为空',
+      name: apiKeyField,
+      message: apiKeyMessage,
+      validate: (input: string) => input.length > 0 || 'API Key 不能为空',
    },
    {
      type: 'list',
      name: 'model',
      message: '选择默认模型:',
-      choices: [
-        { name: 'Claude Sonnet 4 (推荐，平衡性能和成本)', value: 'claude-sonnet-4-20250514' },
-        { name: 'Claude Opus 4 (最强，成本较高)', value: 'claude-opus-4-20250514' },
-        { name: 'Claude 3.5 Haiku (快速，成本低)', value: 'claude-3-5-haiku-20241022' },
-      ],
-      default: 'claude-sonnet-4-20250514',
+      choices: modelChoices,
+      default: DEFAULT_MODELS[provider as ProviderType],
    },
    {
      type: 'number',
@@ -118,7 +170,7 @@ export async function initConfig(): Promise<void> {
    },
  ]);

-  saveConfig(answers);
+  saveConfig({ provider, ...answers });
  console.log('\n✅ 配置已保存到', CONFIG_FILE);
  console.log('现在可以运行 ai-assist 开始使用了！\n');
 }