zrb 1.0.0a3__tar.gz → 1.0.0a4__tar.gz

{zrb-1.0.0a3 → zrb-1.0.0a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0a3
+Version: 1.0.0a4
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -13,11 +13,16 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: rag
 Requires-Dist: autopep8 (>=2.0.4,<3.0.0)
+Requires-Dist: beautifulsoup4 (>=4.12.3,<5.0.0)
 Requires-Dist: black (>=24.10.0,<24.11.0)
+Requires-Dist: chromadb (>=0.5.20,<0.6.0) ; extra == "rag"
+Requires-Dist: fastapi[standard] (>=0.115.5,<0.116.0)
 Requires-Dist: isort (>=5.13.2,<5.14.0)
 Requires-Dist: libcst (>=1.5.0,<2.0.0)
 Requires-Dist: litellm (>=1.52.12,<2.0.0)
+Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb

{zrb-1.0.0a3 → zrb-1.0.0a4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.0.0a3"
+version = "1.0.0a4"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -46,6 +46,14 @@ isort = "~5.13.2"
 requests = "^2.32.3"
 libcst = "^1.5.0"
 litellm = "^1.52.12"
+chromadb = {version = "^0.5.20", optional = true}
+pdfplumber = {version = "^0.11.4", optional = true}
+beautifulsoup4 = "^4.12.3"
+fastapi = {extras = ["standard"], version = "^0.115.5"}
+
+[tool.poetry.extras]
+# poetry install -E rag
+rag = ["chromadb", "pdfplumber"]
 
 [tool.poetry.dev-dependencies]
 flake8 = "~7.1.1"

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/__init__.py

@@ -11,7 +11,7 @@ from zrb.callback.any_callback import AnyCallback
 from zrb.callback.callback import Callback
 from zrb.cmd.cmd_result import CmdResult
 from zrb.cmd.cmd_val import Cmd, CmdPath
-from zrb.config import SHOULD_LOAD_BUILTIN
+from zrb.config import LOAD_BUILTIN
 from zrb.content_transformer.any_content_transformer import AnyContentTransformer
 from zrb.content_transformer.content_transformer import ContentTransformer
 from zrb.context.any_context import AnyContext
@@ -99,7 +99,7 @@ assert Scheduler
 assert cli
 assert Xcom
 
-if SHOULD_LOAD_BUILTIN:
+if LOAD_BUILTIN:
     from zrb import builtin
 
     assert builtin

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/__main__.py

@@ -1,12 +1,14 @@
 import sys
 
 from zrb.runner.cli import cli
-from zrb.util.cli.style import stylize_error
+from zrb.util.cli.style import stylize_error, stylize_warning
 from zrb.util.group import NodeNotFoundError
 
 
 def serve_cli():
     try:
         cli.run(sys.argv[1:])
+    except KeyboardInterrupt:
+        print(stylize_warning("\nStopped"), file=sys.stderr)
     except NodeNotFoundError as e:
         print(stylize_error(f"{e}"), file=sys.stderr)

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/__init__.py

@@ -7,7 +7,7 @@ from zrb.builtin.git import (
     prune_local_branches,
 )
 from zrb.builtin.git_subtree import git_add_subtree, git_pull_subtree, git_push_subtree
-from zrb.builtin.llm import llm_chat
+from zrb.builtin.llm.llm_chat import llm_chat
 from zrb.builtin.md5 import hash_md5, sum_md5
 from zrb.builtin.project.add.fastapp import add_fastapp_to_project
 from zrb.builtin.project.create.create import create_project
@@ -15,7 +15,7 @@ from zrb.builtin.python import format_python_code
 from zrb.builtin.shell.autocomplete.bash import make_bash_autocomplete
 from zrb.builtin.shell.autocomplete.subcmd import get_shell_subcommands
 from zrb.builtin.shell.autocomplete.zsh import make_zsh_autocomplete
-from zrb.builtin.todo import todo_add, todo_edit, todo_list
+from zrb.builtin.todo import todo_add, todo_complete, todo_edit, todo_list, todo_log
 
 assert create_project
 assert add_fastapp_to_project
@@ -39,3 +39,5 @@ assert git_push_subtree
 assert todo_list
 assert todo_add
 assert todo_edit
+assert todo_complete
+assert todo_log

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/base64.py

@@ -1,5 +1,3 @@
-import base64
-
 from zrb.builtin.group import base64_group
 from zrb.context.any_context import AnyContext
 from zrb.input.str_input import StrInput
@@ -14,6 +12,8 @@ from zrb.task.make_task import make_task
     alias="encode",
 )
 def encode_base64(ctx: AnyContext) -> str:
+    import base64
+
     result = base64.b64encode(ctx.input.text.encode()).decode()
     ctx.print(result)
     return result
@@ -27,6 +27,8 @@ def encode_base64(ctx: AnyContext) -> str:
     alias="decode",
 )
 def decode_base64(ctx: AnyContext) -> str:
