Update cli - add generate models

duck_chat/__init__.py

@@ -1,3 +1,5 @@
 __version__ = "v1.3.2"
-from .api import DuckChat  # noqa
-from .models import ModelType  # noqa
+from .api import DuckChat
+from .models import ModelType
+
+__all__ = ["DuckChat", "ModelType"]

duck_chat/cli.py

@@ -1,3 +1,5 @@
+import argparse
+import asyncio
 import readline
 import sys
 import tomllib
@@ -185,6 +187,12 @@ def read_model_from_conf(self) -> ModelType:
 
 
 def safe_entry_point() -> None:
-    import asyncio
-
-    asyncio.run(CLI().run())
+    parser = argparse.ArgumentParser(description="A simple CLI tool.")
+    parser.add_argument("--generate", action="store_true", help="Generate new models")
+    args = parser.parse_args()
+    if args.generate:
+        from .models.generate_models import main as generator
+
+        generator()
+    else:
+        asyncio.run(CLI().run())

duck_chat/models/__init__.py

@@ -1,2 +1,4 @@
-from .models import Role, Message, History  # noqa
-from .model_type import ModelType  # noqa
+from .model_type import ModelType
+from .models import History, Message, Role
+
+__all__ = ["History", "ModelType", "Message", "Role"]

duck_chat/models/generate_models.py

@@ -32,7 +32,7 @@ def get_html() -> str:
     return html
 
 
-def parse_html(html) -> dict[str, str]:
+def parse_html(html: str) -> dict[str, str]:
     """Get models from html page (labels tags)"""
 
     # Parse the content of the webpage
@@ -58,11 +58,12 @@ def write_models(data: dict[str, str], path: Path) -> None:
             f.write(f'    {k} = "{v}"\n')
 
 
-def main():
+def main() -> None:
     html = get_html()
     data = parse_html(html)
-    path = Path().absolute() / "duck_chat" / "models" / "model_type.py"
+    path = Path(__file__).resolve().parent / "model_type.py"
     write_models(data, path)
+    print(f"Generate new models on {path}")
 
 
 if __name__ == "__main__":