zrb 1.2.1__py3-none-any.whl → 1.2.2__py3-none-any.whl

zrb/builtin/llm/llm_chat.py

@@ -5,13 +5,26 @@ from typing import Any
 from zrb.builtin.group import llm_group
 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.web import open_web_route, query_internet
+from zrb.builtin.llm.tool.file import (
+    list_file,
+    read_source_code,
+    read_text_file,
+    write_text_file,
+)
+from zrb.builtin.llm.tool.web import (
+    create_search_internet_tool,
+    open_web_page,
+    search_arxiv,
+    search_wikipedia,
+)
 from zrb.config import (
     LLM_ALLOW_ACCESS_INTERNET,
+    LLM_ALLOW_ACCESS_LOCAL_FILE,
     LLM_ALLOW_ACCESS_SHELL,
     LLM_HISTORY_DIR,
     LLM_MODEL,
     LLM_SYSTEM_PROMPT,
+    SERP_API_KEY,
 )
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.input.bool_input import BoolInput
@@ -117,11 +130,21 @@ llm_chat: LLMTask = llm_group.add_task(
     alias="chat",
 )
 
+
+if LLM_ALLOW_ACCESS_LOCAL_FILE:
+    llm_chat.add_tool(read_source_code)
+    llm_chat.add_tool(list_file)
+    llm_chat.add_tool(read_text_file)
+    llm_chat.add_tool(write_text_file)
+
 if LLM_ALLOW_ACCESS_SHELL:
     llm_chat.add_tool(run_shell_command)
 
 if LLM_ALLOW_ACCESS_INTERNET:
-    llm_chat.add_tool(open_web_route)
-    llm_chat.add_tool(query_internet)
+    llm_chat.add_tool(open_web_page)
+    llm_chat.add_tool(search_wikipedia)
+    llm_chat.add_tool(search_arxiv)
+    if SERP_API_KEY != "":
+        llm_chat.add_tool(create_search_internet_tool(SERP_API_KEY))
     llm_chat.add_tool(get_current_location)
     llm_chat.add_tool(get_current_weather)

zrb/builtin/llm/tool/api.py

@@ -1,13 +1,13 @@
 import json
 from typing import Annotated, Literal
 
-import requests
-
 
 def get_current_location() -> (
     Annotated[str, "JSON string representing latitude and longitude"]
 ):  # noqa
     """Get the user's current location."""
+    import requests
+
     return json.dumps(requests.get("http://ip-api.com/json?fields=lat,lon").json())
 
 
@@ -17,6 +17,8 @@ def get_current_weather(
     temperature_unit: Literal["celsius", "fahrenheit"],
 ) -> str:
     """Get the current weather in a given location."""
