refactor: 解耦AI模块代码，为后续拓展AI功能，复用代码提供基础

Crokily · Crokily · commit d2e764af165b · 2025-09-26T17:09:07.000+10:00
diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
@@ -1,89 +1,52 @@
-import { createOpenAI } from "@ai-sdk/openai";
-import { createGoogleGenerativeAI } from "@ai-sdk/google";
-import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
 import { streamText, UIMessage, convertToModelMessages } from "ai";
+import { getModel, requiresApiKey, type AIProvider } from "@/lib/ai/models";
+import { buildSystemMessage } from "@/lib/ai/prompt";
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 30;
 
-export async function POST(req: Request) {
-  const {
-    messages,
-    system,
-    pageContext,
-    provider,
-    apiKey,
-  }: {
-    messages: UIMessage[];
-    system?: string; // System message forwarded from AssistantChatTransport
-    tools?: unknown; // Frontend tools forwarded from AssistantChatTransport
-    pageContext?: {
-      title?: string;
-      description?: string;
-      content?: string;
-      slug?: string;
-    };
-    provider?: "openai" | "gemini" | "intern";
-    apiKey?: string;
-  } = await req.json();
-
-  // Check if API key is provided (not required for intern provider)
-  if (provider !== "intern" && (!apiKey || apiKey.trim() === "")) {
-    return Response.json(
-      {
-        error:
-          "API key is required. Please configure your API key in the settings.",
-      },
-      { status: 400 },
-    );
-  }
+interface ChatRequest {
+  messages: UIMessage[];
+  system?: string; // System message forwarded from AssistantChatTransport
+  tools?: unknown; // Frontend tools forwarded from AssistantChatTransport
+  pageContext?: {
+    title?: string;
+    description?: string;
+    content?: string;
+    slug?: string;
+  };
+  provider?: AIProvider;
+  apiKey?: string;
+}
 
+export async function POST(req: Request) {
   try {
-    // Build system message with page context
-    let systemMessage =
-      system ||
-      `You are a helpful AI assistant for a documentation website. 
-    You can help users understand the documentation, answer questions about the content, 
-    and provide guidance on the topics covered in the docs. Be concise and helpful.`;
+    const {
+      messages,
+      system,
+      pageContext,
+      provider = "openai", // Default to OpenAI
+      apiKey,
+    }: ChatRequest = await req.json();
 
-    // Add current page context if available
-    if (pageContext?.content) {
-      systemMessage += `\n\n--- CURRENT PAGE CONTEXT ---\n`;
-      if (pageContext.title) {
-        systemMessage += `Page Title: ${pageContext.title}\n`;
-      }
-      if (pageContext.description) {
-        systemMessage += `Page Description: ${pageContext.description}\n`;
-      }
-      if (pageContext.slug) {
-        systemMessage += `Page URL: /docs/${pageContext.slug}\n`;
-      }
-      systemMessage += `Page Content:\n${pageContext.content}`;
-      systemMessage += `\n--- END OF CONTEXT ---\n\nWhen users ask about "this page", "current page", or refer to the content they're reading, use the above context to provide accurate answers. You can summarize, explain, or answer specific questions about the current page content.`;
+    // Validate API key for providers that require it
+    if (requiresApiKey(provider) && (!apiKey || apiKey.trim() === "")) {
+      return Response.json(
+        {
+          error:
+            "API key is required. Please configure your API key in the settings.",
+        },
+        { status: 400 },
+      );
     }
 
-    // Select model based on provider
-    let model;
-    if (provider === "gemini") {
-      const customGoogle = createGoogleGenerativeAI({
-        apiKey: apiKey,
-      });
-      model = customGoogle("models/gemini-2.0-flash");
-    } else if (provider === "intern") {
-      const intern = createOpenAICompatible({
-        name: "intern",
-        baseURL: "https://chat.intern-ai.org.cn/api/v1/",
-        apiKey: process.env.INTERN_KEY,
-      });
-      model = intern("intern-s1");
-    } else {
-      // Default to OpenAI
-      const customOpenAI = createOpenAI({
-        apiKey: apiKey,
-      });
-      model = customOpenAI("gpt-4.1-nano");
-    }
+    // Build system message with page context
+    const systemMessage = buildSystemMessage(system, pageContext);
+
+    // Get AI model instance based on provider
+    const model = getModel(provider, apiKey);
 
+    // Generate streaming response
     const result = streamText({
       model: model,
       system: systemMessage,
@@ -93,6 +56,12 @@ export async function POST(req: Request) {
     return result.toUIMessageStreamResponse();
   } catch (error) {
     console.error("Chat API error:", error);
+
+    // Handle specific model creation errors
+    if (error instanceof Error && error.message.includes("API key")) {
+      return Response.json({ error: error.message }, { status: 400 });
+    }
+
     return Response.json(
       { error: "Failed to process chat request" },
       { status: 500 },
diff --git a/lib/ai/models.ts b/lib/ai/models.ts
@@ -0,0 +1,43 @@
+import { createOpenAIModel } from "./providers/openai";
+import { createGeminiModel } from "./providers/gemini";
+import { createInternModel } from "./providers/intern";
+
+export type AIProvider = "openai" | "gemini" | "intern";
+
+/**
+ * Model factory that returns the appropriate AI model based on provider
+ * @param provider - The AI provider to use
+ * @param apiKey - API key (not required for intern provider)
+ * @returns Configured AI model instance
+ */
+export function getModel(provider: AIProvider, apiKey?: string) {
+  switch (provider) {
+    case "openai":
+      if (!apiKey || apiKey.trim() === "") {
+        throw new Error("OpenAI API key is required");
+      }
+      return createOpenAIModel(apiKey);
+
+    case "gemini":
+      if (!apiKey || apiKey.trim() === "") {
+        throw new Error("Gemini API key is required");
+      }
+      return createGeminiModel(apiKey);
+
+    case "intern":
+      // Intern provider doesn't need API key from user
+      return createInternModel();
+
+    default:
+      throw new Error(`Unsupported AI provider: ${provider}`);
+  }
+}
+
+/**
+ * Check if the given provider requires an API key from the user
+ * @param provider - The AI provider to check
+ * @returns True if API key is required, false otherwise
+ */
+export function requiresApiKey(provider: AIProvider): boolean {
+  return provider !== "intern";
+}
diff --git a/lib/ai/prompt.ts b/lib/ai/prompt.ts
@@ -0,0 +1,46 @@
+interface PageContext {
+  title?: string;
+  description?: string;
+  content?: string;
+  slug?: string;
+}
+
+/**
+ * Build system message with page context for AI assistant
+ * @param customSystem - Custom system message (optional)
+ * @param pageContext - Current page context (optional)
+ * @returns Complete system message string
+ */
+export function buildSystemMessage(
+  customSystem?: string,
+  pageContext?: PageContext,
+): string {
+  // Default system message for documentation assistant
+  let systemMessage =
+    customSystem ||
+    `You are a helpful AI assistant for a documentation website. 
+    You can help users understand the documentation, answer questions about the content, 
+    and provide guidance on the topics covered in the docs. Be concise and helpful.`;
+
+  // Add current page context if available
+  if (pageContext?.content) {
+    systemMessage += `\n\n--- CURRENT PAGE CONTEXT ---\n`;
+
+    if (pageContext.title) {
+      systemMessage += `Page Title: ${pageContext.title}\n`;
+    }
+
+    if (pageContext.description) {
+      systemMessage += `Page Description: ${pageContext.description}\n`;
+    }
+
+    if (pageContext.slug) {
+      systemMessage += `Page URL: /docs/${pageContext.slug}\n`;
+    }
+
+    systemMessage += `Page Content:\n${pageContext.content}`;
+    systemMessage += `\n--- END OF CONTEXT ---\n\nWhen users ask about "this page", "current page", or refer to the content they're reading, use the above context to provide accurate answers. You can summarize, explain, or answer specific questions about the current page content.`;
+  }
+
+  return systemMessage;
+}
diff --git a/lib/ai/providers/gemini.ts b/lib/ai/providers/gemini.ts
@@ -0,0 +1,15 @@
+import { createGoogleGenerativeAI } from "@ai-sdk/google";
+
+/**
+ * Create Google Gemini model instance
+ * @param apiKey - Google Gemini API key provided by user
+ * @returns Configured Gemini model instance
+ */
+export function createGeminiModel(apiKey: string) {
+  const customGoogle = createGoogleGenerativeAI({
+    apiKey: apiKey,
+  });
+
+  // Use the specific model configured for this project
+  return customGoogle("models/gemini-2.0-flash");
+}
diff --git a/lib/ai/providers/intern.ts b/lib/ai/providers/intern.ts
@@ -0,0 +1,17 @@
+import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
+
+/**
+ * Create Intern-AI model instance
+ * Uses environment variable INTERN_KEY for API key
+ * @returns Configured Intern-AI model instance
+ */
+export function createInternModel() {
+  const intern = createOpenAICompatible({
+    name: "intern",
+    baseURL: "https://chat.intern-ai.org.cn/api/v1/",
+    apiKey: process.env.INTERN_KEY,
+  });
+
+  // Use the specific model configured for this project
+  return intern("intern-s1");
+}
diff --git a/lib/ai/providers/openai.ts b/lib/ai/providers/openai.ts
@@ -0,0 +1,15 @@
+import { createOpenAI } from "@ai-sdk/openai";
+
+/**
+ * Create OpenAI model instance
+ * @param apiKey - OpenAI API key provided by user
+ * @returns Configured OpenAI model instance
+ */
+export function createOpenAIModel(apiKey: string) {
+  const customOpenAI = createOpenAI({
+    apiKey: apiKey,
+  });
+
+  // Use the specific model configured for this project
+  return customOpenAI("gpt-4.1-nano");
+}
