Packaging redo

.github/workflows/ci.yml

@@ -28,20 +28,20 @@ jobs:
     - name: Install Python dependencies
       run: |
         . venv/bin/activate
-        pip install -r requirements.txt
+        pip install .
     - name: Linting
       run: |
         . venv/bin/activate
         make cilint
     - name: Running r2ai from the shell
       run: |
         . venv/bin/activate
-        python main.py -h
+        python -m r2ai.cli -h
     - name: Running r2ai from r2pipe
       run: |
         . venv/bin/activate
-        r2 -q -c '#!pipe python main.py -h' /bin/ls
+        r2 -q -c '#!pipe python -m r2ai.cli -h' /bin/ls
     - name: Running r2ai from r2
       run: |
         . venv/bin/activate
-        r2 -q -i main.py -c r2ai /bin/ls
+        r2 -q -i r2ai/plugin.py -c r2ai /bin/ls

Makefile

@@ -12,7 +12,6 @@ PIP=$(PYTHON) -m pip
 LINTED=r2ai/code_block.py
 LINTED+=r2ai/bubble.py
 LINTED+=r2ai/const.py
-LINTED+=main.py
 LINTED+=r2ai/backend/kobaldcpp.py
 # LINTED+=r2ai/index.py
 # LINTED+=r2ai/voice.py

main.py

@@ -1,69 +1 @@
-"""Entrypoint for the r2ai plugin and repl."""
-# pylint: disable=import-outside-toplevel
-# pylint: disable=unused-import
-# pylint: disable=missing-function-docstring
-
-import sys
-import os
-import argparse
-
-script_dir = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(script_dir)
-pv = f"{sys.version_info.major}.{sys.version_info.minor}" # pylint: disable=invalid-name
-sys.path.append(script_dir + f"/venv/lib/python{pv}/site-packages/")
-
-from r2ai.main import main as r2ai_main # pylint: disable=wrong-import-position
-from r2ai.main import register_r2plugin # pylint: disable=wrong-import-position
-
-def is_valid_file(parser, arg):
-    if not os.path.isfile(arg):
-        parser.error(f"The file {arg} does not exist!")
-
-def massage_args(args):
-    runrepl = True
-    if args.command is None:
-        args.command = []
-    if args.webserver:
-        args.command.append("-w")
-    if args.eval:
-        if args.eval == "default":
-            args.command.append("-e")
-            runrepl = False
-        else:
-            args.command.append(f"-e {args.eval}")
-    if args.port:
-        if args.port == "default":
-            runrepl = False
-            args.command.append("-e http.port")
-        else:
-            args.command.append(f"-e http.port={args.port}")
-    if args.model:
-        if args.model == "default":
-            args.command.append("-mm")
-            runrepl = False
-        else:
-            args.command.append(f"-m {args.model}")
-    return runrepl, args
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("bin", nargs="?", type=str)
-    parser.add_argument("-w", "--webserver", action="store_true",
-        help="Start the r2ai webserver. Same as r2ai -c=-w")
-    parser.add_argument("-p", "--port", type=str, nargs="?", const="default",
-        help="Change listen port number")
-    parser.add_argument("-e", "--eval", type=str, nargs="?", const="default",
-        help="Change configuration variable")
-    parser.add_argument("-m", "--model", type=str, nargs="?", const="default",
-        help="Select model name")
-    parser.add_argument("-c", "--command", action="append",
-        help="Command to be executed. Can be passed multiple times.")
-    runrepl, args = massage_args(parser.parse_args())
-    r2ai_main(args, args.command, runrepl)
-
-if __name__ == "__main__":
-    try:
-        import r2lang # pylint: disable=import-error
-        register_r2plugin()
-    except ImportError:
-        main()
+import r2ai.cli

pyproject.toml

@@ -9,43 +9,51 @@ authors = [
   {name = "pancake", email = "pancake@nopcode.org"}
 ]
 