+    import requests
+
     resp = requests.get(
         "https://api.open-meteo.com/v1/forecast",
         params={

zrb/builtin/llm/tool/file.py

@@ -0,0 +1,39 @@
+import os
+
+from zrb.util.file import read_file, write_file
+
+
+def list_file(
+    directory: str = ".",
+    extensions: list[str] = [".py", ".go", ".js", ".ts", ".java", ".c", ".cpp"],
+) -> list[str]:
+    """List all files in a directory"""
+    all_files: list[str] = []
+    for root, _, files in os.walk(directory):
+        for filename in files:
+            for extension in extensions:
+                if filename.lower().endswith(extension):
+                    all_files.append(os.path.join(root, filename))
+    return all_files
+
+
+def read_text_file(file: str) -> str:
+    """Read a text file"""
+    return read_file(os.path.abspath(file))
+
+
+def write_text_file(file: str, content: str):
+    """Write a text file"""
+    return write_file(os.path.abspath(file), content)
+
+
+def read_source_code(
+    directory: str = ".",
+    extensions: list[str] = [".py", ".go", ".js", ".ts", ".java", ".c", ".cpp"],
+) -> list[str]:
+    """Read source code in a directory"""
+    files = list_file(directory, extensions)
+    for index, file in enumerate(files):
+        content = read_text_file(file)
+        files[index] = f"# {file}\n```\n{content}\n```"
+    return files

zrb/builtin/llm/tool/rag.py

@@ -1,7 +1,9 @@
+import fnmatch
 import hashlib
 import json
 import os
 import sys
+from collections.abc import Callable
 
 import ulid
 
@@ -15,6 +17,20 @@ from zrb.util.cli.style import stylize_error, stylize_faint
 from zrb.util.file import read_file
 
 
+class RAGFileReader:
+    def __init__(self, glob_pattern: str, read: Callable[[str], str]):
+        self.glob_pattern = glob_pattern
+        self.read = read
+
+    def is_match(self, file_name: str):
+        if os.sep not in self.glob_pattern and (
+            os.altsep is None or os.altsep not in self.glob_pattern
+        ):
+            # Pattern like "*.txt" – match only the basename.
+            return fnmatch.fnmatch(os.path.basename(file_name), self.glob_pattern)
+        return fnmatch.fnmatch(file_name, self.glob_pattern)
+
+
 def create_rag_from_directory(
     tool_name: str,
     tool_description: str,
@@ -25,6 +41,7 @@ def create_rag_from_directory(
     chunk_size: int = RAG_CHUNK_SIZE,
     overlap: int = RAG_OVERLAP,
     max_result_count: int = RAG_MAX_RESULT_COUNT,
+    file_reader: list[RAGFileReader] = [],
 ):
     async def retrieve(query: str) -> str:
         from chromadb import PersistentClient
@@ -36,35 +53,31 @@ def create_rag_from_directory(
             path=vector_db_path, settings=Settings(allow_reset=True)
         )
         collection = client.get_or_create_collection(vector_db_collection)
-
         # Track file changes using a hash-based approach
         hash_file_path = os.path.join(vector_db_path, "file_hashes.json")
         previous_hashes = _load_hashes(hash_file_path)
         current_hashes = {}
-
+        # Get updated_files
         updated_files = []
-
         for root, _, files in os.walk(document_dir_path):
             for file in files:
                 file_path = os.path.join(root, file)
                 file_hash = _compute_file_hash(file_path)
                 relative_path = os.path.relpath(file_path, document_dir_path)
                 current_hashes[relative_path] = file_hash
-
                 if previous_hashes.get(relative_path) != file_hash:
                     updated_files.append(file_path)
-
+        # Upsert updated_files to vector db
         if updated_files:
             print(
                 stylize_faint(f"Updating {len(updated_files)} changed files"),
                 file=sys.stderr,
             )
-
             for file_path in updated_files:
                 try:
                     relative_path = os.path.relpath(file_path, document_dir_path)
                     collection.delete(where={"file_path": relative_path})
-                    content = _read_file_content(file_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]
@@ -92,14 +105,13 @@ def create_rag_from_directory(
                         stylize_error(f"Error processing {file_path}: {e}"),
                         file=sys.stderr,
                     )
-
             _save_hashes(hash_file_path, current_hashes)
         else:
             print(
                 stylize_faint("No changes detected. Skipping database update."),
                 file=sys.stderr,
             )
-
+        # Vectorize query and get related document chunks
         print(stylize_faint("Vectorizing query"), file=sys.stderr)
         embedding_result = list(embedding_model.embed([query]))
         query_vector = embedding_result[0]
@@ -123,7 +135,22 @@ def _compute_file_hash(file_path: str) -> str:
     return hash_md5.hexdigest()
 
 
-def _read_file_content(file_path: str) -> str:
+def _load_hashes(file_path: str) -> dict:
+    if os.path.exists(file_path):
+        with open(file_path, "r") as f:
+            return json.load(f)
+    return {}
+
+
+def _save_hashes(file_path: str, hashes: dict):
+    with open(file_path, "w") as f:
+        json.dump(hashes, f)
+
+
+def _read_txt_content(file_path: str, file_reader: list[RAGFileReader]):
+    for reader in file_reader:
+        if reader.is_match(file_path):
+            return reader.read(file_path)
     if file_path.lower().endswith(".pdf"):
         return _read_pdf(file_path)
     return read_file(file_path)
@@ -136,15 +163,3 @@ def _read_pdf(file_path: str) -> str:
         return "\n".join(
             page.extract_text() for page in pdf.pages if page.extract_text()
         )
-
-
-def _load_hashes(file_path: str) -> dict:
-    if os.path.exists(file_path):
-        with open(file_path, "r") as f:
-            return json.load(f)
-    return {}
-
-
-def _save_hashes(file_path: str, hashes: dict):
-    with open(file_path, "w") as f:
-        json.dump(hashes, f)

zrb/builtin/llm/tool/web.py

@@ -1,8 +1,9 @@
 import json
+from collections.abc import Callable
 from typing import Annotated
 
 
-def open_web_route(url: str) -> str:
+def open_web_page(url: str) -> str:
     """Get content from a web page."""
     import requests
 
@@ -19,30 +20,55 @@ def open_web_route(url: str) -> str:
     return json.dumps(parse_html_text(response.text))
 
 
-def query_internet(
+def create_search_internet_tool(serp_api_key: str) -> Callable[[str, int], str]:
+    def search_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://serpapi.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",
+                "api_key": serp_api_key,
+            },
+        )
+        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))
+
+    return search_internet
+
+
+def search_wikipedia(query: Annotated[str, "Search query"]) -> str:
+    """Search on wikipedia"""
+    import requests
+
+    params = {"action": "query", "list": "search", "srsearch": query, "format": "json"}
+    response = requests.get("https://en.wikipedia.org/w/api.php", params=params)
+    return response.json()
+
+
+def search_arxiv(
     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."""
+    """Search on Arxiv"""
     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))
+    params = {"search_query": f"all:{query}", "start": 0, "max_results": num_results}
+    response = requests.get("http://export.arxiv.org/api/query", params=params)
+    return response.content
 
 
 def parse_html_text(html_text: str) -> dict[str, str]:

zrb/config.py

@@ -92,7 +92,8 @@ LLM_HISTORY_DIR = os.getenv(
 LLM_HISTORY_FILE = os.getenv(
     "ZRB_LLM_HISTORY_FILE", os.path.join(LLM_HISTORY_DIR, "history.json")
 )
-LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_FILE", "1"))
+LLM_ALLOW_ACCESS_LOCAL_FILE = to_boolean(os.getenv("ZRB_LLM_ACCESS_LOCAL_FILE", "1"))
+LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_SHELL", "1"))
 LLM_ALLOW_ACCESS_INTERNET = to_boolean(os.getenv("ZRB_LLM_ACCESS_INTERNET", "1"))
 # noqa See: https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models
 RAG_EMBEDDING_MODEL = os.getenv(
@@ -101,6 +102,7 @@ RAG_EMBEDDING_MODEL = os.getenv(
 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"))
+SERP_API_KEY = os.getenv("SERP_API_KEY", "")
 
 
 BANNER = f"""

zrb/content_transformer/content_transformer.py

@@ -1,4 +1,5 @@
 import fnmatch
+import os
 import re
 from collections.abc import Callable
 
@@ -40,7 +41,12 @@ class ContentTransformer(AnyContentTransformer):
                     return True
             except re.error:
                 pass
-            return fnmatch.fnmatch(file_path, pattern)
+            if os.sep not in pattern and (
+                os.altsep is None or os.altsep not in pattern
+            ):
+                # Pattern like "*.txt" – match only the basename.
+                return fnmatch.fnmatch(file_path, os.path.basename(file_path))
+            return fnmatch.fnmatch(file_path, file_path)
 
     def transform_file(self, ctx: AnyContext, file_path: str):
         if callable(self._transform_file):

zrb/context/context.py

@@ -88,7 +88,7 @@ class Context(AnyContext):
             return template
         return int(self.render(template))
 
-    def render_float(self, template: str) -> float:
+    def render_float(self, template: str | float) -> float:
         if isinstance(template, float):
             return template
         return float(self.render(template))
@@ -102,9 +102,10 @@ class Context(AnyContext):
         flush: bool = True,
         plain: bool = False,
     ):
