From 0bef69b14071079ba8f073f6ac60e301f0f2db99 Mon Sep 17 00:00:00 2001
From: Daniel O'Connell <github@ahiru.pl>
Date: Fri, 29 Sep 2023 14:48:09 +0200
Subject: [PATCH] Models configured via env variables

---
 api/src/stampy_chat/chat.py       | 24 +++++++++---------------
 api/src/stampy_chat/env.py        |  4 ++++
 api/src/stampy_chat/get_blocks.py |  6 +-----
 3 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/api/src/stampy_chat/chat.py b/api/src/stampy_chat/chat.py
index a2c7cf6..59a4845 100644
--- a/api/src/stampy_chat/chat.py
+++ b/api/src/stampy_chat/chat.py
@@ -1,12 +1,14 @@
-# ------------------------------- env, constants -------------------------------
-from dataclasses import asdict
-from typing import List, Dict, Callable
-import openai
-import re
-from sqlalchemy.orm import PropComparator
-import tiktoken
 import time
+import json
+import re
+import time
+from dataclasses import asdict
+from typing import List, Dict
 
+import openai
+import tiktoken
+
+from stampy_chat.env import COMPLETIONS_MODEL
 from stampy_chat.followups import multisearch_authored
 from stampy_chat.get_blocks import get_top_k_blocks, Block
 from stampy_chat import logging
@@ -15,11 +17,6 @@ from stampy_chat import logging
 logger = logging.getLogger(__name__)
 
 
-# OpenAI models
-EMBEDDING_MODEL = "text-embedding-ada-002"
-COMPLETIONS_MODEL = "gpt-3.5-turbo"
-# COMPLETIONS_MODEL = "gpt-4"
-
 STANDARD_K = 20 if COMPLETIONS_MODEL == 'gpt-4' else 10
 
 # parameters
@@ -139,9 +136,6 @@ def construct_prompt(query: str, mode: str, history: Prompt, context: List[Block
     return prompt
 
 # ------------------------------- completion code -------------------------------
-import time
-import json
-
 
 def check_openai_moderation(prompt: Prompt, query: str):
     prompt_string = '\n\n'.join([message["content"] for message in prompt])
diff --git a/api/src/stampy_chat/env.py b/api/src/stampy_chat/env.py
index 37b7389..51e31ec 100644
--- a/api/src/stampy_chat/env.py
+++ b/api/src/stampy_chat/env.py
@@ -18,6 +18,10 @@ DISCORD_LOGGING_URL = os.environ.get('LOGGING_URL')
 OPENAI_API_KEY   = os.environ.get('OPENAI_API_KEY')
 openai.api_key = OPENAI_API_KEY # non-optional
 
+### Models ###
+EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "text-embedding-ada-002")
+COMPLETIONS_MODEL = os.environ.get("COMPLETIONS_MODEL", "gpt-3.5-turbo")
+
 ### Pinecone ###
 PINECONE_API_KEY     = os.environ.get('PINECONE_API_KEY')
 PINECONE_ENVIRONMENT = os.environ.get('PINECONE_ENVIRONMENT', "us-east1-gcp")
diff --git a/api/src/stampy_chat/get_blocks.py b/api/src/stampy_chat/get_blocks.py
index e09b71c..0908132 100644
--- a/api/src/stampy_chat/get_blocks.py
+++ b/api/src/stampy_chat/get_blocks.py
@@ -7,16 +7,12 @@ import regex as re
 import requests
 import time
 from typing import List, Tuple
-from stampy_chat.env import PINECONE_NAMESPACE, REMOTE_CHAT_INSTANCE
+from stampy_chat.env import PINECONE_NAMESPACE, REMOTE_CHAT_INSTANCE, EMBEDDING_MODEL
 from stampy_chat import logging
 
 
 logger = logging.getLogger(__name__)
 
-# ---------------------------------- constants ---------------------------------
-
-EMBEDDING_MODEL = "text-embedding-ada-002"
-
 # ------------------------------------ types -----------------------------------
 
 @dataclasses.dataclass