From ecfa26102825d12a0711cf0ae7ffe2c3ea6ed1b2 Mon Sep 17 00:00:00 2001
From: shouldnotappearcalm <1191465097@qq.com>
Date: Sun, 13 Oct 2024 13:30:50 +0800
Subject: [PATCH] feat(): support base url and model custom

---
 app/lib/.server/llm/model.ts       | 25 ++++++++++++++++++++++---
 app/lib/.server/llm/stream-text.ts |  3 +--
 worker-configuration.d.ts          |  2 ++
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/app/lib/.server/llm/model.ts b/app/lib/.server/llm/model.ts
index f0d695c47..16348247c 100644
--- a/app/lib/.server/llm/model.ts
+++ b/app/lib/.server/llm/model.ts
@@ -1,9 +1,28 @@
 import { createAnthropic } from '@ai-sdk/anthropic';
+import { getAPIKey } from '~/lib/.server/llm/api-key';
+import { env } from 'node:process';
 
-export function getAnthropicModel(apiKey: string) {
+export function getAnthropicModel(env: Env) {
   const anthropic = createAnthropic({
-    apiKey,
+    apiKey: getAPIKey(env),
+    baseURL: getBaseUrl(env) 
   });
+  return anthropic(getModel(env));
+}
+
+export function getBaseUrl(cloudflareEnv: Env) {
+  /**
+   * The `cloudflareEnv` is only used when deployed or when previewing locally.
+   * In development the environment variables are available through `env`.
+   */
+  return env.ANTHROPIC_BASE_URL || cloudflareEnv.ANTHROPIC_BASE_URL;
+}
+
 
-  return anthropic('claude-3-5-sonnet-20240620');
+export function getModel(cloudflareEnv: Env) {
+  /**
+   * The `cloudflareEnv` is only used when deployed or when previewing locally.
+   * In development the environment variables are available through `env`.
+   */
+  return env.ANTHROPIC_MODEL || cloudflareEnv.ANTHROPIC_MODEL || 'claude-3-5-sonnet-20240620';
 }
diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
index cf937fd00..7003b68a7 100644
--- a/app/lib/.server/llm/stream-text.ts
+++ b/app/lib/.server/llm/stream-text.ts
@@ -1,5 +1,4 @@
 import { streamText as _streamText, convertToCoreMessages } from 'ai';
-import { getAPIKey } from '~/lib/.server/llm/api-key';
 import { getAnthropicModel } from '~/lib/.server/llm/model';
 import { MAX_TOKENS } from './constants';
 import { getSystemPrompt } from './prompts';
@@ -23,7 +22,7 @@ export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;
 
 export function streamText(messages: Messages, env: Env, options?: StreamingOptions) {
   return _streamText({
-    model: getAnthropicModel(getAPIKey(env)),
+    model: getAnthropicModel(env),
     system: getSystemPrompt(),
     maxTokens: MAX_TOKENS,
     headers: {
diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts
index 606a4e521..c519099ef 100644
--- a/worker-configuration.d.ts
+++ b/worker-configuration.d.ts
@@ -1,3 +1,5 @@
 interface Env {
   ANTHROPIC_API_KEY: string;
+  ANTHROPIC_BASE_URL: string;
+  ANTHROPIC_MODEL: string;
 }