chore: 优化agent配置

TBXark · Oct 23, 2024 · a687303 · a687303
1 parent db6c8bc
commit a687303
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 17 deletions.
diff --git a/src/agent/cohere.ts b/src/agent/cohere.ts
@@ -7,11 +7,6 @@ export class Cohere implements ChatAgent {
     readonly name = 'cohere';
     readonly modelKey = 'COHERE_CHAT_MODEL';
 
-    static COHERE_ROLE_MAP: Record<string, string> = {
-        assistant: 'CHATBOT',
-        user: 'USER',
-    };
-
     readonly enable = (context: AgentUserConfig): boolean => {
         return !!(context.COHERE_API_KEY);
     };

diff --git a/src/agent/gemini.ts b/src/agent/gemini.ts
@@ -12,7 +12,7 @@ export class Gemini implements ChatAgent {
     };
 
     readonly model = (ctx: AgentUserConfig): string => {
-        return ctx.GOOGLE_COMPLETIONS_MODEL;
+        return ctx.GOOGLE_CHAT_MODEL;
     };
 
     readonly request = async (params: LLMChatParams, context: AgentUserConfig, onStream: ChatStreamTextHandler | null): Promise<ResponseMessage[]> => {

diff --git a/src/agent/types.ts b/src/agent/types.ts
@@ -22,21 +22,16 @@ export interface LLMChatParams {
 export type ResponseMessage = CoreAssistantMessage | CoreToolMessage;
 
 export type ChatAgentRequest = (params: LLMChatParams, context: AgentUserConfig, onStream: ChatStreamTextHandler | null) => Promise<ResponseMessage[]>;
+export type ImageAgentRequest = (prompt: string, context: AgentUserConfig) => Promise<string | Blob>;
 
-export interface ChatAgent {
+export interface Agent<AgentRequest> {
     name: string;
     modelKey: string;
     enable: (context: AgentUserConfig) => boolean;
-    request: ChatAgentRequest;
+    request: AgentRequest;
     model: (ctx: AgentUserConfig) => string;
 }
 
-export type ImageAgentRequest = (prompt: string, context: AgentUserConfig) => Promise<string | Blob>;
+export type ChatAgent = Agent<ChatAgentRequest>;
 
-export interface ImageAgent {
-    name: string;
-    modelKey: string;
-    enable: (context: AgentUserConfig) => boolean;
-    request: ImageAgentRequest;
-    model: (ctx: AgentUserConfig) => string;
-}
+export type ImageAgent = Agent<ImageAgentRequest>;
diff --git a/src/config/config.ts b/src/config/config.ts
@@ -151,7 +151,7 @@ export class GeminiConfig {
     // Google Gemini API Base
     GOOGLE_API_BASE = 'https://generativelanguage.googleapis.com/v1beta';
     // Google Gemini Model
-    GOOGLE_COMPLETIONS_MODEL = 'gemini-pro';
+    GOOGLE_CHAT_MODEL = 'gemini-pro';
 }
 
 // -- Mistral 配置 --

diff --git a/src/config/env.ts b/src/config/env.ts
@@ -177,6 +177,11 @@ class Environment extends EnvironmentConfig {
         if (source.GOOGLE_COMPLETIONS_API && !this.USER_CONFIG.GOOGLE_API_BASE) {
             this.USER_CONFIG.GOOGLE_API_BASE = source.GOOGLE_COMPLETIONS_API.replace(/\/models\/?$/, '');
         }
+
+        if (source.GOOGLE_COMPLETIONS_MODEL && !this.USER_CONFIG.GOOGLE_CHAT_MODEL) {
+            this.USER_CONFIG.GOOGLE_CHAT_MODEL = source.GOOGLE_COMPLETIONS_MODEL;
+        }
+
         // 兼容旧版 AZURE_COMPLETIONS_API
         if (source.AZURE_COMPLETIONS_API && !this.USER_CONFIG.AZURE_CHAT_MODEL) {
             const url = new URL(source.AZURE_COMPLETIONS_API);