use poetry

ollama.py → ollama/__main__.py

-from ollama.cmd.cli import main
+from ollama.cmd import cli
+
 if __name__ == '__main__':
-    main()
+    cli.main()

ollama/cmd/cli.py

+import os
 import json
 from pathlib import Path
 from argparse import ArgumentParser
@@ -15,7 +16,7 @@ def main():
     server.set_parser(subparsers.add_parser("serve"))
 
     list_parser = subparsers.add_parser("list")
-    list_parser.set_defaults(fn=list)
+    list_parser.set_defaults(fn=list_models)
 
     generate_parser = subparsers.add_parser("generate")
     generate_parser.add_argument("model")
@@ -23,17 +24,22 @@ def main():
     generate_parser.set_defaults(fn=generate)
 
     add_parser = subparsers.add_parser("add")
-    add_parser.add_argument("file")
+    add_parser.add_argument("model")
     add_parser.set_defaults(fn=add)
 
     args = parser.parse_args()
     args = vars(args)
 
-    fn = args.pop("fn")
-    fn(**args)
+    try:
+        fn = args.pop("fn")
+        fn(**args)
+    except KeyError:
+        parser.print_help()
+    except Exception as e:
+        print(e)
 
 
-def list(*args, **kwargs):
+def list_models(*args, **kwargs):
     for m in model.models(*args, **kwargs):
         print(m)
 
@@ -47,5 +53,5 @@ def generate(*args, **kwargs):
             print(choices[0].get("text", ""), end="")
 
 
-def add(*args, **kwargs):
-    engine.add(*args, **kwargs)
+def add(model, models_home):
+    os.rename(model, Path(models_home) / Path(model).name)

ollama/engine.py

 import os
 import json
 import sys
-import shutil
 from contextlib import contextmanager
 from llama_cpp import Llama as LLM
 from template import template
@@ -62,9 +61,3 @@ def load(model, models_home=".", llms={}):
 def unload(model, llms={}):
     if model in llms:
         llms.pop(model)
-
-
-def add(file, models_home=".", *args, **kwargs):
-    if not os.path.exists(file):
-        raise ValueError("Model file {model} not found")
-    shutil.move(file, models_home)

pyproject.toml

+[tool.poetry]
+name = "ollama"
+version = "0.0.1"
+description = ""
+authors = ["Your Name <you@example.com>"]
+readme = "README.md"
+packages = [{include = "ollama"}]
+scripts = {ollama = "ollama.cmd.cli:main"}
+
+[tool.poetry.dependencies]
+python = "^3.11"
+llama-cpp-python = "^0.1.66"
+
+aiohttp = {version = "^3.8.4", optional = true}
+aiohttp-cors = {version = "^0.7.0", optional = true}
+
+[tool.poetry.extras]
+server = ["aiohttp", "aiohttp_cors"]
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+flake8 = "^6.0.0"
+
+[tool.poetry.group.lsp]
+optional = true
+
+[tool.poetry.group.lsp.dependencies]
+pyright = "^1.1.316"
+python-lsp-server = "^1.7.3"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"