zrb 1.5.17__py3-none-any.whl → 1.6.0__py3-none-any.whl

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 LOAD_BUILTIN
+from zrb.config import CFG
 from zrb.content_transformer.any_content_transformer import AnyContentTransformer
 from zrb.content_transformer.content_transformer import ContentTransformer
 from zrb.context.any_context import AnyContext
@@ -109,7 +109,7 @@ assert Xcom
 assert web_config
 assert User
 
-if LOAD_BUILTIN:
+if CFG.LOAD_BUILTIN:
     from zrb import builtin
 
     assert builtin

zrb/__main__.py

@@ -2,7 +2,7 @@ import logging
 import os
 import sys
 
-from zrb.config import INIT_MODULES, INIT_SCRIPTS, LOGGER, LOGGING_LEVEL
+from zrb.config import CFG
 from zrb.runner.cli import cli
 from zrb.util.cli.style import stylize_error, stylize_faint, stylize_warning
 from zrb.util.group import NodeNotFoundError
@@ -23,29 +23,29 @@ class FaintFormatter(logging.Formatter):
 
 
 def serve_cli():
-    LOGGER.setLevel(LOGGING_LEVEL)
+    CFG.LOGGER.setLevel(CFG.LOGGING_LEVEL)
     # Remove existing handlers to avoid duplicates/default formatting
-    for handler in LOGGER.handlers[:]:
-        LOGGER.removeHandler(handler)
+    for handler in CFG.LOGGER.handlers[:]:
+        CFG.LOGGER.removeHandler(handler)
     handler = logging.StreamHandler()
     handler.setFormatter(FaintFormatter())
-    LOGGER.addHandler(handler)
+    CFG.LOGGER.addHandler(handler)
     # --- End Logging Configuration ---
     try:
         # load init modules
-        for init_module in INIT_MODULES:
-            LOGGER.info(f"Loading {init_module}")
+        for init_module in CFG.INIT_MODULES:
+            CFG.LOGGER.info(f"Loading {init_module}")
             load_module(init_module)
         zrb_init_path_list = _get_zrb_init_path_list()
         # load init scripts
-        for init_script in INIT_SCRIPTS:
+        for init_script in CFG.INIT_SCRIPTS:
             abs_init_script = os.path.abspath(os.path.expanduser(init_script))
             if abs_init_script not in zrb_init_path_list:
-                LOGGER.info(f"Loading {abs_init_script}")
+                CFG.LOGGER.info(f"Loading {abs_init_script}")
                 load_file(abs_init_script, -1)
         # load zrb init
         for zrb_init_path in zrb_init_path_list:
-            LOGGER.info(f"Loading {zrb_init_path}")
+            CFG.LOGGER.info(f"Loading {zrb_init_path}")
             load_file(zrb_init_path)
         # run the CLI
         cli.run(sys.argv[1:])
@@ -69,8 +69,8 @@ def _get_zrb_init_path_list() -> list[str]:
         dir_path_list.append(current_path)
     zrb_init_path_list = []
     for current_path in dir_path_list[::-1]:
-        zrb_init_path = os.path.join(current_path, "zrb_init.py")
-        LOGGER.info(f"Finding {zrb_init_path}")
+        zrb_init_path = os.path.join(current_path, CFG.INIT_FILE_NAME)
+        CFG.LOGGER.info(f"Finding {zrb_init_path}")
         if os.path.isfile(zrb_init_path):
             zrb_init_path_list.append(zrb_init_path)
     return zrb_init_path_list

zrb/builtin/__init__.py

