Support alternative endpoints for kobaldcpp

radareorg · May 14, 2024 · a9fecdd · a9fecdd
1 parent 1d96a4f
commit a9fecdd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git 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
@@ -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
@@ -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