added support for custom openai api

docker-compose.dev.yaml

@@ -12,6 +12,8 @@ services:
       - BING_API_KEY=${BING_API_KEY}
       - SERPER_API_KEY=${SERPER_API_KEY}
       - OPENAI_API_KEY=${OPENAI_API_KEY}
+      - CUSTOM_HOST=${CUSTOM_HOST:-http://host.docker.internal:1234/v1}
+      - CUSTOM_API_KEY=${CUSTOM_API_KEY}
       - GROQ_API_KEY=${GROQ_API_KEY}
       - ENABLE_LOCAL_MODELS=${ENABLE_LOCAL_MODELS:-True}
       - SEARCH_PROVIDER=${SEARCH_PROVIDER:-tavily}

src/backend/chat.py

@@ -60,6 +60,11 @@ def get_llm(model: ChatModel) -> LLM:
         )
     elif model == ChatModel.LLAMA_3_70B:
         return Groq(model=model_mappings[model])
+    elif model == ChatModel.CUSTOM:
+        return OpenAI(model=model_mappings[model], # can use https://docs.llamaindex.ai/en/stable/api_reference/llms/openai_like/ instead
+                      api_base=os.environ.get("CUSTOM_HOST", "http://localhost:1234/v1"),
+                      api_key=os.environ.get("CUSTOM_API_KEY", "custom-key")  
+        )
     else:
         raise ValueError(f"Unknown model: {model}")
 

src/backend/constants.py

@@ -9,6 +9,7 @@
 LOCAL_GEMMA_MODEL = "gemma:7b"
 LOCAL_MISTRAL_MODEL = "mistral"
 LOCAL_PHI3_14B = "phi3:14b"
+CUSTOM_MODEL = "gpt-4" # llama_index really wants to know model context length so 8k by default.
 
 
 class ChatModel(str, Enum):
@@ -21,6 +22,7 @@ class ChatModel(str, Enum):
     LOCAL_GEMMA = "gemma"
     LOCAL_MISTRAL = "mistral"
     LOCAL_PHI3_14B = "phi3:14b"
+    CUSTOM = "custom-api"
 
 
 model_mappings: dict[ChatModel, str] = {
@@ -31,4 +33,5 @@ class ChatModel(str, Enum):
     ChatModel.LOCAL_GEMMA: LOCAL_GEMMA_MODEL,
     ChatModel.LOCAL_MISTRAL: LOCAL_MISTRAL_MODEL,
     ChatModel.LOCAL_PHI3_14B: LOCAL_PHI3_14B,
+    ChatModel.CUSTOM: CUSTOM_MODEL,
 }

src/backend/related_queries.py

@@ -36,6 +36,14 @@ def instructor_client(model: ChatModel) -> instructor.AsyncInstructor:
         )
     elif model == ChatModel.LLAMA_3_70B:
         return instructor.from_groq(groq.AsyncGroq(), mode=instructor.Mode.JSON)
+    elif model == ChatModel.CUSTOM:
+        return instructor.from_openai(
+            openai.AsyncOpenAI(
+                base_url=os.environ.get("CUSTOM_HOST", "http://localhost:1234/v1"),
+                api_key=os.environ.get("CUSTOM_API_KEY", "custom-key")
+            ),
+            mode=instructor.Mode.JSON,
+        )
     else:
         raise ValueError(f"Unknown model: {model}")
 

src/backend/validators.py

@@ -24,6 +24,10 @@ def validate_model(model: ChatModel):
         LOCAL_MODELS_ENABLED = strtobool(os.getenv("ENABLE_LOCAL_MODELS", False))
         if not LOCAL_MODELS_ENABLED:
             raise ValueError("Local models are not enabled")
+    elif model == ChatModel.CUSTOM:
+        CUSTOM_HOST = os.getenv("CUSTOM_HOST")
+        if not CUSTOM_HOST:
+            raise ValueError("CUSTOM_HOST environment variable not found")
     else:
         raise ValueError("Invalid model")
     return True

src/frontend/generated/types.gen.ts

@@ -13,6 +13,7 @@ export enum ChatModel {
   GEMMA = "gemma",
   MISTRAL = "mistral",
   LOCAL_PHI3_14B = "phi3:14b",
+  CUSTOM = "custom-api",
 }
 
 export type ChatRequest = {

src/frontend/src/components/model-selection.tsx

@@ -82,6 +82,13 @@ const modelMap: Record<ChatModel, Model> = {
     smallIcon: <FlameIcon className="w-4 h-4 text-green-500" />,
     icon: <FlameIcon className="w-5 h-5 text-green-500" />,
   },
+  [ChatModel.CUSTOM]: {
+    name: "Custom API",
+    description: "Custom API",
+    value: ChatModel.CUSTOM,
+    smallIcon: <MagicWandIcon className="w-4 h-4 text-[#FF0080]" />,
+    icon: <MagicWandIcon className="w-5 h-5 text-[#FF0080]" />,
+  },
 };
 
 const localModelMap: Partial<Record<ChatModel, Model>> = _.pickBy(