From 1be355436d0f6747f6148a5ec1d4f0603ed28d63 Mon Sep 17 00:00:00 2001
From: Jim Su <jimsu@protonmail.com>
Date: Fri, 29 Mar 2024 19:34:32 -0400
Subject: [PATCH] Expose LiteLLM model names in backend (#370)

---
 frontend/.env                              |  3 ++-
 frontend/README.md                         |  2 +-
 frontend/src/components/BannerSettings.tsx | 16 ++++++++++++----
 frontend/src/services/settingsService.ts   | 10 ++++++++--
 frontend/src/socket/socket.ts              |  4 ++--
 opendevin/server/listen.py                 | 16 ++++++++++++++++
 6 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/frontend/.env b/frontend/.env
index 0770f1bf02..8fecc59f85 100644
--- a/frontend/.env
+++ b/frontend/.env
@@ -1 +1,2 @@
-VITE_TERMINAL_WS_URL="ws://localhost:3000/ws"
\ No newline at end of file
+VITE_URL="http://localhost:3000"
+VITE_WS_URL="ws://localhost:3000/ws"
\ No newline at end of file
diff --git a/frontend/README.md b/frontend/README.md
index 8b4b82ff2e..ca2be99516 100644
--- a/frontend/README.md
+++ b/frontend/README.md
@@ -30,6 +30,6 @@ To learn React, check out the [React documentation](https://reactjs.org/).
 
 The OpenDevin terminal is powered by [Xterm.js](https://github.com/xtermjs/xterm.js).
 
-The terminal listens for events over a WebSocket connection. The WebSocket URL is specified by the environment variable `VITE_TERMINAL_WS_URL` (prepending `VITE_` to environment variable names is necessary to expose them).
+The terminal listens for events over a WebSocket connection. The WebSocket URL is specified by the environment variable `VITE_WS_URL` (prepending `VITE_` to environment variable names is necessary to expose them).
 
 A simple websocket server can be found in the `/server` directory.
diff --git a/frontend/src/components/BannerSettings.tsx b/frontend/src/components/BannerSettings.tsx
index 6a254178e7..ed62937e9e 100644
--- a/frontend/src/components/BannerSettings.tsx
+++ b/frontend/src/components/BannerSettings.tsx
@@ -1,13 +1,21 @@
-import React, { ChangeEvent } from "react";
+import React, { ChangeEvent, useEffect, useState } from "react";
 import {
   AGENTS,
-  MODELS,
+  INITIAL_MODELS,
   changeAgent,
   changeModel,
+  fetchModels,
 } from "../services/settingsService";
 import "./BannerSettings.css";
 
 function ModelSelect(): JSX.Element {
+  const [models, setModels] = useState<string[]>(INITIAL_MODELS);
+  useEffect(() => {
+    fetchModels().then((fetchedModels) => {
+      setModels(fetchedModels);
+    });
+  }, []);
+
   return (
     <select
       onChange={(e: ChangeEvent<HTMLSelectElement>) =>
@@ -15,8 +23,8 @@ function ModelSelect(): JSX.Element {
       }
       className="select max-w-xs bg-base-300 xl:w-full w-1/3"
     >
-      {MODELS.map((model) => (
-        <option>{model}</option>
+      {models.map((model) => (
+        <option key={model}>{model}</option>
       ))}
     </select>
   );
diff --git a/frontend/src/services/settingsService.ts b/frontend/src/services/settingsService.ts
index 0bf69ecd2e..db540ffe88 100644
--- a/frontend/src/services/settingsService.ts
+++ b/frontend/src/services/settingsService.ts
@@ -3,7 +3,13 @@ import { appendAssistantMessage } from "../state/chatSlice";
 import { setInitialized } from "../state/taskSlice";
 import store from "../store";
 
-export const MODELS = [
+const { VITE_URL } = import.meta.env;
+export async function fetchModels() {
+  const response = await fetch(`${VITE_URL}/litellm-models`);
+  return response.json();
+}
+
+export const INITIAL_MODELS = [
   "gpt-3.5-turbo-1106",
   "gpt-4-0125-preview",
   "claude-3-haiku-20240307",
@@ -11,7 +17,7 @@ export const MODELS = [
   "claude-3-sonnet-20240229",
 ];
 
-export type Model = (typeof MODELS)[number];
+export type Model = (typeof INITIAL_MODELS)[number];
 
 export const AGENTS = ["LangchainsAgent", "CodeActAgent"];
 
diff --git a/frontend/src/socket/socket.ts b/frontend/src/socket/socket.ts
index 27f5012691..748f28566d 100644
--- a/frontend/src/socket/socket.ts
+++ b/frontend/src/socket/socket.ts
@@ -6,10 +6,10 @@ import { handleObservationMessage } from "./observations";
 
 type SocketMessage = ActionMessage | ObservationMessage;
 
-const WS_URL = import.meta.env.VITE_TERMINAL_WS_URL;
+const WS_URL = import.meta.env.VITE_WS_URL;
 if (!WS_URL) {
   throw new Error(
-    "The environment variable VITE_TERMINAL_WS_URL is not set. Please set it to the WebSocket URL of the terminal server.",
+    "The environment variable VITE_WS_URL is not set. Please set it to the WebSocket URL of the terminal server.",
   );
 }
 
diff --git a/opendevin/server/listen.py b/opendevin/server/listen.py
index 9031a326ff..96dfa7a30b 100644
--- a/opendevin/server/listen.py
+++ b/opendevin/server/listen.py
@@ -1,9 +1,19 @@
 from opendevin.server.session import Session
 from fastapi import FastAPI, WebSocket
+from fastapi.middleware.cors import CORSMiddleware
 import agenthub # noqa F401 (we import this to get the agents registered)
+import litellm
 
 app = FastAPI()
 
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:3001"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 # This endpoint receives events from the client (i.e. the browser)
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
@@ -12,3 +22,9 @@ async def websocket_endpoint(websocket: WebSocket):
     # TODO: should this use asyncio instead of await?
     await session.start_listening()
 
+@app.get("/litellm-models")
+async def get_litellm_models():
+    """
+    Get all models supported by LiteLLM.
+    """
+    return litellm.model_list