From 3e90f8328d028d58426792961b56ad83f289d511 Mon Sep 17 00:00:00 2001
From: Benson Wong <mostlygeek@gmail.com>
Date: Fri, 4 Oct 2024 12:28:50 -0700
Subject: [PATCH] add /v1/models endpoint and proxy everything to llama-server

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

diff --git a/proxy/manager.go b/proxy/manager.go
index b33dd6f..4611069 100644
--- a/proxy/manager.go
+++ b/proxy/manager.go
@@ -28,10 +28,41 @@ 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" {
+
+		// Transform Models to the desired JSON structure
+		jsonOutput := map[string]interface{}{
+			"object": "list",
+			"data":   []map[string]string{},
+		}
+
+		for id := range pm.config.Models {
+			modelData := map[string]string{
+				"id":     id,
+				"object": "model",
+			}
+			jsonOutput["data"] = append(jsonOutput["data"].([]map[string]string), modelData)
+		}
+
+		// Marshal the JSON output for display
+		result, err := json.MarshalIndent(jsonOutput, "", "  ")
+		if err != nil {
+			http.Error(w, "JSON marshal error", http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(result)
+
 	} else {
-		http.Error(w, "endpoint not supported", http.StatusNotFound)
+
+		pm.proxyRequest(w, r)
 	}
 }