zrb 1.0.0b4__tar.gz → 1.0.0b6__tar.gz

{zrb-1.0.0b4 → zrb-1.0.0b6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0b4
+Version: 1.0.0b6
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -19,10 +19,11 @@ 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.6,<0.116.0)
+Requires-Dist: fastembed (>=0.5.1,<0.6.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: pydantic-ai (>=0.0.19,<0.0.20)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.3.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)

{zrb-1.0.0b4 → zrb-1.0.0b6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.0.0b4"
+version = "1.0.0b6"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -44,13 +44,14 @@ black = "~24.10.0"
 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.6"}
 python-jose = {extras = ["cryptography"], version = "^3.3.0"}
 ulid-py = "^1.1.0"
+pydantic-ai = "^0.0.19"
+fastembed = "^0.5.1"
 
 [tool.poetry.extras]
 # poetry install -E rag

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/builtin/llm/llm_chat.py

@@ -92,8 +92,8 @@ llm_chat: LLMTask = llm_group.add_task(
             ),
             BoolInput(
                 "start-new",
-                description="Start new conversation session",
-                prompt="Forget everything and start new conversation session",
+                description="Start new conversation (LLM will forget everything)",
+                prompt="Start new conversation (LLM will forget everything)",
                 default_str="false",
                 allow_positional_parsing=False,
             ),

zrb-1.0.0b6/src/zrb/builtin/llm/previous-session.js

@@ -0,0 +1,21 @@
+async function updatePreviousSession(event) {
+    const currentInput = event.target;
+    if (currentInput.name === "CURRENT_INPUT_NAME") {
+        return
+    }
+    const previousSessionInput = submitTaskForm.querySelector('[name="CURRENT_INPUT_NAME"]');
+    if (previousSessionInput) {
+        const currentSessionName = cfg.SESSION_NAME
+        previousSessionInput.value = currentSessionName;
+    }
+}
+
+document.getElementById("submit-task-form").querySelectorAll("input[name], textarea[name]").forEach((element) => {
+    element.addEventListener("input", updatePreviousSession);
+    element.addEventListener("keyup", updatePreviousSession);
+});
+
+document.getElementById("submit-task-form").querySelectorAll("select[name]").forEach((element) => {
+    element.addEventListener("change", updatePreviousSession);
+});
+

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/builtin/llm/tool/rag.py

@@ -3,7 +3,6 @@ import json
 import os
 import sys
 
-import litellm
 import ulid
 
 from zrb.config import (
@@ -30,7 +29,9 @@ def create_rag_from_directory(
     async def retrieve(query: str) -> str:
         from chromadb import PersistentClient
         from chromadb.config import Settings
+        from fastembed import TextEmbedding
 
+        embedding_model = TextEmbedding(model_name=model)
         client = PersistentClient(
             path=vector_db_path, settings=Settings(allow_reset=True)
         )
@@ -75,10 +76,8 @@ def create_rag_from_directory(
                                 ),
                                 file=sys.stderr,
                             )
-                            response = await litellm.aembedding(
-                                model=model, input=[chunk]
-                            )
-                            vector = response["data"][0]["embedding"]
+                            embedding_result = list(embedding_model.embed([chunk]))
+                            vector = embedding_result[0]
                             collection.upsert(
                                 ids=[chunk_id],
                                 embeddings=[vector],
@@ -102,9 +101,8 @@ def create_rag_from_directory(
             )
 
         print(stylize_faint("Vectorizing query"), file=sys.stderr)
-        query_response = await litellm.aembedding(model=model, input=[query])
-        query_vector = query_response["data"][0]["embedding"]
-
+        embedding_result = list(embedding_model.embed([query]))
+        query_vector = embedding_result[0]
         print(stylize_faint("Searching documents"), file=sys.stderr)
         results = collection.query(
             query_embeddings=query_vector,

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/config.py

@@ -78,22 +78,14 @@ WEB_AUTH_REFRESH_TOKEN_EXPIRE_MINUTES = int(
 )
 LLM_MODEL = os.getenv("ZRB_LLM_MODEL", "ollama_chat/llama3.1")
 
-_DEFAULT_PROMPT = """
-You are a helpful assistant and you have access to several tools.
-Your goal is to provide a final answer by executing a series of planning, actions, reasoning, and evaluations.
-
-Breakdown user request into several actionable tasks. For example, when user ask about current weather on current location, you should get the current location first.
-
-DO NOT TRY TO SIMULATE TOOL OUTPUT.
-
-ERROR HANDLING
-1. If you receive an error, read the error message carefully and identify the specific issue.
-2. Adjust your response accordingly and perform the review process again before resubmitting.
-
-REMINDER:
-- ALWAYS double-check your response format and function arguments before submitting.
-- DON'T make up answers.
-""".strip()
+_DEFAULT_PROMPT = (
+    "You are a helpful AI assistant capable of using various tools to answer user queries. When solving a problem:\n"
+    "1. Carefully analyze the user's request and identify what information is needed to provide a complete answer.\n"
+    "2. Determine which available tools can help you gather the necessary information.\n"
+    "3. Call tools strategically and in a logical sequence to collect required data.\n"
+    "4. If a tool provides incomplete information, intelligently decide which additional tool or approach to use.\n"
+    "5. Always aim to provide the most accurate and helpful response possible."
+)
 LLM_SYSTEM_PROMPT = os.getenv("ZRB_LLM_SYSTEM_PROMPT", _DEFAULT_PROMPT)
 LLM_HISTORY_DIR = os.getenv(
     "ZRB_LLM_HISTORY_DIR", os.path.expanduser(os.path.join("~", ".zrb-llm-history"))
@@ -103,7 +95,10 @@ LLM_HISTORY_FILE = os.getenv(
 )
 LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_FILE", "1"))
 LLM_ALLOW_ACCESS_INTERNET = to_boolean(os.getenv("ZRB_LLM_ACCESS_INTERNET", "1"))
-RAG_EMBEDDING_MODEL = os.getenv("ZRB_RAG_EMBEDDING_MODEL", "ollama/nomic-embed-text")
+# noqa See: https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models
+RAG_EMBEDDING_MODEL = os.getenv(
+    "ZRB_RAG_EMBEDDING_MODEL", "nomic-ai/nomic-embed-text-v1.5-Q"
+)
 RAG_CHUNK_SIZE = int(os.getenv("ZRB_RAG_CHUNK_SIZE", "1024"))
 RAG_OVERLAP = int(os.getenv("ZRB_RAG_OVERLAP", "128"))
 RAG_MAX_RESULT_COUNT = int(os.getenv("ZRB_RAG_MAX_RESULT_COUNT", "5"))

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/runner/web_route/node_page/task/view.html

@@ -63,7 +63,7 @@
             <article>
                 <form id="submit-task-form" onsubmit="submitNewSessionForm(event)">
                     {task_inputs}
-                    <button>🚀 Run</button>
+                    <button>🚀 Run New Session</button>
                 </form>
             </article>
             <article>

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/runner/web_route/static/resources/session/current-session.js

@@ -4,18 +4,22 @@ const CURRENT_SESSION = {
         const logTextarea = document.getElementById("log-textarea");
         const submitTaskForm = document.getElementById("submit-task-form");
         let isFinished = false;
+        let isInputUpdated = false;
         let errorCount = 0;
         while (!isFinished) {
             try {
                 const data = await this.getCurrentSession();
                 // update inputs
-                const dataInputs = data.input;
-                for (const inputName in dataInputs) {
-                    const inputValue = dataInputs[inputName];
-                    const input = submitTaskForm.querySelector(`[name="${inputName}"]`);
-                    if (input) {
-                        input.value = inputValue;
+                if (!isInputUpdated) {
+                    const dataInputs = data.input;
+                    for (const inputName in dataInputs) {
+                        const inputValue = dataInputs[inputName];
+                        const input = submitTaskForm.querySelector(`[name="${inputName}"]`);
+                        if (input) {
+                            input.value = inputValue;
+                        }
                     }
+                    isInputUpdated = true;
                 }
                 resultLineCount = data.final_result.split("\n").length;
                 resultTextarea.rows = resultLineCount <= 5 ? resultLineCount : 5;

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/runner/web_route/static/resources/session/event.js

@@ -20,7 +20,7 @@ window.addEventListener("load", async function () {
 
 
 const submitTaskForm = document.getElementById("submit-task-form");
-submitTaskForm.addEventListener("change", async function(event) {
+async function handleInputUpdate(event) {
     const currentInput = event.target;
     const inputs = Array.from(submitTaskForm.querySelectorAll("input[name], textarea[name], select[name]"));
     const inputMap = {};
@@ -59,6 +59,7 @@ submitTaskForm.addEventListener("change", async function(event) {
                 if (value === "") {
                     return;
                 }
+                console.log(input, data);
                 input.value = value;
             });
         } else {
@@ -67,6 +68,14 @@ submitTaskForm.addEventListener("change", async function(event) {
     } catch (error) {
         console.error("Error during fetch:", error);
     }
+}
+
+submitTaskForm.querySelectorAll("input[name], textarea[name]").forEach((element) => {
+    element.addEventListener("input", handleInputUpdate);
+    element.addEventListener("keyup", handleInputUpdate);
+});
+submitTaskForm.querySelectorAll("select[name]").forEach((element) => {
+    element.addEventListener("change", handleInputUpdate);
 });
 
 

zrb-1.0.0b6/src/zrb/task/llm_task.py

@@ -0,0 +1,199 @@
+import json
+import os
+from collections.abc import Callable
+from typing import Any
+
+from pydantic_ai import Agent, Tool
+from pydantic_ai.messages import ModelMessagesTypeAdapter
+from pydantic_ai.settings import ModelSettings
+
+from zrb.attr.type import StrAttr
+from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.util.attr import get_str_attr
+from zrb.util.cli.style import stylize_faint
+from zrb.util.file import read_file, write_file
+from zrb.util.run import run_async
+
+ListOfDict = list[dict[str, Any]]
+ToolOrCallable = Tool | Callable
+
+
+class LLMTask(BaseTask):
+    def __init__(
+        self,
+        name: str,
+        color: int | None = None,
+        icon: str | None = None,
+        description: str | None = None,
+        cli_only: bool = False,
+        input: list[AnyInput | None] | AnyInput | None = None,
+        env: list[AnyEnv | None] | AnyEnv | None = None,
+        model: StrAttr | None = LLM_MODEL,
+        model_settings: (
+            ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
+        ) = None,
+        render_model: bool = True,
+        agent: Agent | Callable[[AnySharedContext], Agent] | None = None,
+        system_prompt: StrAttr | None = LLM_SYSTEM_PROMPT,
+        render_system_prompt: bool = True,
+        message: StrAttr | None = None,
+        tools: (
+            list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
+        ) = [],
+        conversation_history: (
+            ListOfDict | Callable[[AnySharedContext], ListOfDict]
+        ) = [],
+        conversation_history_reader: (
+            Callable[[AnySharedContext], ListOfDict] | None
+        ) = None,
+        conversation_history_writer: (
+            Callable[[AnySharedContext, ListOfDict], None] | None
+        ) = None,
+        conversation_history_file: StrAttr | None = None,
+        render_history_file: bool = True,
+        execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+        retries: int = 2,
+        retry_period: float = 0,
+        readiness_check: list[AnyTask] | AnyTask | None = None,
+        readiness_check_delay: float = 0.5,
+        readiness_check_period: float = 5,
+        readiness_failure_threshold: int = 1,
+        readiness_timeout: int = 60,
+        monitor_readiness: bool = False,
+        max_call_iteration: int = 20,
+        upstream: list[AnyTask] | AnyTask | None = None,
+        fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
+    ):
+        super().__init__(
+            name=name,
+            color=color,
+            icon=icon,
+            description=description,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            execute_condition=execute_condition,
+            retries=retries,
+            retry_period=retry_period,
+            readiness_check=readiness_check,
+            readiness_check_delay=readiness_check_delay,
+            readiness_check_period=readiness_check_period,
+            readiness_failure_threshold=readiness_failure_threshold,
+            readiness_timeout=readiness_timeout,
+            monitor_readiness=monitor_readiness,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+        )
+        self._model = model
+        self._model_settings = (model_settings,)
+        self._agent = agent
+        self._render_model = render_model
+        self._system_prompt = system_prompt
+        self._render_system_prompt = render_system_prompt
+        self._message = message
+        self._tools = tools
+        self._additional_tools: list[ToolOrCallable] = []
+        self._conversation_history = conversation_history
+        self._conversation_history_reader = conversation_history_reader
+        self._conversation_history_writer = conversation_history_writer
+        self._conversation_history_file = conversation_history_file
+        self._render_history_file = render_history_file
+        self._max_call_iteration = max_call_iteration
+
+    def add_tool(self, tool: ToolOrCallable):
+        self._additional_tools.append(tool)
+
+    async def _exec_action(self, ctx: AnyContext) -> Any:
+        history = await self._read_conversation_history(ctx)
+        user_prompt = self._get_message(ctx)
+        agent = self._get_agent(ctx)
+        result = await agent.run(
+            user_prompt=user_prompt,
+            message_history=ModelMessagesTypeAdapter.validate_python(history),
+        )
+        new_history = json.loads(result.all_messages_json())
+        for history in new_history:
+            ctx.print(stylize_faint(json.dumps(history)))
+        await self._write_conversation_history(ctx, new_history)
+        return result.data
+
+    async def _write_conversation_history(
+        self, ctx: AnyContext, conversations: list[Any]
+    ):
+        if self._conversation_history_writer is not None:
+            await run_async(self._conversation_history_writer(ctx, conversations))
+        history_file = self._get_history_file(ctx)
+        if history_file != "":
+            write_file(history_file, json.dumps(conversations, indent=2))
+
+    def _get_model_settings(self, ctx: AnyContext) -> ModelSettings | None:
+        if isinstance(self._model_settings, ModelSettings):
+            return self._model_settings
+        if callable(self._model_settings):
+            return self._model_settings(ctx)
+        return None
+
+    def _get_agent(self, ctx: AnyContext) -> Agent:
+        if isinstance(self._agent, Agent):
+            return self._agent
+        if callable(self._agent):
+            return self._agent(ctx)
+        tools_or_callables = list(
+            self._tools(ctx) if callable(self._tools) else self._tools
+        )
+        tools_or_callables.extend(self._additional_tools)
+        tools = [
+            tool if isinstance(tool, Tool) else Tool(tool, takes_ctx=False)
+            for tool in tools_or_callables
+        ]
+        return Agent(
+            self._get_model(ctx),
+            system_prompt=self._get_system_prompt(ctx),
+            tools=tools,
+        )
+
+    def _get_model(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx, self._model, "ollama_chat/llama3.1", auto_render=self._render_model
+        )
+
+    def _get_system_prompt(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._system_prompt,
+            "You are a helpful assistant",
+            auto_render=self._render_system_prompt,
+        )
+
+    def _get_message(self, ctx: AnyContext) -> str:
+        return get_str_attr(ctx, self._message, "How are you?", auto_render=True)
+
+    async def _read_conversation_history(self, ctx: AnyContext) -> ListOfDict:
+        if self._conversation_history_reader is not None:
+            return await run_async(self._conversation_history_reader(ctx))
+        if callable(self._conversation_history):
+            return self._conversation_history(ctx)
+        history_file = self._get_history_file(ctx)
+        if (
+            len(self._conversation_history) == 0
+            and history_file != ""
+            and os.path.isfile(history_file)
+        ):
+            return json.loads(read_file(history_file))
+        return self._conversation_history
+
+    def _get_history_file(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._conversation_history_file,
+            "",
+            auto_render=self._render_history_file,
+        )

{zrb-1.0.0b4 → zrb-1.0.0b6}/src/zrb/util/llm/tool.py

@@ -1,6 +1,6 @@
 import inspect
 from collections.abc import Callable
-from typing import Any, get_type_hints
+from typing import Annotated, Any, Literal, get_type_hints
 
 
 def callable_to_tool_schema(callable_obj: Callable) -> dict[str, Any]:
@@ -21,10 +21,14 @@ def callable_to_tool_schema(callable_obj: Callable) -> dict[str, Any]:
     # Build parameter schema
     param_schema = {"type": "object", "properties": {}, "required": []}
     for param_name, param in sig.parameters.items():
-        param_type = hints.get(param_name, str)  # Default type is string
-        param_schema["properties"][param_name] = {
-            "type": _python_type_to_json_type(param_type)
-        }
+        # Get the type hint or default to str
+        param_type = hints.get(param_name, str)
+
+        # Handle annotated types (e.g., Annotated[str, "description"])
+        json_type, param_metadata = _process_type_annotation(param_type)
+        param_schema["properties"][param_name] = param_metadata
+
+        # Mark required parameters
         if param.default is inspect.Parameter.empty:
             param_schema["required"].append(param_name)
     return {
@@ -37,6 +41,30 @@ def callable_to_tool_schema(callable_obj: Callable) -> dict[str, Any]:
     }
 
 
+def _process_type_annotation(py_type: Any) -> tuple[str, dict]:
+    """
+    Process type annotations and return the JSON Schema type and metadata.
+
+    :param py_type: The type annotation.
+    :return: A tuple of (JSON type, parameter metadata).
+    """
+    if hasattr(py_type, "__origin__") and py_type.__origin__ is Literal:
+        # Handle Literal (enum)
+        enum_values = list(py_type.__args__)
+        return "string", {"type": "string", "enum": enum_values}
+
+    if hasattr(py_type, "__origin__") and py_type.__origin__ is Annotated:
+        # Handle Annotated types
+        base_type = py_type.__args__[0]
+        description = py_type.__args__[1]
+        json_type = _python_type_to_json_type(base_type)
+        return json_type, {"type": json_type, "description": description}
+
+    # Fallback to basic type conversion
+    json_type = _python_type_to_json_type(py_type)
+    return json_type, {"type": json_type}
+
+
 def _python_type_to_json_type(py_type):
     """
     Map Python types to JSON Schema types.

zrb-1.0.0b4/src/zrb/builtin/llm/previous-session.js

@@ -1,13 +0,0 @@
-let hasUpdateCurrentPascalInputName = false;
-document.getElementById("submit-task-form").addEventListener("change", async function(event) {
-    const currentInput = event.target;
-    if (hasUpdateCurrentPascalInputName || currentInput.name === "CURRENT_INPUT_NAME") {
-        return
-    }
-    const previousSessionInput = submitTaskForm.querySelector('[name="CURRENT_INPUT_NAME"]');
-    if (previousSessionInput) {
-        const currentSessionName = cfg.SESSION_NAME
-        previousSessionInput.value = currentSessionName;
-    }
-    hasUpdateCurrentPascalInputName = true;
-});