zrb 1.21.31__py3-none-any.whl → 1.21.43__py3-none-any.whl

zrb/builtin/llm/chat_completion.py

@@ -11,16 +11,24 @@ from zrb.builtin.llm.chat_session_cmd import (
     CLEAR_SUB_CMD,
     HELP_CMD,
     HELP_CMD_DESC,
+    LOAD_CMD,
+    LOAD_CMD_DESC,
     MULTILINE_END_CMD,
     MULTILINE_END_CMD_DESC,
     MULTILINE_START_CMD,
     MULTILINE_START_CMD_DESC,
     QUIT_CMD,
     QUIT_CMD_DESC,
+    RESPONSE_CMD,
+    RESPONSE_CMD_DESC,
+    RESPONSE_SAVE_SUB_CMD_DESC,
     RUN_CLI_CMD,
     RUN_CLI_CMD_DESC,
     SAVE_CMD,
     SAVE_CMD_DESC,
+    SAVE_SUB_CMD,
+    SESSION_CMD,
+    SESSION_CMD_DESC,
     SET_SUB_CMD,
     WORKFLOW_ADD_SUB_CMD_DESC,
     WORKFLOW_CLEAR_SUB_CMD_DESC,
@@ -33,6 +41,8 @@ from zrb.builtin.llm.chat_session_cmd import (
     YOLO_SET_FALSE_CMD_DESC,
     YOLO_SET_TRUE_CMD_DESC,
 )
+from zrb.config.config import CFG
+from zrb.util.match import fuzzy_match
 
 if TYPE_CHECKING:
     from prompt_toolkit.completion import Completer
@@ -51,10 +61,39 @@ def get_chat_completer() -> "Completer":
                 yield completion
             for completion in self._complete_slash_file_command(document):
                 yield completion
+            for completion in self._complete_slash_session_load_command(document):
+                yield completion
             # Appendix
             for completion in self._complete_appendix(document):
                 yield completion
 
+        def _complete_slash_session_load_command(self, document: Document):
+            text = document.text_before_cursor
+            prefixes = []
+            for cmd in LOAD_CMD:
+                prefixes.append(f"{cmd} ")
+
+            for prefix in prefixes:
+                if text.startswith(prefix):
+                    pattern = text[len(prefix) :]
+                    # Use fuzzy_path_search but in LLM_HISTORY_DIR/save-point
+                    save_point_dir = os.path.join(CFG.LLM_HISTORY_DIR, "save-point")
+                    if not os.path.isdir(save_point_dir):
+                        return
+
+                    potential_options = self._fuzzy_path_search(
+                        pattern, root=save_point_dir, dirs=False, files=True
+                    )
+                    for option in potential_options:
+                        if option.startswith(save_point_dir):
+                            rel_option = os.path.relpath(option, save_point_dir)
+                        else:
+                            rel_option = option
+                        yield Completion(
+                            f"{prefix}{rel_option}",
+                            start_position=-len(text),
+                        )
+
         def _complete_slash_file_command(self, document: Document):
             text = document.text_before_cursor
             prefixes = []
@@ -64,7 +103,7 @@ def get_chat_completer() -> "Completer":
             for prefix in prefixes:
                 if text.startswith(prefix):
                     pattern = text[len(prefix) :]
-                    potential_options = self._fuzzy_path_search(pattern, dirs=False)
+                    potential_options = self._fuzzy_path_search(pattern, dirs=True)
                     for prefixed_option in [
                         f"{prefix}{option}" for option in potential_options
                     ]:
@@ -91,7 +130,7 @@ def get_chat_completer() -> "Completer":
             if not token.startswith(prefix):
                 return
             pattern = token[len(prefix) :]
-            potential_options = self._fuzzy_path_search(pattern, dirs=False)
+            potential_options = self._fuzzy_path_search(pattern, dirs=True)
             for prefixed_option in [
                 f"{prefix}{option}" for option in potential_options
             ]:
@@ -117,8 +156,16 @@ def get_chat_completer() -> "Completer":
                     cmd_options[f"{cmd} {subcmd}"] = WORKFLOW_CLEAR_SUB_CMD_DESC
                 for subcmd in SET_SUB_CMD:
                     cmd_options[f"{cmd} {subcmd}"] = WORKFLOW_SET_SUB_CMD_DESC
+            for cmd in SESSION_CMD:
+                cmd_options[cmd] = SESSION_CMD_DESC
             for cmd in SAVE_CMD:
                 cmd_options[cmd] = SAVE_CMD_DESC
+            for cmd in LOAD_CMD:
+                cmd_options[cmd] = LOAD_CMD_DESC
+            for cmd in RESPONSE_CMD:
+                cmd_options[cmd] = RESPONSE_CMD_DESC
+                for subcmd in SAVE_SUB_CMD:
+                    cmd_options[f"{cmd} {subcmd}"] = RESPONSE_SAVE_SUB_CMD_DESC
             for cmd in ATTACHMENT_CMD:
                 cmd_options[cmd] = ATTACHMENT_CMD_DESC
                 for subcmd in ADD_SUB_CMD:
@@ -156,47 +203,14 @@ def get_chat_completer() -> "Completer":
             - dirs/files booleans let you restrict results
             - returns list of relative paths (from root), sorted best-first
             """
-            search_pattern = pattern
-            if root is None:
-                # Determine root and adjust pattern if necessary
-                expanded_pattern = os.path.expanduser(pattern)
-                if os.path.isabs(expanded_pattern) or pattern.startswith("~"):
-                    # For absolute paths, find the deepest existing directory
-                    if os.path.isdir(expanded_pattern):
-                        root = expanded_pattern
-                        search_pattern = ""
-                    else:
-                        root = os.path.dirname(expanded_pattern)
-                        while root and not os.path.isdir(root) and len(root) > 1:
-                            root = os.path.dirname(root)
-                        if not os.path.isdir(root):
-                            root = "."  # Fallback
-                            search_pattern = pattern
-                        else:
-                            try:
-                                search_pattern = os.path.relpath(expanded_pattern, root)
-                                if search_pattern == ".":
-                                    search_pattern = ""
-                            except ValueError:
-                                search_pattern = os.path.basename(pattern)
-                else:
-                    root = "."
-                    search_pattern = pattern
-            # Normalize pattern -> tokens split on path separators or whitespace
-            search_pattern = search_pattern.strip()
-            if search_pattern:
-                raw_tokens = [t for t in search_pattern.split(os.path.sep) if t]
-            else:
-                raw_tokens = []
-            # prepare tokens (case)
-            if not case_sensitive:
-                tokens = [t.lower() for t in raw_tokens]
-            else:
-                tokens = raw_tokens
+            root, search_pattern = self._get_root_and_search_pattern(pattern, root)
             # specific ignore list
             try:
-                is_recursive = os.path.abspath(os.path.expanduser(root)).startswith(
-                    os.path.abspath(os.getcwd())
+                abs_root = os.path.abspath(os.path.expanduser(root))
+                abs_cwd = os.path.abspath(os.getcwd())
+                # Check if root is a subdirectory of cwd or is cwd itself
+                is_recursive = (
+                    abs_root.startswith(abs_cwd + os.sep) or abs_root == abs_cwd
                 )
             except Exception:
                 is_recursive = False
@@ -230,29 +244,8 @@ def get_chat_completer() -> "Completer":
                         ):
                             continue
                     cand = display_path.replace(os.sep, "/")  # unify separator
-                    cand_cmp = cand if case_sensitive else cand.lower()
-                    last_pos = 0
-                    score = 0.0
-                    matched_all = True
-                    for token in tokens:
-                        # try contiguous substring search first
-                        idx = cand_cmp.find(token, last_pos)
-                        if idx != -1:
-                            # good match: reward contiguous early matches
-                            score += idx  # smaller idx preferred
-                            last_pos = idx + len(token)
-                        else:
-                            # fallback to subsequence matching
-                            pos = self._find_subsequence_pos(cand_cmp, token, last_pos)
-                            if pos is None:
-                                matched_all = False
-                                break
-                            # subsequence match is less preferred than contiguous substring
-                            score += pos + 0.5 * len(token)
-                            last_pos = pos + len(token)
-                    if matched_all:
-                        # prefer shorter paths when score ties, so include length as tiebreaker
-                        score += 0.01 * len(cand)
+                    matched, score = fuzzy_match(cand, search_pattern)
+                    if matched:
                         out = (
                             cand
                             if os.path.abspath(cand) == cand
@@ -263,25 +256,42 @@ def get_chat_completer() -> "Completer":
             candidates.sort(key=lambda x: (x[0], x[1]))
             return [p for _, p in candidates[:max_results]]
 
-        def _find_subsequence_pos(
-            self, hay: str, needle: str, start: int = 0
-        ) -> int | None:
-            """
-            Try to locate needle in hay as a subsequence starting at `start`.
-            Returns the index of the first matched character of the subsequence or None if not
-            match.
-            """
-            if not needle:
-                return start
-            i = start
-            j = 0
-            first_pos = None
-            while i < len(hay) and j < len(needle):
-                if hay[i] == needle[j]:
-                    if first_pos is None:
-                        first_pos = i
-                    j += 1
-                i += 1
-            return first_pos if j == len(needle) else None
+        def _get_root_and_search_pattern(
+            self,
+            pattern: str,
+            root: str | None = None,
+        ) -> tuple[str, str]:
+            search_pattern = pattern
+            if root is None:
+                # Determine root and adjust pattern if necessary
+                expanded_pattern = os.path.expanduser(pattern)
+                if os.path.isabs(expanded_pattern) or pattern.startswith("~"):
+                    # For absolute paths, find the deepest existing directory
+                    if os.path.isdir(expanded_pattern):
+                        root = expanded_pattern
+                        return (root, "")
+                    root = os.path.dirname(expanded_pattern)
+                    while root and not os.path.isdir(root) and len(root) > 1:
+                        root = os.path.dirname(root)
+                    if not os.path.isdir(root):
+                        return (".", pattern)  # Fallback
+                    try:
+                        search_pattern = os.path.relpath(expanded_pattern, root)
+                        if search_pattern == ".":
+                            return (root, "")
+                    except ValueError:
+                        return (root, os.path.basename(pattern))
+                # Handle redundant current directory prefixes (e.g., ./ or .\)
+                if pattern.startswith(f".{os.sep}"):
+                    return (".", pattern[len(f".{os.sep}") :])
+                if os.sep != "/" and pattern.startswith("./"):
+                    return (".", pattern[2:])
+                return (".", pattern)
+
+            if pattern.startswith(f".{os.sep}"):
+                pattern = pattern[len(f".{os.sep}") :]
+            elif os.sep != "/" and pattern.startswith("./"):
+                pattern = pattern[2:]
+            return (root, pattern)
 
     return ChatCompleter()

zrb/builtin/llm/chat_session.py

@@ -5,27 +5,39 @@ from typing import TYPE_CHECKING, Any
 from zrb.builtin.llm.chat_session_cmd import (
     ATTACHMENT_CMD,
     HELP_CMD,
+    LOAD_CMD,
     MULTILINE_END_CMD,
     MULTILINE_START_CMD,
     QUIT_CMD,
+    RESPONSE_CMD,
     RUN_CLI_CMD,
     SAVE_CMD,
+    SESSION_CMD,
     WORKFLOW_CMD,
     YOLO_CMD,
     get_new_attachments,
     get_new_workflows,
     get_new_yolo_mode,
+    handle_response_cmd,
+    handle_session,
     is_command_match,
     print_commands,
     print_current_attachments,
     print_current_workflows,
     print_current_yolo_mode,
     run_cli_command,
-    save_final_result,
 )
 from zrb.builtin.llm.chat_trigger import llm_chat_trigger
+from zrb.builtin.llm.history import get_last_session_name
+from zrb.builtin.llm.xcom_names import (
+    LLM_ASK_ERROR_XCOM_NAME,
+    LLM_ASK_RESULT_XCOM_NAME,
+    LLM_ASK_SESSION_XCOM_NAME,
+)
 from zrb.config.llm_config import llm_config
 from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.task.llm.workflow import get_llm_loaded_workflow_xcom
 from zrb.util.cli.markdown import render_markdown
 
 if TYPE_CHECKING:
@@ -44,30 +56,39 @@ async def read_user_prompt(ctx: AnyContext) -> str:
     reader: PromptSession[Any] | StreamReader = await _setup_input_reader(is_tty)
     multiline_mode = False
     is_first_time = True
-    current_workflows: str = ctx.input.workflows
+    current_workflows: str = ctx.input.workflow
     current_yolo_mode: bool | str = ctx.input.yolo
     current_attachments: str = ctx.input.attach
     user_inputs: list[str] = []
     final_result: str = ""
     should_end = False
+    start_new: bool = ctx.input.start_new
+    if not start_new and ctx.input.previous_session == "":
+        session = ctx.session
+        if session is not None:
+            # Automatically inject last session name as previous session
+            last_session_name = get_last_session_name()
+            session.shared_ctx.input["previous_session"] = last_session_name
+            session.shared_ctx.input["previous-session"] = last_session_name
+    current_session_name: str | None = ctx.input.previous_session
     while not should_end:
-        await asyncio.sleep(0.01)
-        previous_session_name: str | None = (
-            ctx.input.previous_session if is_first_time else ""
-        )
-        start_new: bool = ctx.input.start_new if is_first_time else False
+        await asyncio.sleep(0)
         if is_first_time and ctx.input.message.strip() != "":
             user_input = ctx.input.message
         else:
             # Get user input based on mode
             if not multiline_mode:
                 ctx.print("💬 >>", plain=True)
-            user_input = await llm_chat_trigger.wait(reader, ctx)
+            user_input = await llm_chat_trigger.wait(
+                ctx, reader, current_session_name, is_first_time
+            )
             if not multiline_mode:
                 ctx.print("", plain=True)
         # At this point, is_first_time has to be False
         if is_first_time:
             is_first_time = False
+        # Add additional workflows activated by LLM in the previous session
+        current_workflows = _get_new_workflows_from_xcom(ctx, current_workflows)
         # Handle user input (including slash commands)
         if multiline_mode:
             if is_command_match(user_input, MULTILINE_END_CMD):
@@ -87,8 +108,8 @@ async def read_user_prompt(ctx: AnyContext) -> str:
                 current_workflows = get_new_workflows(current_workflows, user_input)
                 print_current_workflows(ctx, current_workflows)
                 continue
-            elif is_command_match(user_input, SAVE_CMD):
-                save_final_result(ctx, user_input, final_result)
+            elif is_command_match(user_input, RESPONSE_CMD):
+                handle_response_cmd(ctx, user_input, final_result)
                 continue
             elif is_command_match(user_input, ATTACHMENT_CMD):
                 current_attachments = get_new_attachments(
@@ -101,25 +122,35 @@ async def read_user_prompt(ctx: AnyContext) -> str:
                 print_current_yolo_mode(ctx, current_yolo_mode)
                 continue
             elif is_command_match(user_input, RUN_CLI_CMD):
-                run_cli_command(ctx, user_input)
+                await run_cli_command(ctx, user_input)
                 continue
             elif is_command_match(user_input, HELP_CMD):
                 print_commands(ctx)
                 continue
+            elif (
+                is_command_match(user_input, SESSION_CMD)
+                or is_command_match(user_input, SAVE_CMD)
+                or is_command_match(user_input, LOAD_CMD)
+            ):
+                current_session_name, start_new = handle_session(
+                    ctx, current_session_name, start_new, user_input
+                )
             else:
                 user_inputs.append(user_input)
         # Trigger LLM
         user_prompt = "\n".join(user_inputs)
         user_inputs = []
-        result = await _trigger_ask_and_wait_for_result(
+        result, current_session_name = await _trigger_ask_and_wait_for_result(
             ctx=ctx,
             user_prompt=user_prompt,
             attach=current_attachments,
             workflows=current_workflows,
             yolo_mode=current_yolo_mode,
-            previous_session_name=previous_session_name,
+            previous_session_name=current_session_name,
             start_new=start_new,
         )
+        # After the first trigger, we no longer need to force start_new
+        start_new = False
         current_attachments = ""
         final_result = final_result if result is None else result
         if ctx.is_web_mode or not is_tty:
@@ -127,6 +158,23 @@ async def read_user_prompt(ctx: AnyContext) -> str:
     return final_result
 
 
+def _get_new_workflows_from_xcom(ctx: AnyContext, current_workflows: str):
+    llm_loaded_workflow_xcom = get_llm_loaded_workflow_xcom(ctx)
+    new_workflow_names = [
+        workflow_name.strip()
+        for workflow_name in current_workflows.split(",")
+        if workflow_name.strip() != ""
+    ]
+    while len(llm_loaded_workflow_xcom) > 0:
+        additional_workflow_names = [
+            workflow_name
+            for workflow_name in llm_loaded_workflow_xcom.pop()
+            if workflow_name not in new_workflow_names
+        ]
+        new_workflow_names += additional_workflow_names
+    return ",".join(new_workflow_names)
+
+
 async def _setup_input_reader(
     is_interactive: bool,
 ) -> "PromptSession[Any] | StreamReader":
@@ -151,7 +199,7 @@ async def _trigger_ask_and_wait_for_result(
     yolo_mode: bool | str,
     previous_session_name: str | None = None,
     start_new: bool = False,
-) -> str | None:
+) -> tuple[str | None, str | None]:
     """
     Triggers the LLM ask task and waits for the result via XCom.
 
@@ -162,22 +210,27 @@ async def _trigger_ask_and_wait_for_result(
         start_new: Whether to start a new conversation (optional).
 
     Returns:
-        The result from the LLM task, or None if the user prompt is empty.
+        The result from the LLM task and the session name.
     """
     if user_prompt.strip() == "":
-        return None
+        return None, previous_session_name
     await _trigger_ask(
         ctx, user_prompt, attach, workflows, yolo_mode, previous_session_name, start_new
     )
     result = await _wait_ask_result(ctx)
+
+    resolved_session_name = previous_session_name
+    if result is not None:
+        resolved_session_name = await _wait_ask_session_name(ctx)
+
     md_result = render_markdown(result) if result is not None else ""
     ctx.print("\n🤖 >>", plain=True)
     ctx.print(md_result, plain=True)
     ctx.print("", plain=True)
-    return result
+    return result, resolved_session_name
 
 
-def get_llm_ask_input_mapping(callback_ctx: AnyContext):
+def get_llm_ask_input_mapping(callback_ctx: AnyContext | AnySharedContext):
     """
     Generates the input mapping for the LLM ask task from the callback context.
 
@@ -200,7 +253,7 @@ def get_llm_ask_input_mapping(callback_ctx: AnyContext):
         "previous-session": data.get("previous_session_name"),
         "message": data.get("message"),
         "attach": data.get("attach"),
-        "workflows": data.get("workflows"),
+        "workflow": data.get("workflow"),
         "yolo": data.get("yolo"),
     }
 
@@ -223,15 +276,13 @@ async def _trigger_ask(
         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,
             "attach": attach,
-            "workflows": workflows,
+            "workflow": workflows,
             "yolo": yolo_mode,
         }
     )
