From 4e870f161c2902325b5593aa1e675541668e4a63 Mon Sep 17 00:00:00 2001
From: Yingbei <yingbei@acorn.io>
Date: Thu, 1 Aug 2024 15:45:46 -0700
Subject: [PATCH] update for local model usage

---
 .env                      |  5 ++++-
 README.md                 | 19 +++++++++++++++++--
 docker-compose-remote.yml |  1 +
 docker-compose.yaml       |  1 +
 pkg/server/task/run.go    |  1 +
 5 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/.env b/.env
index 1351cb1..d916909 100644
--- a/.env
+++ b/.env
@@ -4,7 +4,6 @@ DB_PASSWORD=admin123
 DB_HOST=db
 
 MICROSOFT_CLIENT_ID=${CLIENT_ID}
-OPENAI_API_KEY=${OPENAI_API_KEY}
 MICROSOFT_CLIENT_SECRET=${CLIENT_SECRET}
 MICROSOFT_JWT_KEY=${JWT_SECRET_KEY}
 MICROSOFT_TENANT_ID=${TENANT_ID}
@@ -13,3 +12,7 @@ DEVELOPMENT=true
 
 PUBLIC_URL=${PUBLIC_URL}
 UI_SERVER=http://ui:3000
+
+OPENAI_API_KEY=${OPENAI_API_KEY}
+OPENAI_BASE_URL=https://api.openai.com/v1/
+DEFAULT_MODEL=gpt-4o
\ No newline at end of file
diff --git a/README.md b/README.md
index 5e1a942..f48338a 100644
--- a/README.md
+++ b/README.md
@@ -52,7 +52,6 @@ DB_PASSWORD=admin123
 DB_HOST=db
 
 MICROSOFT_CLIENT_ID=${CLIENT_ID}
-OPENAI_API_KEY=${OPENAI_API_KEY}
 MICROSOFT_CLIENT_SECRET=${CLIENT_SECRET}
 MICROSOFT_JWT_KEY=${JWT_SECRET_KEY}
 MICROSOFT_TENANT_ID=${TENANT_ID}
@@ -61,14 +60,30 @@ DEVELOPMENT=true
 
 PUBLIC_URL=${PUBLIC_URL}
 UI_SERVER=http://ui:3000
+
+OPENAI_API_KEY=${OPENAI_API_KEY}
+OPENAI_BASE_URL=https://api.openai.com/v1/
+DEFAULT_MODEL=gpt-4o
 ```
 
 Then run docker-compose:
 
 ```bash
-docker compose -f docker-compose-local.yaml up
+docker compose -f docker-compose.yaml up
 ```
 
 Go to `http://localhost:8080` and you can start logging in and using the app.
 
+### Using a Different Model
+
+The default model is `gpt-4o`. To use a different OpenAI model, update the `DEFAULT_MODEL` in the `.env` file.
+
+To connect to an OpenAI-compatible local model server, such as llama.cpp, ollama, or Rubra's tool.cpp, update the `OPENAI_API_KEY`, `OPENAI_BASE_URL`, and `DEFAULT_MODEL` accordingly. For example:
+
+```
+OPENAI_API_KEY=sk-123
+OPENAI_BASE_URL=http://host.docker.internal:1234/v1
+DEFAULT_MODEL=rubra-meta-llama-3-8b-instruct.Q8_0.gguf
+```
+
 ---
diff --git a/docker-compose-remote.yml b/docker-compose-remote.yml
index 825ad03..a90352d 100644
--- a/docker-compose-remote.yml
+++ b/docker-compose-remote.yml
@@ -39,6 +39,7 @@ services:
       PUBLIC_URL: ${PUBLIC_URL}
       UI_SERVER: ${UI_SERVER}
       OPENAI_API_KEY: ${OPENAI_API_KEY}
+      DEFAULT_MODEL: ${DEFAULT_MODEL}
     env_file:
       - .env
 
diff --git a/docker-compose.yaml b/docker-compose.yaml
index bbfbf85..abfd109 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -43,6 +43,7 @@ services:
       UI_SERVER: ${UI_SERVER}
       OPENAI_API_KEY: ${OPENAI_API_KEY}
       DEVELOPMENT: "true"
+      DEFAULT_MODEL: ${DEFAULT_MODEL}
     env_file:
       - .env
 
diff --git a/pkg/server/task/run.go b/pkg/server/task/run.go
index 24b6e4d..283bc90 100644
--- a/pkg/server/task/run.go
+++ b/pkg/server/task/run.go
@@ -53,6 +53,7 @@ func (h *Handler) RunTask(w http.ResponseWriter, r *http.Request) {
 
 	client, err := gptscript.NewGPTScript(gptscript.GlobalOptions{
 		OpenAIAPIKey: os.Getenv("OPENAI_API_KEY"),
+		DefaultModel: os.Getenv("DEFAULT_MODEL"),
 		Env:          append(os.Environ(), fmt.Sprintf("GPTSCRIPT_GRAPH_MICROSOFT_COM_BEARER_TOKEN=%v", user.Token)),
 		HashID:       uuid.UUID(user.ID.Bytes).String(),
 	})