From 6c3819022ce320881a6d6e5ef719263a17991d2a Mon Sep 17 00:00:00 2001
From: Benson Wong <mostlygeek@gmail.com>
Date: Mon, 21 Oct 2024 15:38:12 -0700
Subject: [PATCH] Add compatibility with OpenAI /v1/models endpoint to list
 models

---
 proxy/manager.go | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/proxy/manager.go b/proxy/manager.go
index 17a641e..5c67507 100644
--- a/proxy/manager.go
+++ b/proxy/manager.go
@@ -31,15 +31,37 @@ func New(config *Config) *ProxyManager {
 func (pm *ProxyManager) HandleFunc(w http.ResponseWriter, r *http.Request) {
 
 	// https://github.com/ggerganov/llama.cpp/blob/master/examples/server/README.md#api-endpoints
-
 	if r.URL.Path == "/v1/chat/completions" {
 		// extracts the `model` from json body
 		pm.proxyChatRequest(w, r)
+	} else if r.URL.Path == "/v1/models" {
+		pm.listModels(w, r)
 	} else {
 		pm.proxyRequest(w, r)
 	}
 }
 
+func (pm *ProxyManager) listModels(w http.ResponseWriter, r *http.Request) {
+	data := []interface{}{}
+	for id := range pm.config.Models {
+		data = append(data, map[string]interface{}{
+			"id":       id,
+			"object":   "model",
+			"created":  time.Now().Unix(),
+			"owned_by": "llama-swap",
+		})
+	}
+
+	// Set the Content-Type header to application/json
+	w.Header().Set("Content-Type", "application/json")
+
+	// Encode the data as JSON and write it to the response writer
+	if err := json.NewEncoder(w).Encode(map[string]interface{}{"data": data}); err != nil {
+		http.Error(w, "Error encoding JSON", http.StatusInternalServerError)
+		return
+	}
+}
+
 func (pm *ProxyManager) swapModel(requestedModel string) error {
 	pm.Lock()
 	defer pm.Unlock()