zrb 1.0.0a2__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

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/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")
+)

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

@@ -0,0 +1,97 @@
+import os
+
+from fastapp_template._zrb.config import (
+    ACTIVATE_VENV_SCRIPT,
+    APP_DIR,
+    MICROSERVICES_ENV_VARS,
+    MONOLITH_ENV_VARS,
+)
+
+from zrb import Cmd, CmdTask, EnvFile, EnvMap, StrInput, Task
+
+
+def create_migration(name: str, module: str) -> Task:
+    return CmdTask(
+        name=f"create-app-name-{name}-migration",
+        description=f"🧩 Create App Name {name.capitalize()} DB migration",
+        input=StrInput(
+            name="message",
+            description="Migration message",
+            prompt="Migration message",
+            allow_empty=False,
+        ),
+        env=[
+            EnvFile(path=os.path.join(APP_DIR, "template.env")),
+            EnvMap(
+                vars={
+                    "APP_DB_URL": f"sqlite:///{APP_DIR}/.migration.{module}.db",
+                    "APP_NAME_MODULES": f"{module}",
+                }
+            ),
+        ],
+        cwd=APP_DIR,
+        cmd=[
+            ACTIVATE_VENV_SCRIPT,
+            f"cd {os.path.join(APP_DIR, 'module', module)}",
+            "alembic upgrade head",
+            Cmd(
+                "alembic revision --autogenerate -m {double_quote(ctx.input.message)}",
+                auto_render=True,
+            ),
+        ],
+        render_cmd=False,
+        retries=2,
+    )
+
+
+def migrate_module(name: str, module: str, as_microservices: bool) -> Task:
+    env_vars = MICROSERVICES_ENV_VARS if as_microservices else MONOLITH_ENV_VARS
+    return CmdTask(
+        name=(
+            f"migrate-app-name-{name}"
+            if as_microservices
+            else f"migrate-{name}-on-monolith"
+        ),
+        description=f"🧩 Run App Name {name.capitalize()} DB migration",
+        env=[
+            EnvFile(path=os.path.join(APP_DIR, "template.env")),
+            EnvMap(
+                vars={
+                    **env_vars,
+                    "APP_NAME_MODULES": f"{module}",
+                }
+            ),
+        ],
+        cwd=APP_DIR,
+        cmd=[
+            ACTIVATE_VENV_SCRIPT,
+            f"cd {os.path.join(APP_DIR, 'module', module)}",
+            "alembic upgrade head",
+        ],
+        render_cmd=False,
+        retries=2,
+    )
+
+
+def run_microservice(name: str, port: int, module: str) -> Task:
+    return CmdTask(
+        name=f"run-app-name-{name}",
+        description=f"🧩 Run App Name {name.capitalize()}",
+        env=[
+            EnvFile(path=os.path.join(APP_DIR, "template.env")),
+            EnvMap(
+                vars={
+                    **MICROSERVICES_ENV_VARS,
+                    "APP_NAME_PORT": f"{port}",
+                    "APP_NAME_MODULES": f"{module}",
+                }
+            ),
+        ],
+        cwd=APP_DIR,
+        cmd=[
+            ACTIVATE_VENV_SCRIPT,
+            'fastapi dev main.py --port "${APP_NAME_PORT}"',
+        ],
+        render_cmd=False,
+        retries=2,
+    )

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