@@ -247,12 +298,18 @@ async def _wait_ask_result(ctx: AnyContext) -> str | None:
     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)
-        if "ask_error" in ctx.xcom and len(ctx.xcom.ask_error) > 0:
-            ctx.xcom.ask_error.pop()
+    while (
+        LLM_ASK_RESULT_XCOM_NAME not in ctx.xcom
+        or len(ctx.xcom[LLM_ASK_RESULT_XCOM_NAME]) == 0
+    ):
+        await asyncio.sleep(0)
+        if (
+            LLM_ASK_ERROR_XCOM_NAME in ctx.xcom
+            and len(ctx.xcom[LLM_ASK_ERROR_XCOM_NAME]) > 0
+        ):
+            ctx.xcom[LLM_ASK_ERROR_XCOM_NAME].pop()
             return None
-    return ctx.xcom.ask_result.pop()
+    return ctx.xcom[LLM_ASK_RESULT_XCOM_NAME].pop()
 
 
 async def _wait_ask_session_name(ctx: AnyContext) -> str:
@@ -265,6 +322,9 @@ async def _wait_ask_session_name(ctx: AnyContext) -> str:
     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()
+    while (
+        LLM_ASK_SESSION_XCOM_NAME not in ctx.xcom
+        or len(ctx.xcom[LLM_ASK_SESSION_XCOM_NAME]) == 0
+    ):
+        await asyncio.sleep(0)
+    return ctx.xcom[LLM_ASK_SESSION_XCOM_NAME].pop()