From a9fecdd05dea994270f35a53fc568d4063377ba7 Mon Sep 17 00:00:00 2001
From: pancake <pancake@nopcode.org>
Date: Tue, 14 May 2024 19:23:51 +0200
Subject: [PATCH] Support alternative endpoints for kobaldcpp

---
 r2ai/backend/kobaldcpp.py | 4 ++--
 r2ai/interpreter.py       | 6 +++++-
 r2ai/models.py            | 3 +++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/r2ai/backend/kobaldcpp.py b/r2ai/backend/kobaldcpp.py
index 47965b3c..cbe3ac75 100644
--- a/r2ai/backend/kobaldcpp.py
+++ b/r2ai/backend/kobaldcpp.py
@@ -2,8 +2,8 @@
 import json
 import requests
 
-def chat(message):
-    API_ENDPOINT='http://localhost:5001/v1/completions'
+def chat(message, API_ENDPOINT='http://localhost:5001'):
+    API_ENDPOINT+='/v1/completions'
     data = {
       "max_length": 1024,
       "prompt": message,
diff --git a/r2ai/interpreter.py b/r2ai/interpreter.py
index fa41f726..ba67d75b 100644
--- a/r2ai/interpreter.py
+++ b/r2ai/interpreter.py
@@ -922,7 +922,11 @@ def respond(self):
           message += f"User: {content}\n"
         elif role == "assistant":
           message += f"AI: {content}\n"
-      response = kobaldcpp.chat(message)
+      response = ""
+      if ":" in self.model:
+        response = kobaldcpp.chat(message, self.model.split(":")[1:])
+      else:
+        response = kobaldcpp.chat(message)
       if "content" in self.messages[-1]:
         last_message = self.messages[-1]["content"]
       if self.env["chat.reply"] == "true":
diff --git a/r2ai/models.py b/r2ai/models.py
index 6523e4db..375c5898 100644
--- a/r2ai/models.py
+++ b/r2ai/models.py
@@ -47,6 +47,9 @@ def mainmodels():
 
 def models():
     return """
+KobaldCpp:
+-m kobaldcpp
+-m kobaldcpp:http://localhost:5001
 OpenAI:
 -m openai:gpt-3.5-turbo
 -m openai:gpt-4