zrb 0.26.2__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

zrb/builtin/group.py

@@ -1,48 +1,34 @@
-import sys
+from zrb.group.group import Group
+from zrb.runner.cli import cli
 
-from zrb.builtin.base64 import base64_group
-from zrb.builtin.devtool import devtool_group, devtool_install_group
-from zrb.builtin.docker import docker_group
-from zrb.builtin.env import env_group
-from zrb.builtin.explain import explain_group
-from zrb.builtin.git import git_group
-from zrb.builtin.md5 import md5_group
-from zrb.builtin.process import process_group, process_pid_group
-from zrb.builtin.project import (
-    project_add_app_group,
-    project_add_fastapp_group,
-    project_add_group,
-    project_add_task_group,
-    project_group,
-)
-from zrb.builtin.ubuntu import ubuntu_group, ubuntu_install_group
-from zrb.helper.accessories.color import colored
-
-_DEPRECATION_WARNING = """
-DEPRECATED: zrb.builtin.group
-User zrb.builtin instead
-
-```python
-from zrb.builtin import project_group
-```
-"""
+base64_group = cli.add_group(Group(name="base64", description="📄 Base64 operations"))
+git_group = cli.add_group(Group(name="git", description="🌱 Git related commands"))
+git_branch_group = git_group.add_group(
+    Group(name="branch", description="🌿 Git branch related commands")
+)
+git_subtree_group = git_group.add_group(
+    Group(name="subtree", description="🌳 Git subtree related commands")
+)
+llm_group = cli.add_group(Group(name="llm", description="🤖 LLM operations"))
+md5_group = cli.add_group(Group(name="md5", description="🔢 Md5 operations"))
+python_group = cli.add_group(
+    Group(name="python", description="🐍 Python related commands")
+)
+todo_group = cli.add_group(Group(name="todo", description="✅ Todo management"))
 
-print(colored(_DEPRECATION_WARNING, color="red", attrs=["bold"]), file=sys.stderr)
+shell_group = cli.add_group(
+    Group(name="shell", description="💬 Shell related commands")
+)
+shell_autocomplete_group: Group = shell_group.add_group(
+    Group(name="autocomplete", description="⌨️ Shell autocomplete related commands")
+)
 
-assert base64_group
-assert devtool_group
-assert devtool_install_group
-assert docker_group
-assert env_group
-assert explain_group
-assert git_group
-assert md5_group
-assert process_group
-assert process_pid_group
-assert project_group
-assert project_add_group
-assert project_add_app_group
-assert project_add_fastapp_group
-assert project_add_task_group
-assert ubuntu_group
-assert ubuntu_install_group
+project_group = cli.add_group(
+    Group(name="project", description="📁 Project related commands")
+)
+add_to_project_group = project_group.add_group(
+    Group(name="add", description="➕ Add things to project")
+)
+add_fastapp_to_project_group = add_to_project_group.add_group(
+    Group(name="fastapp", description="🚀 Add Fastapp resources")
+)

zrb/builtin/llm/llm_chat.py

@@ -0,0 +1,47 @@
+from zrb.builtin.group import llm_group
+from zrb.builtin.llm.tool.cli import run_shell_command
+from zrb.builtin.llm.tool.web import open_web_page, query_internet
+from zrb.config import (
+    LLM_ALLOW_ACCESS_SHELL,
+    LLM_ALLOW_ACCESS_WEB,
+    LLM_HISTORY_FILE,
+    LLM_MODEL,
+    LLM_SYSTEM_PROMPT,
+)
+from zrb.input.str_input import StrInput
+from zrb.input.text_input import TextInput
+from zrb.task.llm_task import LLMTask
+
+llm_chat: LLMTask = llm_group.add_task(
+    LLMTask(
+        name="llm-chat",
+        input=[
+            StrInput(
+                "model",
+                description="LLM Model",
+                prompt="LLM Model",
+                default_str=LLM_MODEL,
+            ),
+            StrInput(
+                "system-prompt",
+                description="System prompt",
+                prompt="System prompt",
+                default_str=LLM_SYSTEM_PROMPT,
+            ),
+            TextInput("message", description="User message", prompt="Your message"),
+        ],
+        history_file=LLM_HISTORY_FILE,
+        description="Chat with LLM",
+        model="{ctx.input.model}",
+        system_prompt="{ctx.input['system-prompt']}",
+        message="{ctx.input.message}",
+    ),
+    alias="chat",
+)
+
+if LLM_ALLOW_ACCESS_SHELL:
+    llm_chat.add_tool(run_shell_command)
+
+if LLM_ALLOW_ACCESS_WEB:
+    llm_chat.add_tool(open_web_page)
+    llm_chat.add_tool(query_internet)