-[project.urls]
-homepage = "https://www.radare.org/"
-repository = "https://github.com/radareorg/r2ai"
-
-[project.scripts]
-r2ai = "r2ai:main"
-
-[tool.setuptools.packages.find]
-where = [".", "r2ai"]
-include = ["main", "r2ai", "r2ai/ui"]
-namespaces = true
-
-[tool.setuptools.dynamic]
-readme = {file = "README.md", content-type = "text/markdown"}
-
-[build-system]
-requires = ["setuptools"]
-build-backend = "setuptools.build_meta"
-
 dependencies = [
   "rich",
   "r2pipe",
   "inquirer",
   "llama-cpp-python",
-  "huggingface_hub==0.22.0",
+  "huggingface_hub",
   "appdirs",
   "unidecode",
   "jsonref",
   "transformers",
-  "pydantic"
+  "pydantic",
+  "pyreadline3",
+  "tokentrim",
+  "boto3",
+  "colorama",
+  "textual",
+  "litellm"
 ]
 
 [project.optional-dependencies]
 extras = [
+  "chromadb",
   "openai",
   "anthropic",
   "groq",
   "google-generativeai",
   "google-cloud-aiplatform"
 ]
+
+
+[project.urls]
+homepage = "https://www.radare.org/"
+repository = "https://github.com/radareorg/r2ai"
+
+[project.scripts]
+r2ai = "r2ai.main:run" 
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["r2ai*"]
+namespaces = true
+
+[tool.setuptools.dynamic]
+readme = {file = "README.md", content-type = "text/markdown"}
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"

r2ai.r2

@@ -1,3 +1,3 @@
 ?ef [r2ai] Loaded via r2pipe accessible through "$r2ai"
-'$r2ai=#!pipe python main.py
+'$r2ai=#!python r2ai/plugin.py
 # $r2ai -h

r2ai.sh

@@ -13,10 +13,9 @@ RD=`realpath "$D"`
 [ -n "$D" ] && cd "$D"
 if [ ! -d venv ]; then
 	$PYTHON -m venv venv
-	. venv/bin/activate
-	pip3 install -r requirements.txt
+	./venv/bin/pip3 install -e .
 else
-	. venv/bin/activate
+	PYTHON=venv/bin/python3
 fi
-# export PYTHONPATH=$PWD
-exec $PYTHON $D/main.py "$@"
+
+exec $PYTHON -m r2ai.cli "$@"

r2ai/cli.py

@@ -0,0 +1,3 @@
+from .main import run # pylint: disable=wrong-import-position
+
+run()

r2ai/main.py

@@ -5,22 +5,19 @@
 import builtins
 import traceback
 import appdirs
+import argparse
 
 from r2ai.repl import r2ai_singleton
 from r2ai.utils import slurp
 from r2ai.repl import runline, r2ai_repl, help_message
 
-from r2ai.pipe import open_r2
-from .const import R2AI_RCFILE
+from r2ai.pipe import open_r2, get_r2_inst
+from r2ai.const import R2AI_RCFILE
 
 OPENAI_KEY = ""
 HAVE_RLANG = False
 HAVE_R2PIPE = False
 RCFILE_LOADED = False
-within_r2 = False
-
-if "R2CORE" in os.environ:
-    within_r2 = True
 
 def r2ai_rlang_plugin(unused_but_required_argument):
     ai = r2ai_singleton()
@@ -68,12 +65,7 @@ def run_rcfile_once(ai):
         RCFILE_LOADED = True
 
 
-def register_r2plugin():
-    import r2lang
-    r2lang.plugin("core", r2ai_rlang_plugin)
-
 def main(args, commands, dorepl=True):
-    global within_r2
 
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -113,8 +105,7 @@ def main(args, commands, dorepl=True):
 
     ai = r2ai_singleton()
     if "R2PIPE_IN" in os.environ:
-        open_r2(None)
-        within_r2 = True
+        pass
     elif args.bin:
         open_r2(vars(args)["bin"], flags=["-2"])
 
@@ -126,33 +117,45 @@ def main(args, commands, dorepl=True):
                 runline(ai, c)
     if dorepl:
         r2ai_repl(ai)
-    # elif HAVE_RLANG and HAVE_R2PIPE:
-    #     r2ai_repl(ai)
-    #     os.exit(0)
 