+        sep = " " if sep is None else sep
         message = sep.join([f"{value}" for value in values])
         if plain:
-            self.append_to_shared_log(remove_style(message))
+            # self.append_to_shared_log(remove_style(message))
             print(message, sep=sep, end=end, file=file, flush=flush)
             return
         color = self._color
@@ -132,6 +133,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.DEBUG:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_log(f"[DEBUG] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -145,6 +147,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.INFO:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_log(f"[INFO] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -158,6 +161,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.INFO:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_warning(f"[WARNING] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -171,6 +175,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.ERROR:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_error(f"[ERROR] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -184,6 +189,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.CRITICAL:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_error(f"[CRITICAL] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)

zrb/input/text_input.py

@@ -69,16 +69,17 @@ class TextInput(BaseInput):
         )
 
     def _prompt_cli_str(self, shared_ctx: AnySharedContext) -> str:
-        prompt_message = (
-            f"{self.comment_start}{super().prompt_message}{self.comment_end}"
+        prompt_message = super().prompt_message
+        comment_prompt_message = (
+            f"{self.comment_start}{prompt_message}{self.comment_end}"
         )
-        prompt_message_eol = f"{prompt_message}\n"
+        comment_prompt_message_eol = f"{comment_prompt_message}\n"
         default_value = self.get_default_str(shared_ctx)
         with tempfile.NamedTemporaryFile(
             delete=False, suffix=self._extension
         ) as temp_file:
             temp_file_name = temp_file.name
-            temp_file.write(prompt_message_eol.encode())
+            temp_file.write(comment_prompt_message_eol.encode())
             # Pre-fill with default content
             if default_value:
                 temp_file.write(default_value.encode())
@@ -87,7 +88,10 @@ class TextInput(BaseInput):
         subprocess.call([self._editor, temp_file_name])
         # Read the edited content
         edited_content = read_file(temp_file_name)
-        parts = [text.strip() for text in edited_content.split(prompt_message, 1)]
+        parts = [
+            text.strip() for text in edited_content.split(comment_prompt_message, 1)
+        ]
         edited_content = "\n".join(parts).lstrip()
         os.remove(temp_file_name)
+        print(f"{prompt_message}: {edited_content}")
         return edited_content

zrb/task/llm_task.py

@@ -4,10 +4,20 @@ from collections.abc import Callable
 from typing import Any
 
 from pydantic_ai import Agent, Tool
-from pydantic_ai.messages import ModelMessagesTypeAdapter
+from pydantic_ai.messages import (
+    FinalResultEvent,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    ModelMessagesTypeAdapter,
+    PartDeltaEvent,
+    PartStartEvent,
+    TextPartDelta,
+    ToolCallPartDelta,
+)
+from pydantic_ai.models import Model
 from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import StrAttr
+from zrb.attr.type import StrAttr, fstring
 from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
@@ -15,7 +25,7 @@ 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.attr import get_attr, 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
@@ -34,7 +44,9 @@ class LLMTask(BaseTask):
         cli_only: bool = False,
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
-        model: StrAttr | None = LLM_MODEL,
+        model: (
+            Callable[[AnySharedContext], Model | str | fstring] | Model | None
+        ) = LLM_MODEL,
         model_settings: (
             ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
         ) = None,
@@ -93,7 +105,7 @@ class LLMTask(BaseTask):
             successor=successor,
         )
         self._model = model
-        self._model_settings = (model_settings,)
+        self._model_settings = model_settings
         self._agent = agent
         self._render_model = render_model
         self._system_prompt = system_prompt
@@ -108,6 +120,9 @@ class LLMTask(BaseTask):
         self._render_history_file = render_history_file
         self._max_call_iteration = max_call_iteration
 
+    def set_model(self, model: Model | str):
+        self._model = model
+
     def add_tool(self, tool: ToolOrCallable):
         self._additional_tools.append(tool)
 
@@ -115,15 +130,85 @@ class LLMTask(BaseTask):
         history = await self._read_conversation_history(ctx)
         user_prompt = self._get_message(ctx)
         agent = self._get_agent(ctx)
-        result = await agent.run(
+        async with agent.iter(
             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)))
+        ) as agent_run:
+            async for node in agent_run:
+                # Each node represents a step in the agent's execution
+                await self._print_node(ctx, agent_run, node)
+        new_history = json.loads(agent_run.result.all_messages_json())
         await self._write_conversation_history(ctx, new_history)
