Use Client

app.py

@@ -4,28 +4,30 @@ from typing import List, Tuple
 from collections import OrderedDict
 
 import gradio as gr
-from openai import OpenAI
+
+from shared import Client
 
 
 config = json.loads(os.environ['CONFIG'])
 
 
 model_names = list(config.keys())
-personas = list(OrderedDict.fromkeys(persona for name in config for persona in config[name]["personas"]))
-
-
-info = "\n".join([f"{model} ({config[model]['name']}): {list(config[model]['personas'].keys())}" for model in model_names])
 
 
 clients = {}
 for name in config:
-    client = OpenAI(
-        base_url=f"{os.environ[config[name]['api_url']]}/v1",
+    model_personas = config[name].get("personas", {})
+    client = Client(
+        api_url=os.environ[config[name]['api_url']],
         api_key=os.environ[config[name]['api_key']],
+        personas=model_personas
     )
     clients[name] = client
 
 
+personas = list(OrderedDict.fromkeys(persona for name in model_names for persona in clients[name].personas))
+info = "\n".join([f"{model} ({config[model]['name']}): {list(clients[model].personas.keys())}" for model in model_names])
+
 
 def respond(
     message,
@@ -36,12 +38,14 @@ def respond(
     conversational,
     max_tokens,
 ):
+    client = clients[model]
+
     messages = []
 
     try:
-        system_prompt = config[model]["personas"][persona]
+        system_prompt = client.personas[persona]
     except KeyError:
-        supported_personas = list(config[model]["personas"].keys())
+        supported_personas = list(client.personas.keys())
         raise gr.Error(f"Model '{model}' does not support persona '{persona}', only {supported_personas}")
     if system_prompt is not None:
         messages.append({"role": "system", "content": system_prompt})
@@ -55,8 +59,8 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    completion = clients[model].chat.completions.create(
-        model="neongeckocom/NeonLLM",
+    completion = client.client.chat.completions.create(
+        model=client.vllm_model_name,
         messages=messages,
         max_tokens=max_tokens,
         temperature=0,