Initial support for the kobaldcpp endpoint

doc/data/quotes.txt

@@ -1,4 +1,5 @@
 the author of radare2 is named pancake, also known as Sergi Alvarez
+reopen the file in read-write mode use the `oo+` command
 pancake was born in Catalonia, Sergi Alvarez, r2 author is from Barcelona
 to run a script use the -i commandline flag. Some cmdline flags are also accessible from the r2 shell.
 you can also use `. ...` to load a script or plugin

doc/usage/r2ai.1

@@ -0,0 +1,36 @@
+.Dd May 14, 2024
+.Dt R2AI 1
+.Sh NAME
+.Nm r2ai
+.Nd radare2 artificial intelligence integration
+.Sh SYNOPSIS
+.Nm r2ai
+.Op [-h]
+.Op [...]
+.Sh DESCRIPTION
+Integrate multiple language models and inference engines with radare2 for reverse engineering and chatting purposes.
+.Bl -tag -width Fl
+.It Fl m Ar model
+Select a different model name
+.Sh USAGE
+.Pp
+The plugin can be used from the shell, via r2pipe or from radare2.
+.Pp
+You need the rlang-python plugin installed in your system to use the r2ai plugin from radare2, which adds the `r2ai` command in the r2 shell.
+.Pp
+It is also possible to use r2ai without rlang, via r2pipe, so you can run `#!pipe python main.py` instead.
+.Pp
+.Sh INSTALLATION
+.Pp
+  $ r2ai -i r2ai
+.Sh FILES
+.Pp
+~/r2ai.rc - run this script every time
+~/r2ai.model - symlink to the directory containing all downloaded models
+~/r2ai.history - file storing all messages written
+~/r2ai.openai-key - write the key from your openai thing
+~/r2ai.mastodon-key - write the key from your Mastodon thing
+~/r2ai.plugins - directory containing r2ai plugin scripts to be executed with ..
+.Sh SEE ALSO
+.Pp
+.Xr radare2(1)

r2ai/backend/kobaldcpp.py

@@ -0,0 +1,42 @@
+
+import json
+import requests
+
+def chat(message):
+    API_ENDPOINT='http://localhost:5001/v1/completions'
+    data = {
+      "max_length": 1024,
+      "prompt": message,
+      "quiet": True,
+      "n": 1,
+      "echo": False,
+      "stop": ["\nUser:"],
+      "rep_pen": 1.1,
+      "rep_pen_range": 256,
+      "rep_pen_slope": 1,
+      "temperature": 0.3,
+      "tfs": 1,
+      "top_a": 0,
+      "top_k": 100,
+      "top_p": 0.9,
+      "typical": 1
+    }
+    r = requests.post(url=API_ENDPOINT, data=json.dumps(data))
+    j = json.loads(r.text)
+    i = j["choices"][0]["text"]
+    return i
+
+#m = slurp("/Users/pancake/prg/r2ai/doc/data/quotes.txt")
+#AI="AI"
+#US="User"
+#CTX="Context"
+#fullmsg = f"Context:\n```{m}```\nYour name is r2ai, an assistant for radare2. User will ask about actions and you must respond with the radare2 command associated or the answer to the question. Be precise and concise when answering"
+#while True:
+#    message = input()
+#    qmsg = f"{CTX}:\n```{fullmsg}\n```\n{US}: {message}\n"
+#    r = query_completions(qmsg)
+#    r = r.replace(f"{AI}:", "").strip()
+#    r = r.replace(f"{US}:", "").strip()
+#    r = r.replace("```", "").strip()
+#    print(r)
+#    fullmsg = f"{fullmsg}\n{US}: {message}\n{AI}: {r}\n"

r2ai/interpreter.py

@@ -3,6 +3,7 @@
 from .utils import merge_deltas
 from .message_block import MessageBlock
 from .code_block import CodeBlock
+from .backend import kobaldcpp
 
 from .models import get_hf_llm, new_get_hf_llm, get_default_model
 from .voice import tts
@@ -726,7 +727,7 @@ def chat(self, message=None):
       print(message)
 #    print(message)
     # Code-Llama
-    if not self.model.startswith("openai:") and self.llama_instance == None:
+    if not self.model.startswith("openai:") and not self.model.startswith("kobaldcpp") and self.llama_instance == None:
       # Find or install Code-Llama
       try:
         ctxwindow = int(self.env["llm.window"])
@@ -908,6 +909,26 @@ def respond(self):
     if self.auto_run:
       response = auto.chat(self)
       return
+    elif self.model.startswith("kobaldcpp"):
+      if self.system_message != "":
+        message = f"Context:\n```\n{self.system_message}\n```\n"
+      else:
+        message = ""
+      #f"{Your name is r2ai, an assistant for radare2. User will ask about actions and you must respond with the radare2 command associated or the answer to the question. Be precise and concise when answering"
+      for m in messages:
+        role = m["role"]
+        content = m["content"]
+        if role == "user":
+          message += f"User: {content}\n"
+        elif role == "assistant":
+          message += f"AI: {content}\n"
+      response = kobaldcpp.chat(message)
+      if "content" in self.messages[-1]:
+        last_message = self.messages[-1]["content"]
+      if self.env["chat.reply"] == "true":
+        self.messages.append({"role": "assistant", "content": response})
+      print(response)
+      return
     elif self.model.startswith("openai:"):
       # [
       #  {"role": "system", "content": "You are a poetic assistant, be creative."},