-        return result.data
+        return agent_run.result.data
+
+    async def _print_node(self, ctx: AnyContext, agent_run: Any, node: Any):
+        if Agent.is_user_prompt_node(node):
+            # A user prompt node => The user has provided input
+            ctx.print(stylize_faint(f">> UserPromptNode: {node.user_prompt}"))
+        elif Agent.is_model_request_node(node):
+            # A model request node => We can stream tokens from the model"s request
+            ctx.print(
+                stylize_faint(">> ModelRequestNode: streaming partial request tokens")
+            )
+            async with node.stream(agent_run.ctx) as request_stream:
+                is_streaming = False
+                async for event in request_stream:
+                    if isinstance(event, PartStartEvent):
+                        if is_streaming:
+                            ctx.print("", plain=True)
+                        ctx.print(
+                            stylize_faint(
+                                f"[Request] Starting part {event.index}: {event.part!r}"
+                            ),
+                        )
+                        is_streaming = False
+                    elif isinstance(event, PartDeltaEvent):
+                        if isinstance(event.delta, TextPartDelta):
+                            ctx.print(
+                                stylize_faint(f"{event.delta.content_delta}"),
+                                end="",
+                                plain=is_streaming,
+                            )
+                        elif isinstance(event.delta, ToolCallPartDelta):
+                            ctx.print(
+                                stylize_faint(f"{event.delta.args_delta}"),
+                                end="",
+                                plain=is_streaming,
+                            )
+                        is_streaming = True
+                    elif isinstance(event, FinalResultEvent):
+                        if is_streaming:
+                            ctx.print("", plain=True)
+                        ctx.print(
+                            stylize_faint(f"[Result] tool_name={event.tool_name}"),
+                        )
+                        is_streaming = False
+                if is_streaming:
+                    ctx.print("", plain=True)
+        elif Agent.is_call_tools_node(node):
+            # A handle-response node => The model returned some data, potentially calls a tool
+            ctx.print(
+                stylize_faint(
+                    ">> CallToolsNode: streaming partial response & tool usage"
+                )
+            )
+            async with node.stream(agent_run.ctx) as handle_stream:
+                async for event in handle_stream:
+                    if isinstance(event, FunctionToolCallEvent):
+                        ctx.print(
+                            stylize_faint(
+                                f"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})"  # noqa
+                            )
+                        )
+                    elif isinstance(event, FunctionToolResultEvent):
+                        ctx.print(
+                            stylize_faint(
+                                f"[Tools] Tool call {event.tool_call_id!r} returned => {event.result.content}"  # noqa
+                            )
+                        )
+        elif Agent.is_end_node(node):
+            # Once an End node is reached, the agent run is complete
+            ctx.print(stylize_faint(f"{agent_run.result.data}"))
 
     async def _write_conversation_history(
         self, ctx: AnyContext, conversations: list[Any]
@@ -135,11 +220,9 @@ class LLMTask(BaseTask):
             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
+        return self._model_settings
 
     def _get_agent(self, ctx: AnyContext) -> Agent:
         if isinstance(self._agent, Agent):
@@ -158,12 +241,16 @@ class LLMTask(BaseTask):
             self._get_model(ctx),
             system_prompt=self._get_system_prompt(ctx),
             tools=tools,
+            model_settings=self._get_model_settings(ctx),
         )
 