+    import base64
+
     result = base64.b64decode(ctx.input.text.encode()).decode()
     ctx.print(result)
     return result

zrb-1.0.0a3/src/zrb/builtin/llm.py → zrb-1.0.0a4/src/zrb/builtin/llm/llm_chat.py

@@ -1,10 +1,18 @@
 from zrb.builtin.group import llm_group
-from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
+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 = llm_group.add_task(
+llm_chat: LLMTask = llm_group.add_task(
     LLMTask(
         name="llm-chat",
         input=[
@@ -22,6 +30,7 @@ llm_chat = llm_group.add_task(
             ),
             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']}",
@@ -29,3 +38,10 @@ llm_chat = llm_group.add_task(
     ),
     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-1.0.0a4/src/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-1.0.0a4/src/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-1.0.0a4/src/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-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/md5.py

@@ -1,5 +1,3 @@
-import hashlib
-
 from zrb.builtin.group import md5_group
 from zrb.context.any_context import AnyContext
 from zrb.input.str_input import StrInput
@@ -14,6 +12,8 @@ from zrb.task.make_task import make_task
     alias="hash",
 )
 def hash_md5(ctx: AnyContext) -> str:
+    import hashlib
+
     result = hashlib.md5(ctx.input.text.encode()).hexdigest()
     ctx.print(result)
     return result
@@ -27,6 +27,8 @@ def hash_md5(ctx: AnyContext) -> str:
     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()

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/project/add/fastapp.py

@@ -29,7 +29,7 @@ scaffold_fastapp = Scaffolder(
         ),
     ],
     source_path=os.path.join(_DIR, "fastapp_template"),
-    auto_render_source_path=False,
+    render_source_path=False,
     destination_path=lambda ctx: os.path.join(
         ctx.input["project-dir"], ctx.input["app-name"]
     ),

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/project/add/fastapp_template/_zrb/helper.py

@@ -39,7 +39,7 @@ def create_migration(name: str, module: str) -> Task:
                 auto_render=True,
             ),
         ],
-        auto_render_cmd=False,
+        render_cmd=False,
         retries=2,
     )
 
@@ -68,7 +68,7 @@ def migrate_module(name: str, module: str, as_microservices: bool) -> Task:
             f"cd {os.path.join(APP_DIR, 'module', module)}",
             "alembic upgrade head",
         ],
-        auto_render_cmd=False,
+        render_cmd=False,
         retries=2,
     )
 
@@ -92,6 +92,6 @@ def run_microservice(name: str, port: int, module: str) -> Task:
             ACTIVATE_VENV_SCRIPT,
             'fastapi dev main.py --port "${APP_NAME_PORT}"',
         ],
-        auto_render_cmd=False,
+        render_cmd=False,
         retries=2,
     )

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/project/add/fastapp_template/_zrb/main.py

@@ -54,7 +54,7 @@ run_monolith = app_run_group.add_task(
             ACTIVATE_VENV_SCRIPT,
             'fastapi dev main.py --port "${APP_NAME_PORT}"',
         ],
-        auto_render_cmd=False,
+        render_cmd=False,
         retries=2,
     ),
     alias="monolith",

{zrb-1.0.0a3 → zrb-1.0.0a4}/src/zrb/builtin/project/create/create.py

@@ -25,7 +25,7 @@ scaffold_project = Scaffolder(
         ),
     ],
     source_path=os.path.join(_DIR, "project-template"),
-    auto_render_source_path=False,
+    render_source_path=False,
     destination_path="{ctx.input['project-dir']}",
     transform_content={"Project Name": "{ctx.input['project-name'].title()}"},
     retries=0,