zrb 1.13.1__py3-none-any.whl → 1.21.33__py3-none-any.whl

zrb/builtin/llm/chat_trigger.py

@@ -0,0 +1,78 @@
+import asyncio
+import os
+from asyncio import StreamReader
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
+
+from zrb.builtin.llm.chat_completion import get_chat_completer
+from zrb.context.any_context import AnyContext
+from zrb.util.run import run_async
+
+if TYPE_CHECKING:
+    from prompt_toolkit import PromptSession
+
+
+ChatTrigger = Callable[[AnyContext], Coroutine[Any, Any, str] | str]
+
+
+class LLMChatTrigger:
+
+    def __init__(self):
+        self._triggers: list[ChatTrigger] = []
+
+    def add_trigger(self, *trigger: ChatTrigger):
+        self.append_trigger(*trigger)
+
+    def append_trigger(self, *trigger: ChatTrigger):
+        for single_trigger in trigger:
+            self._triggers.append(single_trigger)
+
+    async def wait(
+        self, reader: "PromptSession[Any] | StreamReader", ctx: AnyContext
+    ) -> str:
+        trigger_tasks = [
+            asyncio.create_task(run_async(self._read_next_line(reader, ctx)))
+        ] + [asyncio.create_task(run_async(trigger(ctx))) for trigger in self._triggers]
+        final_result: str = ""
+        try:
+            done, pending = await asyncio.wait(
+                trigger_tasks, return_when=asyncio.FIRST_COMPLETED
+            )
+            for task in done:
+                final_result = await task
+            if pending:
+                for task in pending:
+                    task.cancel()
+            for task in done:
+                break
+        except asyncio.CancelledError:
+            ctx.print("Task cancelled.", plain=True)
+            final_result = "/bye"
+        except KeyboardInterrupt:
+            ctx.print("KeyboardInterrupt detected. Exiting...", plain=True)
+            final_result = "/bye"
+        return final_result
+
+    async def _read_next_line(
+        self, reader: "PromptSession[Any] | StreamReader", ctx: AnyContext
+    ) -> str:
+        """Reads one line of input using the provided reader."""
+        from prompt_toolkit import PromptSession
+
+        try:
+            if isinstance(reader, PromptSession):
+                bottom_toolbar = f"📁 Current directory: {os.getcwd()}"
+                return await reader.prompt_async(
+                    completer=get_chat_completer(), bottom_toolbar=bottom_toolbar
+                )
+            line_bytes = await reader.readline()
+            if not line_bytes:
+                return "/bye"  # Signal to exit
+            user_input = line_bytes.decode().strip()
+            ctx.print(user_input, plain=True)
+            return user_input
+        except KeyboardInterrupt:
+            ctx.print("KeyboardInterrupt detected. Exiting...", plain=True)
+            return "/bye"
+
+
+llm_chat_trigger = LLMChatTrigger()

zrb/builtin/llm/history.py

@@ -3,12 +3,12 @@ import os
 from typing import Any
 
 from zrb.config.config import CFG
-from zrb.context.any_shared_context import AnySharedContext
+from zrb.context.any_context import AnyContext
 from zrb.task.llm.conversation_history_model import ConversationHistory
 from zrb.util.file import read_file, write_file
 
 