-    def _get_model(self, ctx: AnyContext) -> str:
-        return get_str_attr(
+    def _get_model(self, ctx: AnyContext) -> str | Model | None:
+        model = get_attr(
             ctx, self._model, "ollama_chat/llama3.1", auto_render=self._render_model
         )
+        if isinstance(model, (Model, str)) or model is None:
+            return model
+        raise ValueError("Invalid model")
 
     def _get_system_prompt(self, ctx: AnyContext) -> str:
         return get_str_attr(

{zrb-1.2.1.dist-info → zrb-1.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.2.1
+Version: 1.2.2
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -24,7 +24,7 @@ Requires-Dist: isort (>=5.13.2,<5.14.0)
 Requires-Dist: libcst (>=1.5.0,<2.0.0)
 Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
 Requires-Dist: psutil (>=6.1.1,<7.0.0)
-Requires-Dist: pydantic-ai (>=0.0.19,<0.0.20)
+Requires-Dist: pydantic-ai (>=0.0.31,<0.0.32)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)

{zrb-1.2.1.dist-info → zrb-1.2.2.dist-info}/RECORD

@@ -7,12 +7,13 @@ zrb/builtin/base64.py,sha256=1YnSwASp7OEAvQcsnHZGpJEvYoI1Z2zTIJ1bCDHfcPQ,921
 zrb/builtin/git.py,sha256=8_qVE_2lVQEVXQ9vhiw8Tn4Prj1VZB78ZjEJJS5Ab3M,5461
 zrb/builtin/git_subtree.py,sha256=7BKwOkVTWDrR0DXXQ4iJyHqeR6sV5VYRt8y_rEB0EHg,3505
 zrb/builtin/group.py,sha256=-phJfVpTX3_gUwS1u8-RbZUHe-X41kxDBSmrVh4rq8E,1682
-zrb/builtin/llm/llm_chat.py,sha256=UVhfJR0APXRC3tEv5i5vYbEANWqi04QD_WsAiARJ7j4,4494
+zrb/builtin/llm/llm_chat.py,sha256=cNAS_AS-Q5NW6qe8dJZh12b6c0zFKdNEvGfwJxNAUmw,5047
 zrb/builtin/llm/previous-session.js,sha256=xMKZvJoAbrwiyHS0OoPrWuaKxWYLoyR5sguePIoCjTY,816
-zrb/builtin/llm/tool/api.py,sha256=yQ3XV8O7Fx7hHssLSOcmiHDnevPhz9ktWi44HK7zTls,801
+zrb/builtin/llm/tool/api.py,sha256=bXFE7jihdhUscxJH8lu5imwlYH735AalbCyUTl28BaQ,826
 zrb/builtin/llm/tool/cli.py,sha256=to_IjkfrMGs6eLfG0cpVN9oyADWYsJQCtyluUhUdBww,253
-zrb/builtin/llm/tool/rag.py,sha256=PawaLZL-ThctxtBtsQuP3XsgTxQKyCGFqrudCANPJKk,5162
-zrb/builtin/llm/tool/web.py,sha256=N2HYuXbKPUpjVAq_UnQMbUrTIE8u0Ut3TeQadZ7_NJc,2217
+zrb/builtin/llm/tool/file.py,sha256=ibvh0zrsnponwyZvw6bWMUbpwSv5S5WUWCDfQ6BjVwk,1160
+zrb/builtin/llm/tool/rag.py,sha256=vEIThEy0JGwXEiNRLOEJAHAE0l1Qie2qvU3ryioeYMk,6066
+zrb/builtin/llm/tool/web.py,sha256=SDnCtYHZ0Q4DtLbIhc11a0UyyKbTTeW60UfeIKzK35k,3204
 zrb/builtin/md5.py,sha256=0pNlrfZA0wlZlHvFHLgyqN0JZJWGKQIF5oXxO44_OJk,949
 zrb/builtin/project/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/builtin/project/add/fastapp/fastapp_input.py,sha256=MKlWR_LxWhM_DcULCtLfL_IjTxpDnDBkn9KIqNmajFs,310
@@ -207,14 +208,14 @@ zrb/callback/callback.py,sha256=hKefB_Jd1XGjPSLQdMKDsGLHPzEGO2dqrIArLl_EmD0,848
 zrb/cmd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/cmd/cmd_result.py,sha256=L8bQJzWCpcYexIxHBNsXj2pT3BtLmWex0iJSMkvimOA,597
 zrb/cmd/cmd_val.py,sha256=7Doowyg6BK3ISSGBLt-PmlhzaEkBjWWm51cED6fAUOQ,1014
-zrb/config.py,sha256=Kb-GOsUS4poSCds4Wqg9LkscpS7BHXSy3dQmqvsFm2Q,4621
+zrb/config.py,sha256=MfHwcQ4OhCmCw6jXpFI8483Ase6YrqNGBvqYzwnwopw,4753
 zrb/content_transformer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/content_transformer/any_content_transformer.py,sha256=v8ZUbcix1GGeDQwB6OKX_1TjpY__ksxWVeqibwa_iZA,850
-zrb/content_transformer/content_transformer.py,sha256=YU6Xr3G_IaCWKQGsf9z9YlCclbiwcJ7ytQv3wKpPIiI,2125
+zrb/content_transformer/content_transformer.py,sha256=STl77wW-I69QaGzCXjvkppngYFLufow8ybPLSyAvlHs,2404
 zrb/context/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/context/any_context.py,sha256=2hgVKbbDwmwrEl1h1L1FaTUjuUYaDd_b7YRGkaorW6Q,6362
 zrb/context/any_shared_context.py,sha256=p1i9af_CUDz5Mf1h1kBZMAa2AEhf17I3O5IgAcjRLoY,1768
-zrb/context/context.py,sha256=qVMqt2tkLEFSI81mLYb_OSD615KH5jP685aUmHEm3XQ,6319
+zrb/context/context.py,sha256=VGoUwoWyL9d4QqJEhg41S-X8T2jlssGpiC9YSc3Gjqk,6601
 zrb/context/shared_context.py,sha256=47Tnnor1ttpwpe_N07rMNM1jgIYPY9abMe1Q5igkMtE,2735
 zrb/dot_dict/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/dot_dict/dot_dict.py,sha256=ubw_x8I7AOJ59xxtFVJ00VGmq_IYdZP3mUhNlO4nEK0,556
@@ -235,7 +236,7 @@ zrb/input/int_input.py,sha256=w5ewSxstUYv5LBAzvX_E0jIueXXdmnY2ehoQMTtg-EA,1380
 zrb/input/option_input.py,sha256=IrpF0XvFbH5G-IEAnoQ4QOvq7gn2wyT4jKwAdMKwV0s,2058
 zrb/input/password_input.py,sha256=Tu8TZx95717YsHICZ0zBzTUPKPf-K9vGlvRyaOTrFEM,1388
 zrb/input/str_input.py,sha256=NevZHX9rf1g8eMatPyy-kUX3DglrVAQpzvVpKAzf7bA,81
-zrb/input/text_input.py,sha256=wSNiYAx2xYPtl09Dfh_uHws9WG2dRqkS0Jnlm9HvD3s,3145
+zrb/input/text_input.py,sha256=gK5LGa9uEBBLwGqhTVmCVgsAuBYEQt3ANDQaSNtSc78,3300
 zrb/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/runner/cli.py,sha256=G_ILZCFzpV-kRE3dm1kq6BorB51TLJ34Qmhgy5SIMlU,6734
 zrb/runner/common_util.py,sha256=mjEBSmfdY2Sun2U5-8y8gGwF82OiRM8sgiYDOdW9NA4,1338
@@ -298,7 +299,7 @@ zrb/task/base_task.py,sha256=SQRf37bylS586KwyW0eYDe9JZ5Hl18FP8kScHae6y3A,21251
 zrb/task/base_trigger.py,sha256=jC722rDvodaBLeNaFghkTyv1u0QXrK6BLZUUqcmBJ7Q,4581
 zrb/task/cmd_task.py,sha256=pUKRSR4DZKjbmluB6vi7cxqyhxOLfJ2czSpYeQbiDvo,10705
 zrb/task/http_check.py,sha256=Gf5rOB2Se2EdizuN9rp65HpGmfZkGc-clIAlHmPVehs,2565
-zrb/task/llm_task.py,sha256=ptXC3x9Dwn7-4JrGQyEtzOXZ4dNQATDgCeowkvwAu9U,7723
+zrb/task/llm_task.py,sha256=B4qhza-4fk7odI7-rv2rLYvBLt1dmZMNgKu8OK7rajM,11849
 zrb/task/make_task.py,sha256=PD3b_aYazthS8LHeJsLAhwKDEgdurQZpymJDKeN60u0,2265
 zrb/task/rsync_task.py,sha256=GSL9144bmp6F0EckT6m-2a1xG25AzrrWYzH4k3SVUKM,6370
 zrb/task/scaffolder.py,sha256=rME18w1HJUHXgi9eTYXx_T2G4JdqDYzBoNOkdOOo5-o,6806
@@ -339,7 +340,7 @@ zrb/util/string/name.py,sha256=8picJfUBXNpdh64GNaHv3om23QHhUZux7DguFLrXHp8,1163
 zrb/util/todo.py,sha256=1nDdwPc22oFoK_1ZTXyf3638Bg6sqE2yp_U4_-frHoc,16015
 zrb/xcom/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/xcom/xcom.py,sha256=o79rxR9wphnShrcIushA0Qt71d_p3ZTxjNf7x9hJB78,1571
-zrb-1.2.1.dist-info/METADATA,sha256=Oj5aFm5hZeFXkWttWm6MNwk5uBCwVIGPHj2Y5IdoQyo,4198
-zrb-1.2.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zrb-1.2.1.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
-zrb-1.2.1.dist-info/RECORD,,
+zrb-1.2.2.dist-info/METADATA,sha256=GAm6vQds-lw7StRPEkrq3oTUSBO1jIYpXi9s1-KSfjw,4198
+zrb-1.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zrb-1.2.2.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
+zrb-1.2.2.dist-info/RECORD,,