refactor: modularize assistants

pyproject.toml

@@ -1,22 +1,36 @@
+[project]
+name = 'git-draft'
+description = 'Version-controlled code assistant'
+authors = [{name='Matthieu Monsch', email='[email protected]'}]
+license = 'MIT'
+readme = 'README.md'
+dynamic = ['version']
+requires-python = '>=3.12'
+dependencies = [
+  'gitpython >=3.1.44,<4',
+]
+
+[project.optional-dependencies]
+openai = ['openai >=1.64.0,<2']
+
+[project.scripts]
+git-draft = 'git_draft.__main__:main'
+
+[project.urls]
+repository = 'https://github.com/mtth/git-draft'
+documentation = 'https://mtth.github.io/git-draft'
+
 [build-system]
 requires = ['poetry-core']
 build-backend = 'poetry.core.masonry.api'
 
+# Poetry
+
 [tool.poetry]
-name = 'git-draft'
 version = '0.0.0' # Set programmatically
-description = 'Version-controlled code assistant'
-authors = ['Matthieu Monsch <[email protected]>']
-readme = 'README.md'
-repository = 'https://github.com/mtth/git-draft'
-packages = [{include = 'git_draft', from = 'src'}]
-
-[tool.poetry.scripts]
-git-draft = 'git_draft.__main__:main'
+packages = [{include='git_draft', from='src'}]
 
 [tool.poetry.dependencies]
-gitpython = '^3.1.44'
-openai = '^1.64.0'
 python = '>=3.12,<4'
 
 [tool.poetry.group.dev.dependencies]
@@ -27,15 +41,7 @@ mypy = '^1.2.0'
 poethepoet = '^0.25.0'
 pytest = '^7.1.2'
 
-[tool.black]
-line-length = 79
-include = '\.py$'
-
-[tool.flake8]
-ignore = ['E203', 'E501', 'E704', 'W503']
-
-[tool.mypy]
-disable_error_code = 'import-untyped'
+# Poe
 
 [tool.poe.tasks.fix]
 help = 'format source code'
@@ -55,5 +61,17 @@ args = [
   {name='args', help='target folders', positional=true, multiple=true, default='src tests'},
 ]
 
+# Other tools
+
+[tool.black]
+line-length = 79
+include = '\.py$'
+
+[tool.flake8]
+ignore = ['E203', 'E501', 'E704', 'W503']
+
+[tool.mypy]
+disable_error_code = 'import-untyped'
+
 [tool.pytest.ini_options]
 log_level = 'DEBUG'

src/git_draft/__init__.py

@@ -1,11 +1,7 @@
-from .assistant import Assistant, OpenAIAssistant
-from .common import open_editor
-from .manager import Manager, enclosing_repo
+from .assistants import Assistant, Session, Toolbox
 
 __all__ = [
     "Assistant",
-    "OpenAIAssistant",
-    "Manager",
-    "enclosing_repo",
-    "open_editor",
+    "Session",
+    "Toolbox",
 ]

src/git_draft/__main__.py

@@ -5,7 +5,9 @@
 import sys
 import textwrap
 
-from . import Manager, OpenAIAssistant, enclosing_repo, open_editor
+from .assistants import load_assistant
+from .common import open_editor
+from .manager import Manager, enclosing_repo
 
 
 EPILOG = """\
@@ -39,6 +41,12 @@ def callback(_option, _opt, _value, parser) -> None:
 add_command("finalize", help="apply the current draft to the original branch")
 add_command("generate", help="draft a new change from a prompt")
 
+parser.add_option(
+    "-a",
+    "--assistant",
+    dest="ASSISTANT",
+    help="assistant key",
+)
 parser.add_option(
     "-d",
     "--delete",
@@ -65,20 +73,21 @@ def callback(_option, _opt, _value, parser) -> None:
 
 
 def main() -> None:
-    (opts, args) = parser.parse_args()
+    (opts, _args) = parser.parse_args()
 
     repo = enclosing_repo()
     manager = Manager(repo)
 
     command = getattr(opts, "command", "generate")
     if command == "generate":
+        assistant = load_assistant(opts.assistant, {})
         prompt = opts.prompt
         if not prompt:
             if sys.stdin.isatty():
                 prompt = open_editor(textwrap.dedent(EDITOR_PLACEHOLDER))
             else:
                 prompt = sys.stdin.read()
-        manager.generate_draft(prompt, OpenAIAssistant(), reset=opts.reset)
+        manager.generate_draft(prompt, assistant, reset=opts.reset)
     elif command == "finalize":
         manager.finalize_draft(delete=opts.delete)
     elif command == "discard":