zrb 1.2.2__py3-none-any.whl → 1.3.1__py3-none-any.whl

zrb/builtin/llm/llm_chat.py

@@ -2,6 +2,8 @@ import json
 import os
 from typing import Any
 
+from pydantic_ai.models import Model
+
 from zrb.builtin.group import llm_group
 from zrb.builtin.llm.tool.api import get_current_location, get_current_weather
 from zrb.builtin.llm.tool.cli import run_shell_command
@@ -26,11 +28,14 @@ from zrb.config import (
     LLM_SYSTEM_PROMPT,
     SERP_API_KEY,
 )
+from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
+from zrb.input.any_input import AnyInput
 from zrb.input.bool_input import BoolInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
 from zrb.task.llm_task import LLMTask
+from zrb.util.attr import get_attr
 from zrb.util.file import read_file, write_file
 from zrb.util.string.conversion import to_pascal_case
 
@@ -85,23 +90,53 @@ def _write_chat_conversation(
     write_file(last_session_file_path, current_session_name)
 
 
-llm_chat: LLMTask = llm_group.add_task(
-    LLMTask(
-        name="llm-chat",
-        input=[
+class _LLMChat(LLMTask):
+
+    _default_model: Model | str | None = None
+
+    def set_default_model(self, model: Model | str):
+        self._default_model = model
+
+    @property
+    def inputs(self) -> list[AnyInput]:
+        task_inputs = super().inputs
+        model_input_default = LLM_MODEL if self._default_model is None else "default"
+        return [
             StrInput(
                 "model",
                 description="LLM Model",
                 prompt="LLM Model",
-                default=LLM_MODEL,
+                default=model_input_default,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
+            *task_inputs,
+        ]
+
+    def _get_model(self, ctx: AnyContext) -> str | Model | None:
+        if ctx.input.model == "default":
+            if self._default_model is not None:
+                return self._default_model
+            return super()._get_model(ctx)
+        model = get_attr(
+            ctx, ctx.input.model, "ollama_chat/llama3.1", auto_render=self._render_model
+        )
+        if isinstance(model, (Model, str)) or model is None:
+            return model
+        raise ValueError("Invalid model")
+
+
+llm_chat: LLMTask = llm_group.add_task(
+    _LLMChat(
+        name="llm-chat",
+        input=[
             TextInput(
                 "system-prompt",
                 description="System prompt",
                 prompt="System prompt",
                 default=LLM_SYSTEM_PROMPT,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
             BoolInput(
                 "start-new",
@@ -109,6 +144,7 @@ llm_chat: LLMTask = llm_group.add_task(
                 prompt="Start new conversation (LLM will forget everything)",
                 default=False,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
             TextInput("message", description="User message", prompt="Your message"),
             PreviousSessionInput(
@@ -117,12 +153,12 @@ llm_chat: LLMTask = llm_group.add_task(
                 prompt="Previous conversation session (can be empty)",
                 allow_positional_parsing=False,
                 allow_empty=True,
+                always_prompt=False,
             ),
         ],
         conversation_history_reader=_read_chat_conversation,
         conversation_history_writer=_write_chat_conversation,
         description="Chat with LLM",
-        model="{ctx.input.model}",
         system_prompt="{ctx.input['system-prompt']}",
         message="{ctx.input.message}",
         retries=0,

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_task.py

@@ -8,7 +8,7 @@ from my_app_name._zrb.column.add_column_util import (
     update_my_app_name_test_update,
     update_my_app_name_ui,
 )
-from my_app_name._zrb.config import APP_DIR
+from my_app_name._zrb.config import ACTIVATE_VENV_SCRIPT, APP_DIR
 from my_app_name._zrb.format_task import format_my_app_name_code
 from my_app_name._zrb.group import app_create_group
 from my_app_name._zrb.input import (
@@ -17,9 +17,16 @@ from my_app_name._zrb.input import (
     new_column_input,
     new_column_type_input,
 )
-from my_app_name._zrb.util import get_existing_module_names, get_existing_schema_names
+from my_app_name._zrb.util import (
+    cd_module_script,
+    get_existing_module_names,
+    get_existing_schema_names,
+    set_create_migration_db_url_env,
+    set_env,
+)
+from my_app_name._zrb.venv_task import prepare_venv
 
-from zrb import AnyContext, Task, make_task
+from zrb import AnyContext, Cmd, CmdTask, EnvFile, Task, make_task
 
 
 @make_task(
@@ -118,6 +125,34 @@ update_my_app_name_test_delete_task = Task(
 )
 
 
+create_my_app_name_entity_migration = CmdTask(
+    name="create-my-app-name-entity-migration",
+    input=[
+        existing_module_input,
+        existing_entity_input,
+        new_column_input,
+    ],
+    env=EnvFile(path=os.path.join(APP_DIR, "template.env")),
+    cwd=APP_DIR,
+    cmd=[
+        ACTIVATE_VENV_SCRIPT,
+        Cmd(lambda ctx: set_create_migration_db_url_env(ctx.input.module)),
+        Cmd(lambda ctx: set_env("MY_APP_NAME_MODULES", ctx.input.module)),
+        Cmd(lambda ctx: cd_module_script(ctx.input.module)),
+        "alembic upgrade head",
+        Cmd(
+            'alembic revision --autogenerate -m "create_{to_snake_case(ctx.input.entity)}_{to_snake_case(ctx.input.column)}_column"',  # noqa
+        ),
+    ],
+    render_cmd=False,
+    retries=0,
+    upstream=[
+        prepare_venv,
+        update_my_app_name_schema_task,
+    ],
+)
+
+
 add_my_app_name_column = app_create_group.add_task(
     Task(
         name="add-my-app-name-column",
@@ -129,6 +164,7 @@ add_my_app_name_column = app_create_group.add_task(
             update_my_app_name_test_read_task,
             update_my_app_name_test_update_task,
             update_my_app_name_test_delete_task,
+            create_my_app_name_entity_migration,
         ],
         successor=format_my_app_name_code,
         retries=0,

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_util.py

@@ -2,7 +2,7 @@ import os
 import re
 import textwrap
 
-from bs4 import BeautifulSoup, formatter
+from bs4 import BeautifulSoup, Tag, formatter
 from my_app_name._zrb.config import APP_DIR
 
 from zrb.context.any_context import AnyContext
@@ -65,6 +65,7 @@ def update_my_app_name_ui(ctx: AnyContext):
     kebab_entity_name = to_kebab_case(ctx.input.entity)
     snake_column_name = to_snake_case(ctx.input.column)
     human_column_name = to_human_case(ctx.input.column).title()
+    column_type = ctx.input.type
     subroute_file_path = os.path.join(
         APP_DIR,
         "module",
@@ -85,18 +86,21 @@ def update_my_app_name_ui(ctx: AnyContext):
         form_id="crud-create-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     new_code = _add_input_to_form(
         new_code,
         form_id="crud-update-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     new_code = _add_input_to_form(
         new_code,
         form_id="crud-delete-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     # JS Function
     new_code = _alter_js_function_returned_array(
@@ -140,7 +144,7 @@ def _add_th_before_last(html_str, table_id, th_content):
 
 
 def _add_input_to_form(
-    html_str: str, form_id: str, column_label: str, column_name: str
+    html_str: str, form_id: str, column_label: str, column_name: str, column_type: str
 ) -> str:
     soup = BeautifulSoup(html_str, "html.parser")
     # Find the form by id.
@@ -151,9 +155,7 @@ def _add_input_to_form(
     new_label = soup.new_tag("label")
     new_label.append(f"{column_label}: ")
     # Create a new input element with the provided column name.
-    new_input = soup.new_tag(
-        "input", attrs={"type": "text", "name": column_name, "required": "required"}
-    )
+    new_input = _get_html_input(soup, column_name, column_type)
     new_label.append(new_input)
     # Look for a footer element inside the form.
     footer = form.find("footer")
@@ -168,6 +170,26 @@ def _add_input_to_form(
     )
 
 
+def _get_html_input(soup: BeautifulSoup, column_name: str, column_type: str) -> Tag:
+    # Map your custom types to HTML input types.
+    type_mapping = {
+        "str": "text",
+        "int": "number",
+        "float": "number",
+        "bool": "checkbox",
+        "datetime": "datetime-local",
+        "date": "date",
+    }
+    # Get the corresponding HTML input type; default to "text" if not found.
+    html_input_type = type_mapping.get(column_type, "text")
+    # Create the new input tag with the appropriate attributes.
+    new_input = soup.new_tag(
+        "input",
+        attrs={"type": html_input_type, "name": column_name, "required": "required"},
+    )
+    return new_input
+
+
 def _infer_html_indent_width(html_str: str) -> int:
     """
     Infer the indentation width (number of spaces) from the HTML string.
@@ -208,7 +230,7 @@ def _alter_js_function_returned_array(
     #  3. Captures the newline and leading whitespace (indent) of the return statement.
     #  4. Captures the rest of the return line.
     pattern = (
-        r"(function\s+"
+        r"("
         + re.escape(js_function_name)
         + r"\s*\([^)]*\)\s*\{)"  # group1: function header
         r"([\s\S]*?)"  # group2: code before return