-    #     r2lang.plugin("core", r2ai_rlang_plugin)
-
-    # else:
-    #     if "R2CORE" in os.environ:
-    #         print("[R2AI] Please: r2pm -ci rlang-python")
-    #         sys.exit(0)
-
-    #     run_rcfile(ai)
-    #     if len(sys.argv) > 1:
-    #         for arg in sys.argv[1:]:
-    #             if arg.endswith(".py"):
-    #                 exec(slurp(arg), globals())
-    #                 sys.stderr.close()
-    #                 sys.exit(0)
-    #             elif not arg.startswith("/"):
-    #                 runline(ai, arg)
-    #             if arg == "-h" or arg == "-v":
-    #                 sys.exit(0)
-    #             elif arg == "-repl":
-    #                 r2ai_repl(ai)
-    #     elif not within_r2:
-    #         r2ai_repl(ai)
-    #     elif HAVE_R2PIPE:
-    #         r2ai_repl(ai)
-    #     else:
-    #         print("r2ai plugin cannot be loaded. Run `r2pm -ci rlang-python`")
+def massage_args(args):
+    runrepl = True
+    if args.command is None:
+        args.command = []
+    if args.webserver:
+        args.command.append("-w")
+    if args.eval:
+        if args.eval == "default":
+            args.command.append("-e")
+            runrepl = False
+        else:
+            args.command.append(f"-e {args.eval}")
+    if args.port:
+        if args.port == "default":
+            runrepl = False
+            args.command.append("-e http.port")
+        else:
+            args.command.append(f"-e http.port={args.port}")
+    if args.model:
+        if args.model == "default":
+            args.command.append("-mm")
+            runrepl = False
+        else:
+            args.command.append(f"-m {args.model}")
+    return runrepl, args
+
+def run():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("bin", nargs="?", type=str)
+    parser.add_argument("-w", "--webserver", action="store_true",
+        help="Start the r2ai webserver. Same as r2ai -c=-w")
+    parser.add_argument("-p", "--port", type=str, nargs="?", const="default",
+        help="Change listen port number")
+    parser.add_argument("-e", "--eval", type=str, nargs="?", const="default",
+        help="Change configuration variable")
+    parser.add_argument("-m", "--model", type=str, nargs="?", const="default",
+        help="Select model name")
+    parser.add_argument("-c", "--command", action="append",
+        help="Command to be executed. Can be passed multiple times.")
+    runrepl, args = massage_args(parser.parse_args())
+    main(args, args.command, runrepl)

r2ai/pipe.py

@@ -34,7 +34,7 @@ def cmd(self, x):
 try:
     import r2lang
     have_rlang = True
-except Exception:
+except Exception as e:
     import r2pipe
     try:
         if r2pipe.in_r2():
@@ -56,12 +56,11 @@ def cmd(self, x):
             r2lang = FakeLang(None)
 
 def r2singleton():
-    global r2lang
-    return r2lang
+    global r2lang, r2
+    return r2lang or r2
 
 def get_r2_inst():
-    global r2
-    return r2
+    return r2singleton()
 
 @progress_bar("Loading", color="yellow")
 def open_r2(file, flags=[]):

r2ai/plugin.py

@@ -0,0 +1,32 @@
+"""Entrypoint for the r2ai plugin and repl."""
+
+import builtins
+import traceback
+
+import r2lang
+def r2ai_rlang_plugin(unused_but_required_argument):
+    from r2ai.main import run, r2ai_singleton, run_rcfile_once, runline, help_message # pylint: disable=wrong-import-position
+    ai = r2ai_singleton()
+    def _call(s):
+        if not s.startswith("r2ai"):
+            return False
+        try:
+            run_rcfile_once(ai)
+            if len(s) == 4:
+                builtins.print(help_message)
+            else:
+                usertext = s[4:].strip()
+                runline(ai, usertext)
+        except Exception as e:
+            builtins.print(e)
+            traceback.print_exc()
+        return True
+
+    return {
+        "name": "r2ai",
+        "license": "MIT",
+        "desc": "run llama language models inside r2",
+        "call": _call,
+    }
+
+r2lang.plugin("core", r2ai_rlang_plugin)