@@ -0,0 +1,132 @@
+import os
+
+from fastapp_template._zrb.config import ACTIVATE_VENV_SCRIPT, APP_DIR
+from fastapp_template._zrb.group import (
+    app_create_group,
+    app_migrate_group,
+    app_run_group,
+)
+from fastapp_template._zrb.helper import (
+    create_migration,
+    migrate_module,
+    run_microservice,
+)
+from fastapp_template._zrb.venv_task import prepare_venv
+
+from zrb import CmdTask, Env, EnvFile, Task
+
+# 🚀 Run/Migrate All ===========================================================
+
+run_all = app_run_group.add_task(
+    Task(
+        name="run-app-name", description="🟢 Run App Name as monolith and microservices"
+    ),
+    alias="all",
+)
+
+migrate_all = app_migrate_group.add_task(
+    Task(
+        name="migrate-app-name",
+        description="📦 Run App Name DB migration for monolith and microservices",
+    ),
+    alias="all",
+)
+
+create_all_migration = app_create_group.add_task(
+    Task(
+        name="create-app-name-migration", description="📦 Create App Name DB migration"
+    ),
+    alias="migration",
+)
+
+# 🗿 Run/Migrate Monolith =====================================================
+
+run_monolith = app_run_group.add_task(
+    CmdTask(
+        name="run-monolith-app-name",
+        description="🗿 Run App Name as a monolith",
+        env=[
+            EnvFile(path=os.path.join(APP_DIR, "template.env")),
+            Env(name="APP_NAME_MODE", default="monolith"),
+        ],
+        cwd=APP_DIR,
+        cmd=[
+            ACTIVATE_VENV_SCRIPT,
+            'fastapi dev main.py --port "${APP_NAME_PORT}"',
+        ],
+        render_cmd=False,
+        retries=2,
+    ),
+    alias="monolith",
+)
+prepare_venv >> run_monolith >> run_all
+
+migrate_monolith = app_migrate_group.add_task(
+    Task(
+        name="migrate-monolith-app-name",
+        description="🗿 Run App Name DB migration for monolith",
+    ),
+    alias="monolith",
+)
+migrate_monolith >> migrate_all
+
+# 🌐 Run/Migrate Microsevices ==================================================
+
+run_microservices = app_run_group.add_task(
+    Task(
+        name="run-microservices-app-name",
+        description="🌐 Run App Name as microservices",
+    ),
+    alias="microservices",
+)
+run_microservices >> run_all
+
+migrate_microservices = app_migrate_group.add_task(
+    Task(
+        name="migrate-microservices-app-name",
+        description="🌐 Run App Name DB migration for microservices",
+    ),
+    alias="microservices",
+)
+migrate_microservices >> migrate_all
+
+# 📡 Run/Migrate Gateway =======================================================
+
+run_gateway = app_run_group.add_task(
+    run_microservice("gateway", 3001, "gateway"), alias="microservices-gateway"
+)
+prepare_venv >> run_gateway >> run_microservices
+
+create_gateway_migration = app_create_group.add_task(
+    create_migration("gateway", "gateway"), alias="gateway-migration"
+)
+prepare_venv >> create_gateway_migration >> create_all_migration
+
+migrate_monolith_gateway = migrate_module("gateway", "gateway", as_microservices=False)
+prepare_venv >> migrate_monolith_gateway >> [migrate_monolith, run_monolith]
+
+migrate_microservices_gateway = app_migrate_group.add_task(
+    migrate_module("gateway", "gateway", as_microservices=True),
+    alias="microservices-gateway",
+)
+prepare_venv >> migrate_microservices_gateway >> [migrate_microservices, run_gateway]
+
+# 🔐 Run/Migrate Auth ==========================================================
+
+run_auth = app_run_group.add_task(
+    run_microservice("auth", 3002, "auth"), alias="microservices-auth"
+)
+prepare_venv >> run_auth >> run_microservices
+
+create_auth_migration = app_create_group.add_task(
+    create_migration("auth", "auth"), alias="auth-migration"
+)
+prepare_venv >> create_auth_migration >> create_all_migration
+
+migrate_monolith_auth = migrate_module("auth", "auth", as_microservices=False)
+prepare_venv >> migrate_monolith_auth >> [migrate_monolith, run_monolith]
+
+migrate_microservices_auth = app_migrate_group.add_task(
+    migrate_module("auth", "auth", as_microservices=True), alias="microservices-auth"
+)
+prepare_venv >> migrate_microservices_auth >> [migrate_microservices, run_auth]

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

@@ -0,0 +1,22 @@
+import os
+
+from fastapp_template._zrb.config import ACTIVATE_VENV_SCRIPT, APP_DIR
+from fastapp_template._zrb.group import app_group
+
+from zrb import CmdTask
+
+create_venv = CmdTask(
+    name="create-app-name-venv",
+    cwd=APP_DIR,
+    cmd="python -m venv .venv",
+    execute_condition=lambda _: not os.path.isdir(os.path.join(APP_DIR, ".venv")),
+)
+
+prepare_venv = CmdTask(
+    name="prepare-app-name-venv",
+    cmd=[ACTIVATE_VENV_SCRIPT, "pip install -r requirements.txt"],
+    cwd=APP_DIR,
+)
+create_venv >> prepare_venv
+
+app_group.add_task(create_venv, alias="prepare")

zrb/builtin/project/add/fastapp_template/common/app.py

@@ -0,0 +1,18 @@
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+from fastapp_template.common.db_engine import engine
+from fastapp_template.config import APP_MODE, APP_MODULES
+from sqlmodel import SQLModel
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    SQLModel.metadata.create_all(engine)
+    yield
+
+
+app_title = (
+    "App Name" if APP_MODE == "monolith" else f"App Name - {', '.join(APP_MODULES)}"
+)
+app = FastAPI(title=app_title, lifespan=lifespan)