@@ -9,7 +9,7 @@ from zrb.builtin.git import (
 from zrb.builtin.git_subtree import git_add_subtree, git_pull_subtree, git_push_subtree
 from zrb.builtin.http import generate_curl, http_request
 from zrb.builtin.jwt import decode_jwt, encode_jwt, validate_jwt
-from zrb.builtin.llm.llm_chat import llm_chat
+from zrb.builtin.llm.llm_ask import llm_ask
 from zrb.builtin.md5 import hash_md5, sum_md5, validate_md5
 from zrb.builtin.project.add.fastapp.fastapp_task import add_fastapp_to_project
 from zrb.builtin.project.create.project_task import create_project
@@ -55,7 +55,7 @@ assert validate_base64
 assert encode_jwt
 assert decode_jwt
 assert validate_jwt
-assert llm_chat
+assert llm_ask
 assert hash_md5
 assert sum_md5
 assert validate_md5

zrb/builtin/llm/chat_session.py

@@ -0,0 +1,202 @@
+"""
+This module provides functions for managing interactive chat sessions with an LLM.
+
+It handles reading user input, triggering the LLM task, and managing the
+conversation flow via XCom.
+"""
+
+import asyncio
+
+from zrb.context.any_context import AnyContext
+from zrb.util.cli.style import stylize_bold_yellow, stylize_faint
+
+
+async def read_user_prompt(ctx: AnyContext) -> str:
+    """
+    Reads user input from the CLI for an interactive chat session.
+
+    Handles special commands like /bye, /multi, /end, and /help.
+    Triggers the LLM task and waits for the result.
+
+    Args:
+        ctx: The context object for the task.
+
+    Returns:
+        The final result from the LLM session.
+    """
+    _show_info(ctx)
+    final_result = ""
+    ctx.print(stylize_faint("🧑 >> ") + f"{ctx.input.message}", plain=True)
+    result = await _trigger_ask_and_wait_for_result(
+        ctx,
+        user_prompt=ctx.input.message,
+        previous_session_name=ctx.input.previous_session,
+        start_new=ctx.input.start_new,
+    )
+    if ctx.env.get("_ZRB_WEB_ENV", "0") != "0":
+        # Don't run in web environment
+        if result is not None:
+            final_result = result
+        return final_result
+    multiline_mode = False
+    user_inputs = []
+    while True:
+        await asyncio.sleep(0.01)
+        ctx.print(stylize_faint("🧑 >> "), end="", plain=True)
+        user_input = input()
+        # Handle special input
+        if user_input.strip().lower() in ("/bye", "/quit"):
+            user_prompt = "\n".join(user_inputs)
+            user_inputs = []
+            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
+            if result is not None:
+                final_result = result
+            break
+        elif user_input.strip().lower() in ("/multi"):
+            multiline_mode = True
+        elif user_input.strip().lower() in ("/end"):
+            multiline_mode = False
+            user_prompt = "\n".join(user_inputs)
+            user_inputs = []
+            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
+            if result is not None:
+                final_result = result
+        elif user_input.strip().lower() in ("/help", "/info"):
+            _show_info(ctx)
+            continue
+        else:
+            user_inputs.append(user_input)
+            if multiline_mode:
+                continue
+            user_prompt = "\n".join(user_inputs)
+            user_inputs = []
+            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
+            if result is not None:
+                final_result = result
+    return final_result
+
+
+async def _trigger_ask_and_wait_for_result(
+    ctx: AnyContext,
+    user_prompt: str,
+    previous_session_name: str | None = None,
+    start_new: bool = False,
+) -> str | None:
+    """
+    Triggers the LLM ask task and waits for the result via XCom.
+
+    Args:
+        ctx: The context object for the task.
+        user_prompt: The user's message to send to the LLM.
+        previous_session_name: The name of the previous chat session (optional).
+        start_new: Whether to start a new conversation (optional).
+
+    Returns:
+        The result from the LLM task, or None if the user prompt is empty.
+    """
+    if user_prompt.strip() == "":
+        return None
+    await _trigger_ask(ctx, user_prompt, previous_session_name, start_new)
+    result = await _wait_ask_result(ctx)
+    ctx.print(stylize_faint("🤖 >> ") + result, plain=True)
+    return result
+
+
+def get_llm_ask_input_mapping(callback_ctx: AnyContext):
+    """
+    Generates the input mapping for the LLM ask task from the callback context.
+
+    Args:
+        callback_ctx: The context object for the callback.
+
+    Returns:
+        A dictionary containing the input mapping for the LLM ask task.
+    """
+    data = callback_ctx.xcom.ask_trigger.pop()
+    return {
+        "model": callback_ctx.input.model,
+        "base-url": callback_ctx.input.base_url,
+        "api-key": callback_ctx.input.api_key,
+        "system-prompt": callback_ctx.input.system_prompt,
+        "start-new": data.get("start_new"),
+        "previous-session": data.get("previous_session_name"),
+        "message": data.get("message"),
+    }
+
+
+async def _trigger_ask(
+    ctx: AnyContext,
+    user_prompt: str,
+    previous_session_name: str | None = None,
+    start_new: bool = False,
+):
+    """
+    Triggers the LLM ask task by pushing data to the 'ask_trigger' XCom queue.
+
+    Args:
+        ctx: The context object for the task.
+        user_prompt: The user's message to send to the LLM.
+        previous_session_name: The name of the previous chat session (optional).
+        start_new: Whether to start a new conversation (optional).
+    """
+    if previous_session_name is None:
+        previous_session_name = await _wait_ask_session_name(ctx)
+    ctx.xcom["ask_trigger"].push(
+        {
+            "previous_session_name": previous_session_name,
+            "start_new": start_new,
+            "message": user_prompt,
+        }
+    )
+
+
+async def _wait_ask_result(ctx: AnyContext) -> str:
+    """
+    Waits for and retrieves the LLM task result from the 'ask_result' XCom queue.
+
+    Args:
+        ctx: The context object for the task.
+
+    Returns:
+        The result string from the LLM task.
+    """
+    while "ask_result" not in ctx.xcom or len(ctx.xcom.ask_result) == 0:
+        await asyncio.sleep(0.1)
+    return ctx.xcom.ask_result.pop()
+
+
+async def _wait_ask_session_name(ctx: AnyContext) -> str:
+    """
+    Waits for and retrieves the LLM chat session name from the 'ask_session_name' XCom queue.
+
+    Args:
+        ctx: The context object for the task.
+
+    Returns:
+        The session name string.
+    """
+    while "ask_session_name" not in ctx.xcom or len(ctx.xcom.ask_session_name) == 0:
+        await asyncio.sleep(0.1)
+    return ctx.xcom.ask_session_name.pop()
+
+
+def _show_info(ctx: AnyContext):
+    """
+    Displays the available chat session commands to the user.
+
+    Args:
+        ctx: The context object for the task.
+    """
+    ctx.print(
+        stylize_bold_yellow(
+            "\n".join(
+                [
+                    "/bye:   Quit from chat session",
+                    "/multi: Start multiline input",
+                    "/end:   End multiline input",
+                    "/help:  Show this message",
+                ]
+            )
+        ),
+        plain=True,
+    )

zrb/builtin/llm/history.py

@@ -2,7 +2,7 @@ import json
 import os
 from typing import Any
 
-from zrb.config import LLM_HISTORY_DIR
+from zrb.config import CFG
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.task.llm.history import ConversationHistoryData
 from zrb.util.file import read_file, write_file
@@ -17,7 +17,7 @@ def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | Non
         return None  # Indicate no history to load
     previous_session_name = ctx.input.previous_session
     if not previous_session_name:  # Check for empty string or None
-        last_session_file_path = os.path.join(LLM_HISTORY_DIR, "last-session")
+        last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
         if os.path.isfile(last_session_file_path):
             previous_session_name = read_file(last_session_file_path).strip()
             if not previous_session_name:  # Handle empty last-session file
@@ -25,7 +25,7 @@ def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | Non
         else:
             return None  # No previous session specified and no last session found
     conversation_file_path = os.path.join(
-        LLM_HISTORY_DIR, f"{previous_session_name}.json"
+        CFG.LLM_HISTORY_DIR, f"{previous_session_name}.json"
     )
     if not os.path.isfile(conversation_file_path):
         ctx.log_warning(f"History file not found: {conversation_file_path}")
@@ -55,19 +55,19 @@ def write_chat_conversation(
     ctx: AnySharedContext, history_data: ConversationHistoryData
 ):
     """Writes the conversation history data (including context) to a session file."""
-    os.makedirs(LLM_HISTORY_DIR, exist_ok=True)
+    os.makedirs(CFG.LLM_HISTORY_DIR, exist_ok=True)
     current_session_name = ctx.session.name
     if not current_session_name:
         ctx.log_warning("Cannot write history: Session name is empty.")
         return
     conversation_file_path = os.path.join(
-        LLM_HISTORY_DIR, f"{current_session_name}.json"
+        CFG.LLM_HISTORY_DIR, f"{current_session_name}.json"
     )
     try:
         # Use model_dump_json to serialize the Pydantic model
         write_file(conversation_file_path, history_data.model_dump_json(indent=2))
         # Update the last-session pointer
-        last_session_file_path = os.path.join(LLM_HISTORY_DIR, "last-session")
+        last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
         write_file(last_session_file_path, current_session_name)
     except Exception as e:
         ctx.log_error(f"Error writing history file '{conversation_file_path}': {e}")

zrb/builtin/llm/llm_ask.py

@@ -0,0 +1,142 @@
+from zrb.builtin.group import llm_group
+from zrb.builtin.llm.chat_session import get_llm_ask_input_mapping, read_user_prompt
+from zrb.builtin.llm.history import read_chat_conversation, write_chat_conversation
+from zrb.builtin.llm.input import PreviousSessionInput
+from zrb.builtin.llm.tool.api import get_current_location, get_current_weather
+from zrb.builtin.llm.tool.cli import run_shell_command
+from zrb.builtin.llm.tool.file import (
+    apply_diff,
+    list_files,
+    read_from_file,
+    search_files,
+    write_to_file,
+)
+from zrb.builtin.llm.tool.web import (
+    create_search_internet_tool,
+    open_web_page,
+    search_arxiv,
+    search_wikipedia,
+)
+from zrb.callback.callback import Callback
+from zrb.config import CFG
+from zrb.input.bool_input import BoolInput
+from zrb.input.str_input import StrInput
+from zrb.input.text_input import TextInput
+from zrb.task.base_trigger import BaseTrigger
+from zrb.task.llm_task import LLMTask
+
+_llm_ask_inputs = [
+    StrInput(
+        "model",
+        description="LLM Model",
+        prompt="LLM Model",
+        default="",
+        allow_positional_parsing=False,
+        always_prompt=False,
+        allow_empty=True,
+    ),
+    StrInput(
+        "base-url",
+        description="LLM API Base URL",
+        prompt="LLM API Base URL",
+        default="",
+        allow_positional_parsing=False,
+        always_prompt=False,
+        allow_empty=True,
+    ),
+    StrInput(
+        "api-key",
+        description="LLM API Key",
+        prompt="LLM API Key",
+        default="",
+        allow_positional_parsing=False,
+        always_prompt=False,
+        allow_empty=True,
+    ),
+    TextInput(
+        "system-prompt",
+        description="System prompt",
+        prompt="System prompt",
+        default="",
+        allow_positional_parsing=False,
+        always_prompt=False,
+    ),
+    BoolInput(
+        "start-new",
+        description="Start new conversation (LLM will forget everything)",
+        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(
+        "previous-session",
+        description="Previous conversation session",
+        prompt="Previous conversation session (can be empty)",
+        allow_positional_parsing=False,
+        allow_empty=True,
+        always_prompt=False,
+    ),
+]
+
+llm_ask: LLMTask = llm_group.add_task(
+    LLMTask(
+        name="llm-ask",
+        input=_llm_ask_inputs,
+        description="❓ Ask LLM",
+        model=lambda ctx: None if ctx.input.model.strip() == "" else ctx.input.model,
+        model_base_url=lambda ctx: (
+            None if ctx.input.base_url.strip() == "" else ctx.input.base_url
+        ),
+        model_api_key=lambda ctx: (
+            None if ctx.input.api_key.strip() == "" else ctx.input.api_key
+        ),
+        conversation_history_reader=read_chat_conversation,
+        conversation_history_writer=write_chat_conversation,
+        system_prompt=lambda ctx: (
+            None if ctx.input.system_prompt.strip() == "" else ctx.input.system_prompt
+        ),
+        message="{ctx.input.message}",
+        retries=0,
+    ),
+    alias="ask",
+)
+
+llm_group.add_task(
+    BaseTrigger(
+        name="llm-chat",
+        input=_llm_ask_inputs,
+        description="💬 Chat with LLM",
+        queue_name="ask_trigger",
+        action=read_user_prompt,
+        callback=Callback(
+            task=llm_ask,
+            input_mapping=get_llm_ask_input_mapping,
+            result_queue="ask_result",
+            session_name_queue="ask_session_name",
+        ),
+        retries=0,
+        cli_only=True,
+    ),
+    alias="chat",
+)
+
+if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
+    llm_ask.add_tool(list_files)
+    llm_ask.add_tool(read_from_file)
+    llm_ask.add_tool(write_to_file)
+    llm_ask.add_tool(search_files)
+    llm_ask.add_tool(apply_diff)
+
+if CFG.LLM_ALLOW_ACCESS_SHELL:
+    llm_ask.add_tool(run_shell_command)
+
+if CFG.LLM_ALLOW_ACCESS_INTERNET:
+    llm_ask.add_tool(open_web_page)
+    llm_ask.add_tool(search_wikipedia)
+    llm_ask.add_tool(search_arxiv)
+    if CFG.SERP_API_KEY != "":
+        llm_ask.add_tool(create_search_internet_tool(CFG.SERP_API_KEY))
+    llm_ask.add_tool(get_current_location)
+    llm_ask.add_tool(get_current_weather)

zrb/builtin/llm/tool/rag.py

@@ -8,14 +8,7 @@ from textwrap import dedent
 
 import ulid
 
-from zrb.config import (
-    RAG_CHUNK_SIZE,
-    RAG_EMBEDDING_API_KEY,
-    RAG_EMBEDDING_BASE_URL,
-    RAG_EMBEDDING_MODEL,
-    RAG_MAX_RESULT_COUNT,
-    RAG_OVERLAP,
-)
+from zrb.config import CFG
 from zrb.util.cli.style import stylize_error, stylize_faint
 from zrb.util.file import read_file
 
@@ -42,13 +35,13 @@ def create_rag_from_directory(
     document_dir_path: str = "./documents",
     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,
+    chunk_size: int | None = None,
+    overlap: int | None = None,
+    max_result_count: int | None = None,
     file_reader: list[RAGFileReader] = [],
-    openai_api_key: str = RAG_EMBEDDING_API_KEY,
-    openai_base_url: str = RAG_EMBEDDING_BASE_URL,
-    openai_embedding_model: str = RAG_EMBEDDING_MODEL,
+    openai_api_key: str | None = None,
+    openai_base_url: str | None = None,
+    openai_embedding_model: str | None = None,
 ):
     """Create a RAG retrieval tool function for LLM use.
     This factory configures and returns an async function that takes a query,
@@ -64,10 +57,33 @@ def create_rag_from_directory(
 
         # Initialize OpenAI client with custom URL if provided
         client_args = {}
-        if openai_api_key:
-            client_args["api_key"] = openai_api_key
-        if openai_base_url:
-            client_args["base_url"] = openai_base_url
+        # Initialize OpenAI client with custom URL if provided
+        client_args = {}
+        api_key_val = (
+            openai_api_key if openai_api_key is not None else CFG.RAG_EMBEDDING_API_KEY
+        )
+        base_url_val = (
+            openai_base_url
+            if openai_base_url is not None
+            else CFG.RAG_EMBEDDING_BASE_URL
+        )
+        embedding_model_val = (
+            openai_embedding_model
+            if openai_embedding_model is not None
+            else CFG.RAG_EMBEDDING_MODEL
+        )
+        chunk_size_val = chunk_size if chunk_size is not None else CFG.RAG_CHUNK_SIZE
+        overlap_val = overlap if overlap is not None else CFG.RAG_OVERLAP
+        max_result_count_val = (
+            max_result_count
+            if max_result_count is not None
+            else CFG.RAG_MAX_RESULT_COUNT
+        )
+
+        if api_key_val:
+            client_args["api_key"] = api_key_val
+        if base_url_val:
+            client_args["base_url"] = base_url_val
         # Initialize OpenAI client for embeddings
         openai_client = OpenAI(**client_args)
         # Initialize ChromaDB client
@@ -101,8 +117,8 @@ def create_rag_from_directory(
                     collection.delete(where={"file_path": relative_path})
                     content = _read_txt_content(file_path, file_reader)
                     file_id = ulid.new().str
-                    for i in range(0, len(content), chunk_size - overlap):
-                        chunk = content[i : i + chunk_size]
+                    for i in range(0, len(content), chunk_size_val - overlap_val):
+                        chunk = content[i : i + chunk_size_val]
                         if chunk:
                             chunk_id = ulid.new().str
                             print(
@@ -113,7 +129,7 @@ def create_rag_from_directory(
                             )
                             # Get embeddings using OpenAI
                             embedding_response = openai_client.embeddings.create(
-                                input=chunk, model=openai_embedding_model
+                                input=chunk, model=embedding_model_val
                             )
                             vector = embedding_response.data[0].embedding
                             collection.upsert(
@@ -140,13 +156,13 @@ def create_rag_from_directory(
         print(stylize_faint("Vectorizing query"), file=sys.stderr)
         # Get embeddings using OpenAI
         embedding_response = openai_client.embeddings.create(
-            input=query, model=openai_embedding_model
+            input=query, model=embedding_model_val
         )
         query_vector = embedding_response.data[0].embedding
         print(stylize_faint("Searching documents"), file=sys.stderr)
         results = collection.query(
             query_embeddings=query_vector,
-            n_results=max_result_count,
+            n_results=max_result_count_val,
         )
         return json.dumps(results)
 

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/module/gateway/view/static/default/script.js

@@ -41,4 +41,4 @@ function updateAutoTheme() {
 }
 
 updateAutoTheme();
-window.matchMedia('(prefers-color-scheme: dark)').addListener(updateAutoTheme);
+window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', updateAutoTheme);