From 22667064949e060b9f06c7ec3e660ec98fbfd57d Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 5 Feb 2025 09:28:56 +0000
Subject: [PATCH] refactor: move LLM client implementations to utils directory

Co-Authored-By: Han Xiao <han.xiao@jina.ai>
---
 src/config.ts           | 106 +---------------------------------------
 src/utils/llm-client.ts | 104 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+), 105 deletions(-)
 create mode 100644 src/utils/llm-client.ts

diff --git a/src/config.ts b/src/config.ts
index 2529f31..362abab 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -1,5 +1,6 @@
 import dotenv from 'dotenv';
 import { ProxyAgent, setGlobalDispatcher } from 'undici';
+import { GoogleAIWrapper, LocalLLMClient, LLMClient } from './utils/llm-client';
 
 interface ModelConfig {
   model: string;
@@ -15,111 +16,6 @@ interface ToolConfigs {
   agentBeastMode: ModelConfig;
 }
 
-import { GoogleGenerativeAI } from '@google/generative-ai';
-
-interface LLMClientConfig {
-  model: string;
-  temperature: number;
-  generationConfig?: {
-    responseMimeType?: string;
-    responseSchema?: any;
-  };
-}
-
-interface LLMResponse {
-  text(): string;
-  usageMetadata: {
-    totalTokenCount: number;
-  };
-}
-
-interface LLMClient {
-  getGenerativeModel(config: LLMClientConfig): {
-    generateContent(prompt: string): Promise<{
-      response: LLMResponse;
-    }>;
-  };
-}
-
-class GoogleAIWrapper implements LLMClient {
-  private client: GoogleGenerativeAI;
-
-  constructor(apiKey: string) {
-    this.client = new GoogleGenerativeAI(apiKey);
-  }
-
-  getGenerativeModel(config: LLMClientConfig) {
-    const model = this.client.getGenerativeModel({
-      model: config.model,
-      generationConfig: {
-        temperature: config.temperature,
-        ...(config.generationConfig || {})
-      }
-    });
-
-    return {
-      generateContent: async (prompt: string) => {
-        const result = await model.generateContent(prompt);
-        return {
-          response: {
-            text: () => result.response.text(),
-            usageMetadata: {
-              totalTokenCount: result.response.usageMetadata?.totalTokenCount ?? 0
-            }
-          }
-        };
-      }
-    };
-  }
-}
-
-class LocalLLMClient implements LLMClient {
-  constructor(
-    private hostname: string,
-    private port: string,
-    private model: string
-  ) {}
-
-  getGenerativeModel(config: LLMClientConfig) {
-    return {
-      generateContent: async (prompt: string) => {
-        const response = await fetch(`http://${this.hostname}:${this.port}/v1/chat/completions`, {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/json',
-          },
-          body: JSON.stringify({
-            model: this.model,
-            messages: [
-              {
-                role: 'user',
-                content: prompt,
-              },
-            ],
-            temperature: config.temperature,
-            response_format: {
-              type: 'json_schema',
-              json_schema: config.generationConfig?.responseSchema,
-            },
-            max_tokens: 1000,
-            stream: false,
-          }),
-        });
-
-        const data = await response.json();
-        return {
-          response: {
-            text: () => data.choices[0].message.content,
-            usageMetadata: {
-              totalTokenCount: data.usage?.total_tokens || 0,
-            },
-          },
-        };
-      },
-    };
-  }
-}
-
 
 dotenv.config();
 
diff --git a/src/utils/llm-client.ts b/src/utils/llm-client.ts
new file mode 100644
index 0000000..fd8ad31
--- /dev/null
+++ b/src/utils/llm-client.ts
@@ -0,0 +1,104 @@
+import { GoogleGenerativeAI } from '@google/generative-ai';
+
+export interface LLMClientConfig {
+  model: string;
+  temperature: number;
+  generationConfig?: {
+    responseMimeType?: string;
+    responseSchema?: any;
+  };
+}
+
+export interface LLMResponse {
+  text(): string;
+  usageMetadata: {
+    totalTokenCount: number;
+  };
+}
+
+export interface LLMClient {
+  getGenerativeModel(config: LLMClientConfig): {
+    generateContent(prompt: string): Promise<{
+      response: LLMResponse;
+    }>;
+  };
+}
+
+export class GoogleAIWrapper implements LLMClient {
+  private client: GoogleGenerativeAI;
+
+  constructor(apiKey: string) {
+    this.client = new GoogleGenerativeAI(apiKey);
+  }
+
+  getGenerativeModel(config: LLMClientConfig) {
+    const model = this.client.getGenerativeModel({
+      model: config.model,
+      generationConfig: {
+        temperature: config.temperature,
+        ...(config.generationConfig || {})
+      }
+    });
+
+    return {
+      generateContent: async (prompt: string) => {
+        const result = await model.generateContent(prompt);
+        return {
+          response: {
+            text: () => result.response.text(),
+            usageMetadata: {
+              totalTokenCount: result.response.usageMetadata?.totalTokenCount ?? 0
+            }
+          }
+        };
+      }
+    };
+  }
+}
+
+export class LocalLLMClient implements LLMClient {
+  constructor(
+    private hostname: string,
+    private port: string,
+    private model: string
+  ) {}
+
+  getGenerativeModel(config: LLMClientConfig) {
+    return {
+      generateContent: async (prompt: string) => {
+        const response = await fetch(`http://${this.hostname}:${this.port}/v1/chat/completions`, {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+          },
+          body: JSON.stringify({
+            model: this.model,
+            messages: [
+              {
+                role: 'user',
+                content: prompt,
+              },
+            ],
+            temperature: config.temperature,
+            response_format: {
+              type: 'json_schema',
+              json_schema: config.generationConfig?.responseSchema,
+            },
+            max_tokens: 1000,
+            stream: false,
+          }),
+        });
+
+        const data = await response.json();
+        return {
+          response: {
+            text: () => data.choices[0].message.content,
+            usageMetadata: {
+              totalTokenCount: data.usage?.total_tokens || 0,
+            },
+          },
+        };
+      },
+    };
+  }
+}