-def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | None:
+def read_chat_conversation(ctx: AnyContext) -> dict[str, Any] | list | None:
     """Reads conversation history from the session file.
     Returns the raw dictionary or list loaded from JSON, or None if not found/empty.
     The LLMTask will handle parsing this into ConversationHistory.
@@ -51,10 +51,10 @@ def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | Non
         return None
 
 
-def write_chat_conversation(ctx: AnySharedContext, history_data: ConversationHistory):
+def write_chat_conversation(ctx: AnyContext, history_data: ConversationHistory):
     """Writes the conversation history data (including context) to a session file."""
     os.makedirs(CFG.LLM_HISTORY_DIR, exist_ok=True)
-    current_session_name = ctx.session.name
+    current_session_name = ctx.session.name if ctx.session is not None else None
     if not current_session_name:
         ctx.log_warning("Cannot write history: Session name is empty.")
         return

zrb/builtin/llm/llm_ask.py

@@ -1,127 +1,257 @@
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
 from zrb.builtin.group import llm_group
+from zrb.builtin.llm.attachment import get_media_type
 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.api import (
+    create_get_current_location,
+    create_get_current_weather,
+)
 from zrb.builtin.llm.tool.cli import run_shell_command
 from zrb.builtin.llm.tool.code import analyze_repo
 from zrb.builtin.llm.tool.file import (
     analyze_file,
     list_files,
     read_from_file,
-    read_many_files,
     replace_in_file,
     search_files,
-    write_many_files,
     write_to_file,
 )
+from zrb.builtin.llm.tool.note import (
+    read_contextual_note,
+    read_long_term_note,
+    write_contextual_note,
+    write_long_term_note,
+)
 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.config import CFG
+from zrb.config.llm_config import llm_config
+from zrb.context.any_context import AnyContext
+from zrb.input.any_input import AnyInput
 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
+from zrb.util.string.conversion import to_boolean
 
-_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,
-    ),
-    TextInput(
-        "modes",
-        description="Modes",
-        prompt="Modes",
-        default="coding",
-        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
+if TYPE_CHECKING:
+    from pydantic_ai import AbstractToolset, Tool, UserContent
+
+    ToolOrCallable = Tool | Callable
+
+
+def _get_toolset(ctx: AnyContext) -> list["AbstractToolset[None] | str"]:
+    cwd = os.getcwd()
+    toolsets = []
+    for config_path in [
+        os.path.join(cwd, "mcp_config.json"),
+        os.path.join(cwd, "mcp-config.json"),
+    ]:
+        if os.path.isfile(config_path):
+            toolsets.append(config_path)
+    return toolsets
+
+
+def _get_tool(ctx: AnyContext) -> list["ToolOrCallable"]:
+    tools = [
+        read_long_term_note,
+        write_long_term_note,
+        read_contextual_note,
+        write_contextual_note,
+    ]
+    if CFG.LLM_ALLOW_ANALYZE_REPO:
+        tools.append(analyze_repo)
+    if CFG.LLM_ALLOW_ANALYZE_FILE:
+        tools.append(analyze_file)
+    if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
+        tools.append(search_files)
+        tools.append(list_files)
+        tools.append(read_from_file)
+        tools.append(replace_in_file)
+        tools.append(write_to_file)
+    if CFG.LLM_ALLOW_ACCESS_SHELL:
+        tools.append(run_shell_command)
+    if CFG.LLM_ALLOW_OPEN_WEB_PAGE:
+        tools.append(open_web_page)
+    if CFG.LLM_ALLOW_GET_CURRENT_LOCATION:
+        tools.append(create_get_current_location())
+    if CFG.LLM_ALLOW_GET_CURRENT_WEATHER:
+        tools.append(create_get_current_weather())
+    if CFG.SERPAPI_KEY != "" and CFG.LLM_ALLOW_SEARCH_INTERNET:
+        tools.append(create_search_internet_tool())
+    return tools
+
+
+def _get_default_yolo_mode(ctx: AnyContext) -> str:
+    default_value = llm_config.default_yolo_mode
+    if isinstance(default_value, list):
+        return ",".join(default_value)
+    return f"{default_value}"
+
+
+def _render_yolo_mode_input(ctx: AnyContext) -> list[str] | bool:
+    if ctx.input.yolo.strip() == "":
+        return []
+    elements = [element.strip() for element in ctx.input.yolo.split(",")]
+    if len(elements) == 1:
+        try:
+            return to_boolean(elements[0])
+        except Exception:
+            pass
+    return elements
+
+
+def _render_attach_input(ctx: AnyContext) -> "list[UserContent]":
+    from pathlib import Path
+
+    from pydantic_ai import BinaryContent
+
+    attachment_paths: list[str] = [
+        attachment_path.strip()
+        for attachment_path in ctx.input.attach.split(",")
+        if attachment_path.strip() != ""
+    ]
+    if len(attachment_paths) == 0:
+        return []
+    attachments = []
+    for attachment_path in attachment_paths:
+        attachment_path = os.path.abspath(os.path.expanduser(attachment_path))
+        media_type = get_media_type(attachment_path)
+        if media_type is None:
+            ctx.log_error(f"Cannot get media type of {attachment_path}")
+            continue
+        data = Path(attachment_path).read_bytes()
+        attachments.append(BinaryContent(data, media_type=media_type))
+    return attachments
+
+
+def _get_inputs(require_message: bool = True) -> list[AnyInput | None]:
+    return [
+        StrInput(
+            "model",
+            description="LLM Model",
+            prompt="LLM Model",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        model_api_key=lambda ctx: (
-            None if ctx.input.api_key.strip() == "" else ctx.input.api_key
+        StrInput(
+            "base-url",
+            description="LLM API Base URL",
+            prompt="LLM API Base URL",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        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
+        StrInput(
+            "api-key",
+            description="LLM API Key",
+            prompt="LLM API Key",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        modes=lambda ctx: (
-            None if ctx.input.modes.strip() == "" else ctx.input.modes.split(",")
+        TextInput(
+            "system-prompt",
+            description="System prompt",
+            prompt="System prompt",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
         ),
-        message="{ctx.input.message}",
-        retries=0,
+        TextInput(
+            "workflows",
+            description="Workflows",
+            prompt="Workflows",
+            default=lambda ctx: ",".join(llm_config.default_workflows),
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        BoolInput(
+            "start-new",
+            description="Start new session (LLM Agent will forget past conversation)",
+            prompt="Start new session (LLM Agent will forget past conversation)",
+            default=False,
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        StrInput(
+            "yolo",
+            description="YOLO mode (LLM Agent will start in YOLO Mode)",
+            prompt="YOLO mode (LLM Agent will start in YOLO Mode)",
+            default=_get_default_yolo_mode,
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        TextInput(
+            "message",
+            description="User message",
+            prompt="Your message",
+            always_prompt=require_message,
+            allow_empty=not require_message,
+        ),
+        StrInput(
+            name="attach",
+            description="Comma separated attachments",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        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(
+    name="llm-ask",
+    input=_get_inputs(True),
+    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
     ),
-    alias="ask",
+    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
+    ),
+    workflows=lambda ctx: (
+        None if ctx.input.workflows.strip() == "" else ctx.input.workflows.split(",")
+    ),
+    attachment=_render_attach_input,
+    message="{ctx.input.message}",
+    tools=_get_tool,
+    toolsets=_get_toolset,
+    yolo_mode=_render_yolo_mode_input,
+    retries=0,
 )
+llm_group.add_task(llm_ask, alias="ask")
 
 llm_group.add_task(
     BaseTrigger(
         name="llm-chat",
-        input=_llm_ask_inputs,
+        input=_get_inputs(False),
         description="💬 Chat with LLM",
         queue_name="ask_trigger",
         action=read_user_prompt,
@@ -137,25 +267,3 @@ llm_group.add_task(
     ),
     alias="chat",
 )
-
-if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
-    llm_ask.append_tool(
-        analyze_repo,
-        analyze_file,
-        search_files,
-        list_files,
-        read_from_file,
-        read_many_files,
-        replace_in_file,
-        write_to_file,
-        write_many_files,
-    )
-
-if CFG.LLM_ALLOW_ACCESS_SHELL:
-    llm_ask.append_tool(run_shell_command)
-
-if CFG.LLM_ALLOW_ACCESS_INTERNET:
-    llm_ask.append_tool(open_web_page, search_wikipedia, search_arxiv)
-    if CFG.SERPAPI_KEY != "":
-        llm_ask.append_tool(create_search_internet_tool(CFG.SERPAPI_KEY))
-    llm_ask.append_tool(get_current_location, get_current_weather)

zrb/builtin/llm/tool/api.py

@@ -1,63 +1,75 @@
-import json
-from typing import Literal
-
-
-def get_current_location() -> str:
-    """
-    Fetches the user's current geographical location (latitude and longitude) based on their IP address.
-
-    Use this tool when the user asks "Where am I?", "What is my current location?", or has a query that requires knowing their location to be answered.
-
-    Returns:
-        str: A JSON string containing the 'lat' and 'lon' of the current location.
-             Example: '{"lat": 48.8584, "lon": 2.2945}'
-    Raises:
-        requests.RequestException: If the API request to the location service fails.
-    """
-    import requests
-
-    try:
-        response = requests.get("http://ip-api.com/json?fields=lat,lon", timeout=5)
-        response.raise_for_status()
-        return json.dumps(response.json())
-    except requests.RequestException as e:
-        raise requests.RequestException(f"Failed to get location: {e}") from None
-
-
-def get_current_weather(
-    latitude: float,
-    longitude: float,
-    temperature_unit: Literal["celsius", "fahrenheit"],
-) -> str:
-    """
-    Retrieves the current weather conditions for a given geographical location.
-
-    Use this tool when the user asks about the weather. If the user does not provide a location, first use the `get_current_location` tool to determine their location.
-
-    Args:
-        latitude (float): The latitude of the location.
-        longitude (float): The longitude of the location.
-        temperature_unit (Literal["celsius", "fahrenheit"]): The desired unit for the temperature reading.
-
-    Returns:
-        str: A JSON string containing detailed weather data, including temperature, wind speed, and weather code.
-    Raises:
-        requests.RequestException: If the API request to the weather service fails.
-    """
-    import requests
-
-    try:
-        response = requests.get(
-            "https://api.open-meteo.com/v1/forecast",
-            params={
-                "latitude": latitude,
-                "longitude": longitude,
-                "temperature_unit": temperature_unit,
-                "current_weather": True,
-            },
-            timeout=5,
+from typing import Any, Callable, Literal
+
+from zrb.config.llm_config import llm_config
+
+
+def create_get_current_location() -> Callable:
+    if llm_config.default_current_location_tool is not None:
+        return llm_config.default_current_location_tool
+
+    def get_current_location() -> dict[str, float]:
+        """
+        Gets the user's current geographical location (latitude and longitude).
+
+        Example:
+        get_current_location() # Returns {'lat': 48.8584, 'lon': 2.2945}
+
+        Returns:
+            dict: Dictionary with 'lat' and 'lon' keys.
+        """
+        import requests
+
+        try:
+            response = requests.get("http://ip-api.com/json?fields=lat,lon", timeout=5)
+            response.raise_for_status()
+            return dict(response.json())
+        except requests.RequestException as e:
+            raise requests.RequestException(f"Failed to get location: {e}")
+
+    return get_current_location
+
+
+def create_get_current_weather() -> Callable:
+    if llm_config.default_current_weather_tool is not None:
+        return llm_config.default_current_weather_tool
+
+    def get_current_weather(
+        latitude: float,
+        longitude: float,
+        temperature_unit: Literal["celsius", "fahrenheit"],
+    ) -> dict[str, Any]:
+        """
+        Gets current weather conditions for a given location.
+
+        Example:
+        get_current_weather(
+            latitude=34.0522, longitude=-118.2437, temperature_unit='fahrenheit'
         )
-        response.raise_for_status()
-        return json.dumps(response.json())
-    except requests.RequestException as e:
-        raise requests.RequestException(f"Failed to get weather data: {e}") from None
+
+        Args:
+            latitude (float): Latitude of the location.
+            longitude (float): Longitude of the location.
+            temperature_unit (Literal): 'celsius' or 'fahrenheit'.
+
+        Returns:
+            dict: Detailed weather data.
+        """
+        import requests
+
+        try:
+            response = requests.get(
+                "https://api.open-meteo.com/v1/forecast",
+                params={
+                    "latitude": latitude,
+                    "longitude": longitude,
+                    "temperature_unit": temperature_unit,
+                    "current_weather": True,
+                },
+                timeout=5,
+            )
+            response.raise_for_status()
+            return dict(response.json())
+        except requests.RequestException as e:
+            raise requests.RequestException(f"Failed to get weather data: {e}")
+
+    return get_current_weather