zrb/builtin/llm/tool/cli.py

@@ -0,0 +1,9 @@
+import subprocess
+
+
+def run_shell_command(command: str) -> str:
+    """Running a shell command"""
+    output = subprocess.check_output(
+        command, shell=True, stderr=subprocess.STDOUT, text=True
+    )
+    return output

zrb/builtin/llm/tool/rag.py

@@ -0,0 +1,189 @@
+import json
+import os
+import sys
+from collections.abc import Callable, Iterable
+
+import litellm
+
+from zrb.config import (
+    RAG_CHUNK_SIZE,
+    RAG_EMBEDDING_MODEL,
+    RAG_MAX_RESULT_COUNT,
+    RAG_OVERLAP,
+)
+from zrb.util.cli.style import stylize_error, stylize_faint
+from zrb.util.run import run_async
+
+Document = str | Callable[[], str]
+Documents = Callable[[], Iterable[Document]] | Iterable[Document]
+
+
+def create_rag_from_directory(
+    tool_name: str,
+    tool_description: str,
+    document_dir_path: str = "./documents",
+    model: str = RAG_EMBEDDING_MODEL,
+    vector_db_path: str = "./chroma",
+    vector_db_collection: str = "documents",
+    chunk_size: int = RAG_CHUNK_SIZE,
+    overlap: int = RAG_OVERLAP,
+    max_result_count: int = RAG_MAX_RESULT_COUNT,
+):
+    return create_rag(
+        tool_name=tool_name,
+        tool_description=tool_description,
+        documents=get_rag_documents(os.path.expanduser(document_dir_path)),
+        model=model,
+        vector_db_path=vector_db_path,
+        vector_db_collection=vector_db_collection,
+        reset_db=get_rag_reset_db(
+            document_dir_path=os.path.expanduser(document_dir_path),
+            vector_db_path=os.path.expanduser(vector_db_path),
+        ),
+        chunk_size=chunk_size,
+        overlap=overlap,
+        max_result_count=max_result_count,
+    )
+
+
+def create_rag(
+    tool_name: str,
+    tool_description: str,
+    documents: Documents = [],
+    model: str = RAG_EMBEDDING_MODEL,
+    vector_db_path: str = "./chroma",
+    vector_db_collection: str = "documents",
+    reset_db: Callable[[], bool] | bool = False,
+    chunk_size: int = RAG_CHUNK_SIZE,
+    overlap: int = RAG_OVERLAP,
+    max_result_count: int = RAG_MAX_RESULT_COUNT,
+) -> Callable[[str], str]:
+    async def retrieve(query: str) -> str:
+        import chromadb
+        from chromadb.config import Settings
+
+        is_db_exist = os.path.isdir(vector_db_path)
+        client = chromadb.PersistentClient(
+            path=vector_db_path, settings=Settings(allow_reset=True)
+        )
+        should_reset_db = (
+            await run_async(reset_db()) if callable(reset_db) else reset_db
+        )
+        if (not is_db_exist) or should_reset_db:
+            client.reset()
+            collection = client.get_or_create_collection(vector_db_collection)
+            chunk_index = 0
+            print(stylize_faint("Scanning documents"), file=sys.stderr)
+            docs = await run_async(documents()) if callable(documents) else documents
+            for document in docs:
+                if callable(document):
+                    try:
+                        document = await run_async(document())
+                    except Exception as error:
+                        print(stylize_error(f"Error: {error}"), file=sys.stderr)
+                        continue
+                for i in range(0, len(document), chunk_size - overlap):
+                    chunk = document[i : i + chunk_size]
+                    if len(chunk) > 0:
+                        print(
+                            stylize_faint(f"Vectorize chunk {chunk_index}"),
+                            file=sys.stderr,
+                        )
+                        response = await litellm.aembedding(model=model, input=[chunk])
+                        vector = response["data"][0]["embedding"]
+                        print(
+                            stylize_faint(f"Adding chunk {chunk_index} to db"),
+                            file=sys.stderr,
+                        )
+                        collection.upsert(
+                            ids=[f"id{chunk_index}"],
+                            embeddings=[vector],
+                            documents=[chunk],
+                        )
+                        chunk_index += 1
+        collection = client.get_or_create_collection(vector_db_collection)
+        # Generate embedding for the query
+        print(stylize_faint("Vectorize query"), file=sys.stderr)
+        query_response = await litellm.aembedding(model=model, input=[query])
+        print(stylize_faint("Search documents"), file=sys.stderr)
+        # Search for the top_k most similar documents
+        results = collection.query(
+            query_embeddings=query_response["data"][0]["embedding"],
+            n_results=max_result_count,
+        )
+        return json.dumps(results)
+
+    retrieve.__name__ = tool_name
+    retrieve.__doc__ = tool_description
+    return retrieve
+
+
+def get_rag_documents(document_dir_path: str) -> Callable[[], list[Callable[[], str]]]:
+    def get_documents() -> list[Callable[[], str]]:
+        # Walk through the directory
+        readers = []
+        for root, _, files in os.walk(document_dir_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                if file_path.lower().endswith(".pdf"):
+                    readers.append(_get_pdf_reader(file_path))
+                    continue
+                readers.append(_get_text_reader(file_path))
+        return readers
+
+    return get_documents
+
+
+def _get_text_reader(file_path: str):
+    def read():
+        print(stylize_faint(f"Start reading {file_path}"), file=sys.stderr)
+        with open(file_path, "r", encoding="utf-8") as f:
+            content = f.read()
+        print(stylize_faint(f"Complete reading {file_path}"), file=sys.stderr)
+        return content
+
+    return read
+
+
+def _get_pdf_reader(file_path):
+    def read():
+        import pdfplumber
+
+        print(stylize_faint(f"Start reading {file_path}"), file=sys.stderr)
+        contents = []
+        with pdfplumber.open(file_path) as pdf:
+            for page in pdf.pages:
+                contents.append(page.extract_text())
+        print(stylize_faint(f"Complete reading {file_path}"), file=sys.stderr)
+        return "\n".join(contents)
+
+    return read
+
+
+def get_rag_reset_db(
+    document_dir_path: str, vector_db_path: str = "./chroma"
+) -> Callable[[], bool]:
+    def should_reset_db() -> bool:
+        document_exist = os.path.isdir(document_dir_path)
+        if not document_exist:
+            raise ValueError(f"Document directory not exists: {document_dir_path}")
+        vector_db_exist = os.path.isdir(vector_db_path)
+        if not vector_db_exist:
+            return True
+        document_mtime = _get_most_recent_mtime(document_dir_path)
+        vector_db_mtime = _get_most_recent_mtime(vector_db_path)
+        return document_mtime > vector_db_mtime
+
+    return should_reset_db
+
+
+def _get_most_recent_mtime(directory):
+    most_recent_mtime = 0
+    for root, dirs, files in os.walk(directory):
+        # Check mtime for directories
+        for name in dirs + files:
+            file_path = os.path.join(root, name)
+            mtime = os.path.getmtime(file_path)
+            if mtime > most_recent_mtime:
+                most_recent_mtime = mtime
+    return most_recent_mtime

zrb/builtin/llm/tool/web.py

@@ -0,0 +1,74 @@
+import json
+from typing import Annotated
+
+
+def open_web_page(url: str) -> str:
+    """Get content from a web page."""
+    import requests
+
+    response = requests.get(
+        url,
+        headers={
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"  # noqa
+        },
+    )
+    if response.status_code != 200:
+        raise Exception(
+            f"Error: Unable to retrieve search results (status code: {response.status_code})"  # noqa
+        )
+    return json.dumps(parse_html_text(response.text))
+
+
+def query_internet(
+    query: Annotated[str, "Search query"],
+    num_results: Annotated[int, "Search result count, by default 10"] = 10,
+) -> str:
+    """Search factual information from the internet by using Google."""
+    import requests
+
+    response = requests.get(
+        "https://google.com/search",
+        headers={
+            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"  # noqa
+        },
+        params={
+            "q": query,
+            "num": num_results,
+            "hl": "en",
+            "safe": "off",
+        },
+    )
+    if response.status_code != 200:
+        raise Exception(
+            f"Error: Unable to retrieve search results (status code: {response.status_code})"  # noqa
+        )
+    return json.dumps(parse_html_text(response.text))
+
+
+def parse_html_text(html_text: str) -> dict[str, str]:
+    from bs4 import BeautifulSoup
+
+    ignored_tags = [
+        "script",
+        "link",
+        "meta",
+        "style",
+        "code",
+        "footer",
+        "nav",
+        "header",
+        "aside",
+    ]
+    soup = BeautifulSoup(html_text, "html.parser")
+    links = []
+    for anchor in soup.find_all("a"):
+        if not anchor or "href" not in anchor.attrs:
+            continue
+        link: str = anchor["href"]
+        if link.startswith("#") or link.startswith("/"):
+            continue
+        links.append(link)
+    for tag in soup(ignored_tags):
+        tag.decompose()
+    html_text = soup.get_text(separator=" ", strip=True)
+    return {"content": html_text, "links_on_page": links}

zrb/builtin/md5.py

@@ -0,0 +1,36 @@
+from zrb.builtin.group import md5_group
+from zrb.context.any_context import AnyContext
+from zrb.input.str_input import StrInput
+from zrb.task.make_task import make_task
+
+
+@make_task(
+    name="hash-md5",
+    input=StrInput(name="text", description="Text", prompt="Text to encode"),
+    description="🧩 Hash text to md5",
+    group=md5_group,
+    alias="hash",
+)
+def hash_md5(ctx: AnyContext) -> str:
+    import hashlib
+
+    result = hashlib.md5(ctx.input.text.encode()).hexdigest()
+    ctx.print(result)
+    return result
+
+
+@make_task(
+    name="sum-md5",
+    input=StrInput(name="file", description="File name", prompt="File name"),
+    description="➕ Get md5 checksum of a file",
+    group=md5_group,
+    alias="sum",
+)
+def sum_md5(ctx: AnyContext) -> str:
+    import hashlib
+
+    with open(ctx.input.file, mode="rb") as file:
+        content = file.read()
+    result = hashlib.md5(content).hexdigest()
+    ctx.print(result)
+    return result

zrb/builtin/project/__init__.py

@@ -1,35 +0,0 @@
-from zrb.builtin.project._group import project_group
-from zrb.builtin.project.add import (
-    add_cmd_task,
-    add_docker_compose_task,
-    add_fastapp_application,
-    add_fastapp_crud,
-    add_fastapp_field,
-    add_fastapp_module,
-    add_plugin,
-    add_project_tasks,
-    add_python_app,
-    add_python_task,
-    project_add_app_group,
-    project_add_fastapp_group,
-    project_add_group,
-    project_add_task_group,
-)
-from zrb.builtin.project.create import create_project
-
-assert project_group
-assert project_add_group
-assert project_add_app_group
-assert project_add_fastapp_group
-assert project_add_task_group
-assert create_project
-assert add_plugin
-assert add_project_tasks
-assert add_cmd_task
-assert add_docker_compose_task
-assert add_fastapp_application
-assert add_fastapp_crud
-assert add_fastapp_module
-assert add_fastapp_field
-assert add_python_task
-assert add_python_app

zrb/builtin/project/add/__init__.py

@@ -1,37 +0,0 @@
-from zrb.builtin.project.add._group import project_add_group
-from zrb.builtin.project.add.app import (
-    add_app_generator,
-    add_python_app,
-    project_add_app_group,
-)
-from zrb.builtin.project.add.fastapp import (
-    add_fastapp_application,
-    add_fastapp_crud,
-    add_fastapp_field,
-    add_fastapp_module,
-    project_add_fastapp_group,
-)
-from zrb.builtin.project.add.plugin import add_plugin
-from zrb.builtin.project.add.project_task import add_project_tasks
-from zrb.builtin.project.add.task import (
-    add_cmd_task,
-    add_docker_compose_task,
-    add_python_task,
-    project_add_task_group,
-)
-
-assert project_add_group
-assert project_add_app_group
-assert project_add_fastapp_group
-assert project_add_task_group
-assert add_app_generator
-assert add_python_app
-assert add_fastapp_application
-assert add_fastapp_crud
-assert add_fastapp_field
-assert add_fastapp_module
-assert add_plugin
-assert add_project_tasks
-assert add_cmd_task
-assert add_docker_compose_task
-assert add_python_task

zrb/builtin/project/add/fastapp.py

@@ -0,0 +1,72 @@
+import os
+
+from zrb.builtin.group import add_to_project_group
+from zrb.context.any_context import AnyContext
+from zrb.input.str_input import StrInput
+from zrb.task.make_task import make_task
+from zrb.task.scaffolder import Scaffolder
+from zrb.task.task import Task
+from zrb.util.string.conversion import double_quote
+from zrb.util.string.name import get_random_name
+
+_DIR = os.path.dirname(__file__)
+
+
+scaffold_fastapp = Scaffolder(
+    name="scaffold-fastapp",
+    input=[
+        StrInput(
+            name="project-dir",
+            description="Project directory",
+            prompt="Project directory",
+            default_str=lambda _: os.getcwd(),
+        ),
+        StrInput(
+            name="app-name",
+            description="App name",
+            prompt="App name",
+            default_str=lambda _: get_random_name(separator="_"),
+        ),
+    ],
+    source_path=os.path.join(_DIR, "fastapp_template"),
+    render_source_path=False,
+    destination_path=lambda ctx: os.path.join(
+        ctx.input["project-dir"], ctx.input["app-name"]
+    ),
+    transform_content={
+        "fastapp_template": "{to_snake_case(ctx.input['app-name'])}",
+        "App Name": "{ctx.input['app-name'].title()}",
+        "App name": "{ctx.input['app-name'].capitalize()}",
+        "app-name": "{to_kebab_case(ctx.input['app-name'])}",
+        "app_name": "{to_snake_case(ctx.input['app-name'])}",
+        "APP_NAME": "{to_snake_case(ctx.input['app-name']).upper()}",
+        "secure-password": lambda _: get_random_name(),
+    },
+    retries=0,
+)
+
+
+@make_task(
+    name="register-fastapp-automation",
+)
+def register_fastapp_automation(ctx: AnyContext):
+    project_dir_path = ctx.input["project-dir"]
+    zrb_init_path = os.path.join(project_dir_path, "zrb_init.py")
+    app_dir_path = ctx.input["app-name"]
+    app_automation_file_part = ", ".join(
+        [double_quote(part) for part in [app_dir_path, "_zrb", "main.py"]]
+    )
+    with open(zrb_init_path, "+a") as f:
+        f.write(f"load_file(os.path.join(_DIR, {app_automation_file_part}))\n")
+
+
+scaffold_fastapp >> register_fastapp_automation
+
+add_fastapp_to_project = add_to_project_group.add_task(
+    Task(
+        name="add-fastapp",
+        description="🚀 Add FastApp to project",
+    ),
+    alias="fastapp",
+)
+add_fastapp_to_project << [register_fastapp_automation]

zrb/builtin/project/add/fastapp_template/.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+.venv
+.env
+*.db

zrb/builtin/project/add/fastapp_template/README.md

@@ -0,0 +1,7 @@
+# App Name
+
+# Principle
+
+- Developer should be able to override everything with sane amount of code.
+- Simple tasks should only require small amount of code.
+- A bit of magic is okay as long as transparent and documented.

zrb/builtin/project/add/fastapp_template/_zrb/config.py

@@ -0,0 +1,17 @@
+import os
+import platform
+
+DIR = os.path.dirname(__file__)
+APP_DIR = os.path.dirname(DIR)
+APP_MODULE_NAME = os.path.basename(APP_DIR)
+
+MICROSERVICES_ENV_VARS = {
+    "APP_NAME_MODE": "microservices",
+    "APP_NAME_AUTH_BASE_URL": "http://localhost:3002",
+}
+MONOLITH_ENV_VARS = {"APP_NAME_MODE": "monolith"}
+
+if platform.system() == "Windows":
+    ACTIVATE_VENV_SCRIPT = "Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser; . .venv\Scripts\Activate"  # noqa
+else:
+    ACTIVATE_VENV_SCRIPT = "source .venv/bin/activate"

zrb/builtin/project/add/fastapp_template/_zrb/group.py

@@ -0,0 +1,16 @@
+from zrb import Group
+from zrb.builtin import project_group
+
+app_group = project_group.add_group(
+    Group(name="app-name", description="🚀 Managing App Name")
+)
+
+app_run_group = app_group.add_group(Group(name="run", description="🟢 Run App Name"))
+
+app_migrate_group = app_group.add_group(
+    Group(name="migrate", description="📦 Run App Name DB migration")
+)
+
+app_create_group = app_group.add_group(
+    Group(name="create", description="✨ Create resources for App Name")
+)