From 88f2a95808f443327db908710826102eba858ef3 Mon Sep 17 00:00:00 2001
From: daniel nakov <daniel.nakov@gmail.com>
Date: Fri, 7 Jun 2024 14:26:30 -0400
Subject: [PATCH] Package updates

---
 .gitignore       |  2 ++
 main.py          | 35 ++++++++++++++++++--------------
 pyproject.toml   | 52 ++++++++++++++++++++++++++++++++++++++++++++++++
 r2ai/__init__.py |  5 -----
 setup.py         | 39 ------------------------------------
 5 files changed, 74 insertions(+), 59 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 pyproject.toml
 delete mode 100644 setup.py

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 00000000..7ab273a9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+build
+venv
\ No newline at end of file
diff --git a/main.py b/main.py
index a57c3ba6..a891275e 100644
--- a/main.py
+++ b/main.py
@@ -4,18 +4,23 @@
 import sys
 import os
 
-if "R2AI" in os.environ:
-    print("Cant load r2ai r2 plugin from inside r2ai")
-    sys.exit(0)
-r2aihome = os.path.dirname(__file__)
-try:
-    r2aihome = os.path.dirname(os.readlink(__file__))
-except (OSError, FileNotFoundError):
-    pass
-sys.path.insert(0, r2aihome)
-os.environ["R2AI"] = "1"
-if "VIRTUAL_ENV" in os.environ or "R2CORE" in os.environ:
-    runpy.run_path(os.path.join(r2aihome, 'r2ai', 'main.py'))
-else:
-    ARGS = " ".join(sys.argv[1:])
-    os.system(f"cd {r2aihome}; ./r2ai.sh {ARGS}")
+def main():
+  print('os', os.environ.get('R2AI'))
+  if os.environ.get('R2AI'):
+      print("Cant load r2ai r2 plugin from inside r2ai")
+      sys.exit(0)
+  r2aihome = os.path.dirname(__file__)
+  try:
+      r2aihome = os.path.dirname(os.readlink(__file__))
+  except (OSError, FileNotFoundError):
+      pass
+  sys.path.insert(0, r2aihome)
+  os.environ["R2AI"] = "1"
+  if "VIRTUAL_ENV" in os.environ or "R2CORE" in os.environ:
+      runpy.run_path(os.path.join(r2aihome, 'r2ai', 'main.py'))
+  else:
+      ARGS = " ".join(sys.argv[1:])
+      os.system(f"cd {r2aihome}; ./r2ai.sh {ARGS}")
+
+if(__name__ == "__main__"):
+    main()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..075717a9
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,52 @@
+[project]
+name = "r2ai"
+dynamic = ["readme"]
+version = "0.7.0"
+description = "Applying language models on radare2 for reverse engineering and fun purposes"
+
+license = {text = "MIT License"}
+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"]
+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",
+  "appdirs",
+  "unidecode",
+  "jsonref",
+  "transformers",
+  "pydantic",
+  "torch"
+]
+
+[project.optional-dependencies]
+extras = [
+  "openai",
+  "anthropic",
+  "groq",
+  "google-generativeai",
+  "google-cloud-aiplatform"
+]
diff --git a/r2ai/__init__.py b/r2ai/__init__.py
index 00b7a2ae..e69de29b 100644
--- a/r2ai/__init__.py
+++ b/r2ai/__init__.py
@@ -1,5 +0,0 @@
-import sys
-from .models import models, mainmodels
-sys.modules["r2ai"].models = models
-sys.modules["r2ai"].mainmodels = mainmodels
-VERSION = "0.7.0"
diff --git a/setup.py b/setup.py
deleted file mode 100644
index ba68709f..00000000
--- a/setup.py
+++ /dev/null
@@ -1,39 +0,0 @@
-"""The classic setuptools python file."""
-
-from setuptools import setup
-import r2ai
-
-with open("README.md", encoding = "utf-8") as fd:
-    readme = fd.read()
-
-setup(
-    name ='r2ai',
-    version = r2ai.VERSION,
-    description = "Applying language models on radare2 for reverse engineering and fun purposes",
-    long_description = readme,
-    long_description_content_type = "text/markdown",
-    author = "pancake",
-    author_email = "pancake@nopcode.org",
-    url = "https://www.radare.org/",
-    packages = [
-        'r2ai',
-    ],
-    install_requires = [
-        'rich',
-        'r2pipe',
-        'inquirer',
-        'llama-cpp-python',
-        'huggingface_hub',
-        'appdirs',
-        'unidecode',
-        'jsonref',
-        'transformers',
-        'pydantic',
-        'torch',
-    ],
-    entry_points = {
-        'console_scripts': [
-            'r2ai = r2ai.main:main'
-        ]
-    }
-)