From 1b7ae5457f68f43017f99b70d4ffbfd06a28497c Mon Sep 17 00:00:00 2001
From: rentianyue-jk <rentianyue-jk@360shuke.com>
Date: Tue, 16 Jul 2024 10:22:54 +0800
Subject: [PATCH] support custom models

---
 apps/api/src/lib/LLM-extraction/models.ts | 2 +-
 docker-compose.yaml                       | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/apps/api/src/lib/LLM-extraction/models.ts b/apps/api/src/lib/LLM-extraction/models.ts
index 8de8ee4b..d190c495 100644
--- a/apps/api/src/lib/LLM-extraction/models.ts
+++ b/apps/api/src/lib/LLM-extraction/models.ts
@@ -48,7 +48,7 @@ function prepareOpenAIDoc(
 
 export async function generateOpenAICompletions({
   client,
-  model = "gpt-4o",
+  model = process.env.MODEL_NAME,
   document,
   schema, //TODO - add zod dynamic type checking
   prompt = defaultPrompt,
diff --git a/docker-compose.yaml b/docker-compose.yaml
index b88f3ed8..b785308a 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -12,6 +12,8 @@ x-common-service: &common-service
     - PORT=${PORT:-3002}
     - NUM_WORKERS_PER_QUEUE=${NUM_WORKERS_PER_QUEUE}
     - OPENAI_API_KEY=${OPENAI_API_KEY}
+    - OPENAI_BASE_URL=${OPENAI_BASE_URL}
+    - MODEL_NAME=${MODEL_NAME:-gpt-4o}
     - SLACK_WEBHOOK_URL=${SLACK_WEBHOOK_URL}
     - SERPER_API_KEY=${SERPER_API_KEY}
     - LLAMAPARSE_API_KEY=${LLAMAPARSE_API_KEY}