lionagi 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

lionagi/core/unit/unit_mixin.py

@@ -0,0 +1,1168 @@
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""
+The base directive module.
+"""
+
+import asyncio
+import contextlib
+import re
+from abc import ABC
+
+from typing import Any, Optional
+
+from lionagi.libs import ParseUtil, StringMatch, to_list
+from lionagi.libs.ln_nested import nmerge
+from lionagi.core.collections.abc import ActionError
+from lionagi.core.message import ActionRequest, ActionResponse, Instruction
+from lionagi.core.message.util import _parse_action_request
+from lionagi.core.report.form import Form
+from lionagi.core.unit.util import process_tools
+from lionagi.core.validator.validator import Validator
+
+
+class DirectiveMixin(ABC):
+    """
+    DirectiveMixin is a class for handling chat operations and
+    processing responses.
+    """
+
+    def _create_chat_config(
+        self,
+        system: Optional[str] = None,
+        instruction: Optional[str] = None,
+        context: Optional[str] = None,
+        images: Optional[str] = None,
+        sender: Optional[str] = None,
+        recipient: Optional[str] = None,
+        requested_fields: Optional[list] = None,
+        form: Form = None,
+        tools: bool = False,
+        branch: Optional[Any] = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Create the chat configuration based on the provided parameters.
+
+        Args:
+            system: System message.
+            instruction: Instruction message.
+            context: Context message.
+            sender: Sender identifier.
+            recipient: Recipient identifier.
+            requested_fields: Fields requested in the response.
+            form: Form data.
+            tools: Flag indicating if tools should be used.
+            branch: Branch instance.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            dict: The chat configuration.
+        """
+        branch = branch or self.branch
+
+        if system:
+            branch.add_message(system=system)
+
+        if not form:
+            if recipient == "branch.ln_id":
+                recipient = branch.ln_id
+
+            branch.add_message(
+                instruction=instruction,
+                context=context,
+                sender=sender,
+                recipient=recipient,
+                requested_fields=requested_fields,
+                images=images,
+            )
+        else:
+            instruct_ = Instruction.from_form(form)
+            branch.add_message(instruction=instruct_)
+
+        if "tool_parsed" in kwargs:
+            kwargs.pop("tool_parsed")
+            tool_kwarg = {"tools": tools}
+            kwargs = tool_kwarg | kwargs
+        elif tools and branch.has_tools:
+            kwargs = branch.tool_manager.parse_tool(tools=tools, **kwargs)
+
+        config = {**self.imodel.config, **kwargs}
+        if sender is not None:
+            config["sender"] = sender
+
+        return config
+
+    async def _call_chatcompletion(
+        self, imodel: Optional[Any] = None, branch: Optional[Any] = None, **kwargs
+    ) -> Any:
+        """
+        Calls the chat completion model.
+
+        Args:
+            imodel: The model instance.
+            branch: The branch instance.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The chat completion result.
+        """
+        imodel = imodel or self.imodel
+        branch = branch or self.branch
+        return await imodel.call_chat_completion(branch.to_chat_messages(), **kwargs)
+
+    async def _process_chatcompletion(
+        self,
+        payload: dict,
+        completion: dict,
+        sender: str,
+        invoke_tool: bool = True,
+        branch: Optional[Any] = None,
+        action_request: Optional[Any] = None,
+        costs=None,
+    ) -> Any:
+        """
+        Processes the chat completion response.
+        Currently only support last message for function calling
+
+        Args:
+            payload: The payload data.
+            completion: The completion data.
+            sender: The sender identifier.
+            invoke_tool: Flag indicating if tools should be invoked.
+            branch: The branch instance.
+            action_request: The action request instance.
+
+        Returns:
+            Any: The processed result.
+        """
+        branch = branch or self.branch
+        _msg = None
+
+        if "choices" in completion:
+            payload.pop("messages", None)
+            branch.update_last_instruction_meta(payload)
+            _choices = completion.pop("choices", None)
+
+            def process_completion_choice(choice, price=None):
+                if isinstance(choice, dict):
+                    msg = choice.pop("message", None)
+                    _completion = completion.copy()
+                    _completion.update(choice)
+                    branch.add_message(
+                        assistant_response=msg,
+                        metadata=_completion,
+                        sender=sender,
+                    )
+
+                a = branch.messages[-1]._meta_get(
+                    ["extra", "usage", "prompt_tokens"], 0
+                )
+                b = branch.messages[-1]._meta_get(
+                    ["extra", "usage", "completion_tokens"], 0
+                )
+                m = completion.get("model", None)
+                if m:
+                    ttl = (a * price[0] + b * price[1]) / 1000000
+                    branch.messages[-1]._meta_insert(["extra", "usage", "expense"], ttl)
+                return msg
+
+            if _choices and not isinstance(_choices, list):
+                _choices = [_choices]
+
+            if _choices and isinstance(_choices, list):
+                for _choice in _choices:
+                    _msg = process_completion_choice(_choice, price=costs)
+
+            branch.imodel.status_tracker.num_tasks_succeeded += 1
+        else:
+            branch.imodel.status_tracker.num_tasks_failed += 1
+
+        return await self._process_action_request(
+            _msg=_msg,
+            branch=branch,
+            invoke_tool=invoke_tool,
+            action_request=action_request,
+        )
+
+    async def _process_action_request(
+        self,
+        _msg: Optional[dict] = None,
+        branch: Optional[Any] = None,
+        invoke_tool: bool = True,
+        action_request: Optional[Any] = None,
+    ) -> Any:
+        """
+        Processes an action request from the assistant response.
+
+        Args:
+            _msg: The message data.
+            branch: The branch instance.
+            invoke_tool: Flag indicating if tools should be invoked.
+            action_request: The action request instance.
+
+        Returns:
+            Any: The processed result.
+        """
+        action_request = action_request or _parse_action_request(_msg)
+        if action_request is None:
+            return _msg if _msg else False
+
+        if action_request:
+            for i in action_request:
+                if i.function in branch.tool_manager.registry:
+                    i.recipient = branch.tool_manager.registry[i.function].ln_id
+                else:
+                    raise ActionError(f"Tool {i.function} not found in registry")
+                branch.add_message(action_request=i, recipient=i.recipient)
+
+        if invoke_tool:
+            tasks = []
+            for i in action_request:
+                tool = branch.tool_manager.registry[i.function]
+                tasks.append(asyncio.create_task(tool.invoke(i.arguments)))
+            results = await asyncio.gather(*tasks)
+
+            for idx, item in enumerate(results):
+                if item is not None:
+                    branch.add_message(
+                        action_request=action_request[idx],
+                        func_outputs=item,
+                        sender=action_request[idx].recipient,
+                        recipient=action_request[idx].sender,
+                    )
+
+        return None
+
+    async def _output(
+        self,
+        payload: dict,
+        completion: dict,
+        sender: str,
+        invoke_tool: bool,
+        requested_fields: dict,
+        form: Form = None,
+        return_form: bool = True,
+        strict: bool = False,
+        rulebook: Any = None,
+        use_annotation: bool = True,
+        template_name: str = None,
+        costs=None,
+    ) -> Any:
+        """
+        Outputs the final processed response.
+
+        Args:
+            payload: The payload data.
+            completion: The completion data.
+            sender: The sender identifier.
+            invoke_tool: Flag indicating if tools should be invoked.
+            requested_fields: Fields requested in the response.
+            form: Form data.
+            return_form: Flag indicating if form should be returned.
+            strict: Flag indicating if strict validation should be applied.
+            rulebook: Rulebook instance for validation.
+            use_annotation: Flag indicating if annotations should be used.
+            template_name: Template name for form.
+
+        Returns:
+            Any: The processed response.
+        """
+        _msg = await self._process_chatcompletion(
+            payload=payload,
+            completion=completion,
+            sender=sender,
+            invoke_tool=invoke_tool,
+            costs=costs,
+        )
+
+        if _msg is None:
+            return None
+
+        response_ = self._process_model_response(_msg, requested_fields)
+
+        if form:
+            validator = Validator(rulebook=rulebook) if rulebook else self.validator
+            form = await validator.validate_response(
+                form=form,
+                response=response_,
+                strict=strict,
+                use_annotation=use_annotation,
+            )
+            if template_name:
+                form.template_name = template_name
+
+            return (
+                form
+                if return_form
+                else {
+                    i: form.work_fields[i]
+                    for i in form.requested_fields
+                    if form.work_fields[i] is not None
+                }
+            )
+
+        return response_
+
+    async def _base_chat(
+        self,
+        instruction: Any = None,
+        *,
+        system: Any = None,
+        context: Any = None,
+        sender: Any = None,
+        recipient: Any = None,
+        requested_fields: dict = None,
+        form: Form = None,
+        tools: Any = False,
+        images: Optional[str] = None,
+        invoke_tool: bool = True,
+        return_form: bool = True,
+        strict: bool = False,
+        rulebook: Any = None,
+        imodel: Any = None,
+        use_annotation: bool = True,
+        branch: Any = None,
+        clear_messages: bool = False,
+        return_branch: bool = False,
+        **kwargs,
+    ) -> Any:
+        """
+        Handles the base chat operation by configuring the chat and
+        processing the response.
+
+        Args:
+            instruction: Instruction message.
+            system: System message.
+            context: Context message.
+            sender: Sender identifier.
+            recipient: Recipient identifier.
+            requested_fields: Fields requested in the response.
+            form: Form data.
+            tools: Flag indicating if tools should be used.
+            invoke_tool: Flag indicating if tools should be invoked.
+            return_form: Flag indicating if form should be returned.
+            strict: Flag indicating if strict validation should be applied.
+            rulebook: Rulebook instance for validation.
+            imodel: Model instance.
+            use_annotation: Flag indicating if annotations should be used.
+            branch: Branch instance.
+            clear_messages: Flag indicating if messages should be cleared.
+            return_branch: Flag indicating if branch should be returned.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The processed response and branch.
+        """
+        branch = branch or self.branch
+        if clear_messages:
+            branch.clear()
+            branch.set_system(system)
+
+        config = self._create_chat_config(
+            system=system,
+            instruction=instruction,
+            context=context,
+            sender=sender,
+            recipient=recipient,
+            requested_fields=requested_fields,
+            form=form,
+            tools=tools,
+            branch=branch,
+            images=images,
+            **kwargs,
+        )
+
+        payload, completion = await self._call_chatcompletion(
+            imodel=imodel, branch=branch, **config
+        )
+
+        imodel = imodel or self.imodel
+        out_ = await self._output(
+            payload=payload,
+            completion=completion,
+            sender=sender,
+            invoke_tool=invoke_tool,
+            requested_fields=requested_fields,
+            form=form,
+            return_form=return_form,
+            strict=strict,
+            rulebook=rulebook,
+            use_annotation=use_annotation,
+            costs=imodel.costs,
+        )
+
+        return out_, branch if return_branch else out_
+
+    async def _chat(
+        self,
+        instruction=None,
+        context=None,
+        system=None,
+        sender=None,
+        recipient=None,
+        branch=None,
+        requested_fields=None,
+        form: Form = None,
+        tools=False,
+        invoke_tool=True,
+        return_form=True,
+        strict=False,
+        rulebook=None,
+        imodel=None,
+        images: Optional[str] = None,
+        clear_messages=False,
+        use_annotation=True,
+        timeout: float = None,
+        return_branch=False,
+        **kwargs,
+    ):
+        """
+        Handles the chat operation.
+
+        Args:
+            instruction: Instruction message.
+            context: Context message.
+            system: System message.
+            sender: Sender identifier.
+            recipient: Recipient identifier.
+            branch: Branch instance.
+            requested_fields: Fields requested in the response.
+            form: Form data.
+            tools: Flag indicating if tools should be used.
+            invoke_tool: Flag indicating if tools should be invoked.
+            return_form: Flag indicating if form should be returned.
+            strict: Flag indicating if strict validation should be applied.
+            rulebook: Rulebook instance for validation.
+            imodel: Model instance.
+            clear_messages: Flag indicating if messages should be cleared.
+            use_annotation: Flag indicating if annotations should be used.
+            timeout: Timeout value.
+            return_branch: Flag indicating if branch should be returned.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The processed response.
+        """
+        a = await self._base_chat(
+            context=context,
+            instruction=instruction,
+            system=system,
+            sender=sender,
+            recipient=recipient,
+            requested_fields=requested_fields,
+            form=form,
+            tools=tools,
+            images=images,
+            invoke_tool=invoke_tool,
+            return_form=return_form,
+            strict=strict,
+            rulebook=rulebook,
+            imodel=imodel,
+            use_annotation=use_annotation,
+            timeout=timeout,
+            branch=branch,
+            clear_messages=clear_messages,
+            return_branch=return_branch,
+            **kwargs,
+        )
+
+        if isinstance(a, str):
+            return a
+
+        a = list(a)
+
+        if len(a) == 2 and a[0] == a[1]:
+            return a[0] if not isinstance(a[0], tuple) else a[0][0]
+        if len(a) == 2 and a[0] != a[1]:
+            return a[0], a[1]
+        if len(a) == 1 and isinstance(a[0], tuple):
+            return a[0][0]
+        if len(a) == 1 and not isinstance(a[0], tuple):
+            return a[0]
+
+    async def _direct(
+        self,
+        instruction=None,
+        context=None,
+        form: Form = None,
+        branch=None,
+        tools=None,
+        reason: bool = None,
+        predict: bool = None,
+        score: bool = None,
+        select: bool = None,
+        plan: bool = None,
+        allow_action: bool = None,
+        allow_extension: bool = None,
+        confidence: bool = None,
+        max_extension: int = None,
+        score_num_digits=None,
+        score_range=None,
+        select_choices=None,
+        plan_num_step=None,
+        predict_num_sentences=None,
+        clear_messages=False,
+        return_branch=False,
+        images: Optional[str] = None,
+        verbose=None,
+        **kwargs,
+    ):
+        """
+        Directs the operation based on the provided parameters.
+
+        Args:
+            instruction: Instruction message.
+            context: Context message.
+            form: Form data.
+            branch: Branch instance.
+            tools: Tools data.
+            reason: Flag indicating if reason should be included.
+            predict: Flag indicating if prediction should be included.
+            score: Flag indicating if score should be included.
+            select: Flag indicating if selection should be included.
+            plan: Flag indicating if plan should be included.
+            allow_action: Flag indicating if action should be allowed.
+            allow_extension: Flag indicating if extension should be allowed.
+            confidence: Flag indicating if confidence should be included.
+            max_extension: Maximum extension value.
+            score_num_digits: Number of digits for score.
+            score_range: Range for score.
+            select_choices: Choices for selection.
+            plan_num_step: Number of steps for plan.
+            predict_num_sentences: Number of sentences for prediction.
+            clear_messages: Flag indicating if messages should be cleared.
+            return_branch: Flag indicating if branch should be returned.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The processed response and branch.
+        """
+        a = await self._base_direct(
+            instruction=instruction,
+            context=context,
+            form=form,
+            branch=branch,
+            tools=tools,
+            reason=reason,
+            predict=predict,
+            score=score,
+            select=select,
+            images=images,
+            plan=plan,
+            allow_action=allow_action,
+            allow_extension=allow_extension,
+            confidence=confidence,
+            max_extension=max_extension,
+            score_num_digits=score_num_digits,
+            score_range=score_range,
+            select_choices=select_choices,
+            plan_num_step=plan_num_step,
+            predict_num_sentences=predict_num_sentences,
+            clear_messages=clear_messages,
+            return_branch=return_branch,
+            verbose=verbose,
+            **kwargs,
+        )
+
+        a = list(a)
+        if len(a) == 2 and a[0] == a[1]:
+            return a[0] if not isinstance(a[0], tuple) else a[0][0]
+
+        return a[0], a[1]
+
+    async def _base_direct(
+        self,
+        instruction=None,
+        *,
+        context=None,
+        form: Form = None,
+        branch=None,
+        tools=None,
+        reason: bool = None,
+        predict: bool = None,
+        score: bool = None,
+        select: bool = None,
+        plan: bool = None,
+        allow_action: bool = None,
+        allow_extension: bool = None,
+        confidence: bool = None,
+        max_extension: int = None,
+        score_num_digits=None,
+        score_range=None,
+        select_choices=None,
+        plan_num_step=None,
+        predict_num_sentences=None,
+        clear_messages=False,
+        return_branch=False,
+        images: Optional[str] = None,
+        verbose=None,
+        **kwargs,
+    ):
+        """
+        Handles the base direct operation.
+
+        Args:
+            instruction: Instruction message.
+            context: Context message.
+            form: Form data.
+            branch: Branch instance.
+            tools: Tools data.
+            reason: Flag indicating if reason should be included.
+            predict: Flag indicating if prediction should be included.
+            score: Flag indicating if score should be included.
+            select: Flag indicating if selection should be included.
+            plan: Flag indicating if plan should be included.
+            allow_action: Flag indicating if action should be allowed.
+            allow_extension: Flag indicating if extension should be allowed.
+            confidence: Flag indicating if confidence should be included.
+            max_extension: Maximum extension value.
+            score_num_digits: Number of digits for score.
+            score_range: Range for score.
+            select_choices: Choices for selection.
+            plan_num_step: Number of steps for plan.
+            predict_num_sentences: Number of sentences for prediction.
+            clear_messages: Flag indicating if messages should be cleared.
+            return_branch: Flag indicating if branch should be returned.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The processed response and branch.
+        """
+        # Ensure branch is initialized
+        branch = branch or self.branch
+        if clear_messages:
+            branch.clear()
+
+        # Set a default max_extension if allow_extension is True and max_extension is None
+        if allow_extension and not max_extension:
+            max_extension = 3  # Set a default limit for recursion
+
+        # Process tools if provided
+        if tools:
+            process_tools(tools, branch)
+
+        if allow_action and not tools:
+            tools = True
+
+        tool_schema=None
+        if tools:
+            tool_schema = branch.tool_manager.get_tool_schema(tools)
+
+        if not form:
+            form = self.default_template(
+                instruction=instruction,
+                context=context,
+                reason=reason,
+                predict=predict,
+                score=score,
+                select=select,
+                plan=plan,
+                tool_schema=tool_schema,
+                allow_action=allow_action,
+                allow_extension=allow_extension,
+                max_extension=max_extension,
+                confidence=confidence,
+                score_num_digits=score_num_digits,
+                score_range=score_range,
+                select_choices=select_choices,
+                plan_num_step=plan_num_step,
+                predict_num_sentences=predict_num_sentences,
+            )
+
+        elif form and "tool_schema" not in form._all_fields:
+            form.append_to_input("tool_schema")
+            form.tool_schema = tool_schema
+
+        else:
+            form.tool_schema = tool_schema
+
+        verbose = (
+            verbose
+            if verbose is not None and isinstance(verbose, bool)
+            else self.verbose
+        )
+        if verbose:
+            print("Chatting with model...")
+
+        # Call the base chat method
+        form = await self._chat(
+            form=form,
+            branch=branch,
+            images=images,
+            **kwargs,
+        )
+
+        # Handle actions if allowed and required
+        if allow_action and getattr(form, "action_required", None):
+            actions = getattr(form, "actions", None)
+            if actions:
+                if verbose:
+                    print(
+                        "Found action requests in model response. Processing actions..."
+                    )
+                form = await self._act(form, branch, actions=actions)
+                if verbose:
+                    print("Actions processed!")
+
+        last_form = form
+
+        ctr = 1
+
+        # Handle extensions if allowed and required
+        extension_forms = []
+        max_extension = max_extension if isinstance(max_extension, int) else 3
+        while (
+            allow_extension
+            and max_extension > 0
+            and getattr(last_form, "extension_required", None)
+        ):
+            if getattr(last_form, "is_extension", None):
+                break
+            if verbose:
+                print(f"\nFound extension requests in model response.")
+                print(
+                    f"------------------- Processing extension No.{ctr} -------------------"
+                )
+
+            max_extension -= 1
+
+            # new form cannot be extended, otherwise it will be an infinite loop
+            new_form = await self._extend(
+                tools=tools,
+                reason=reason,
+                predict=predict,
+                score=score,
+                select=select,
+                plan=getattr(last_form, "plan", None),
+                allow_action=allow_action,
+                confidence=confidence,
+                score_num_digits=score_num_digits,
+                score_range=score_range,
+                select_choices=select_choices,
+                predict_num_sentences=predict_num_sentences,
+                **kwargs,
+            )
+
+            if verbose:
+                print(f"------------------- Extension completed -------------------\n")
+
+            extension_forms.extend(new_form)
+            last_form = new_form[-1] if isinstance(new_form, list) else new_form
+            ctr += len(form)
+
+        if extension_forms:
+            if not getattr(form, "extension_forms", None):
+                form._add_field("extension_forms", list, None, [])
+            form.extension_forms.extend(extension_forms)
+            action_responses = [
+                i.action_response
+                for i in extension_forms
+                if getattr(i, "action_response", None) is not None
+            ]
+            if not hasattr(form, "action_response"):
+                form.add_field("action_response", {})
+
+            for action_response in action_responses:
+                nmerge([form.action_response, action_response])
+
+        if "PLEASE_ACTION" in form.answer:
+            if verbose:
+                print("Analyzing action responses and generating answer...")
+
+            answer = await self._chat(
+                "please provide final answer basing on the above"
+                " information, provide answer value as a string only"
+                " do not return as json, do not include other information",
+            )
+
+            if isinstance(answer, dict):
+                a = answer.get("answer", None)
+                if a is not None:
+                    answer = a
+
+            answer = str(answer).strip()
+            if answer.startswith("{") and answer.endswith("}"):
+                answer = answer[1:-1]
+                answer = answer.strip()
+            if '"answer":' in answer:
+                answer.replace('"answer":', "")
+                answer = answer.strip()
+            elif "'answer':" in answer:
+                answer.replace("'answer':", "")
+                answer = answer.strip()
+
+            form.answer = answer
+
+        return form, branch if return_branch else form
+
+    async def _extend(
+        self,
+        tools,
+        reason,
+        predict,
+        score,
+        select,
+        plan,
+        # image,
+        allow_action,
+        confidence,
+        score_num_digits,
+        score_range,
+        select_choices,
+        predict_num_sentences,
+        **kwargs,
+    ):
+        """
+        Handles the extension of the form based on the provided parameters.
+
+        Args:
+            form: Form data.
+            tools: Tools data.
+            reason: Flag indicating if reason should be included.
+            predict: Flag indicating if prediction should be included.
+            score: Flag indicating if score should be included.
+            select: Flag indicating if selection should be included.
+            plan: Flag indicating if plan should be included.
+            allow_action: Flag indicating if action should be allowed.
+            confidence: Flag indicating if confidence should be included.
+            score_num_digits: Number of digits for score.
+            score_range: Range for score.
+            select_choices: Choices for selection.
+            predict_num_sentences: Number of sentences for prediction.
+            allow_extension: Flag indicating if extension should be allowed.
+            max_extension: Maximum extension value.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            list: The extended forms.
+        """
+        extension_forms = []
+
+        # Ensure the next step in the plan is handled
+        directive_kwargs = {
+            "tools": tools,
+            "reason": reason,
+            "predict": predict,
+            "score": score,
+            "select": select,
+            "allow_action": allow_action,
+            "confidence": confidence,
+            "score_num_digits": score_num_digits,
+            "score_range": score_range,
+            "select_choices": select_choices,
+            "predict_num_sentences": predict_num_sentences,
+            **kwargs,
+        }
+
+        if plan:
+            keys = [f"step_{i+1}" for i in range(len(plan))]
+            plan = StringMatch.force_validate_dict(plan, keys)
+
+            # If plan is provided, process each step
+            for i in keys:
+                directive_kwargs["instruction"] = plan[i]
+                last_form = await self._direct(**directive_kwargs)
+                last_form.is_extension = True
+                extension_forms.append(last_form)
+                directive_kwargs["max_extension"] -= 1
+                if not getattr(last_form, "extension_required", None):
+                    break
+
+        else:
+            # Handle single step extension
+            last_form = await self._direct(**directive_kwargs)
+            last_form.is_extension = True
+            extension_forms.append(last_form)
+
+        return extension_forms
+
+    async def _act(self, form, branch, actions=None):
+        """
+        Processes actions based on the provided form and actions.
+
+        Args:
+            form: Form data.
+            branch: Branch instance.
+            actions: Actions data.
+
+        Returns:
+            dict: The updated form.
+        """
+        if getattr(form, "action_performed", None) is True:
+            return form
+
+        keys = [f"action_{i+1}" for i in range(len(actions))]
+        actions = StringMatch.force_validate_dict(actions, keys)
+
+        try:
+            requests = []
+            for k in keys:
+                _func = actions[k]["function"]
+                _func = _func.replace("functions.", "")
+                msg = ActionRequest(
+                    function=_func,
+                    arguments=actions[k]["arguments"],
+                    sender=branch.ln_id,
+                    recipient=branch.tool_manager.registry[_func].ln_id,
+                )
+                requests.append(msg)
+                branch.add_message(action_request=msg)
+
+            if requests:
+                out = await self._process_action_request(
+                    branch=branch, invoke_tool=True, action_request=requests
+                )
+
+                if out is False:
+                    raise ValueError("No requests found.")
+
+                len_actions = len(actions)
+                action_responses = [
+                    i
+                    for i in branch.messages[-len_actions:]
+                    if isinstance(i, ActionResponse)
+                ]
+
+                _action_responses = {}
+                for idx, item in enumerate(action_responses):
+                    _action_responses[f"action_{idx+1}"] = item._to_dict()
+
+                form.append_to_request("action_response")
+                if (a := getattr(form, "action_response", None)) is None:
+                    form.add_field("action_response", {})
+
+                len1 = len(form.action_response)
+                for k, v in _action_responses.items():
+                    while k in form.action_response:
+                        k = f"{k}_1"
+                    form.action_response[k] = v
+
+                if len(form.action_response) > len1:
+                    form.append_to_request("action_performed")
+                    form.action_performed = True
+                return form
+
+        except Exception as e:
+            raise ValueError(f"Error processing action request: {e}")
+
+    async def _select(
+        self,
+        form=None,
+        choices=None,
+        reason=False,
+        confidence_score=None,
+        instruction=None,
+        template=None,
+        context=None,
+        branch=None,
+        **kwargs,
+    ):
+        """
+        Selects a response based on the provided parameters.
+
+        Args:
+            form (Any, optional): Form to create instruction from.
+            choices (Any, optional): Choices for the selection.
+            reason (bool, optional): Whether to include a reason for the selection.
+            confidence_score (Any, optional): Confidence score for the selection.
+            instruction (Any, optional): Instruction for the selection.
+            template (Any, optional): Template for the selection.
+            context (Any, optional): Context to perform the selection on.
+            branch (Any, optional): Branch to use for the selection.
+            **kwargs: Additional arguments for the selection.
+
+        Returns:
+            Any: The selection response.
+        """
+        branch = branch or self.branch
+
+        if not form:
+            form = template(
+                choices=choices,
+                reason=reason,
+                confidence_score=confidence_score,
+                instruction=instruction,
+                context=context,
+            )
+
+        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+
+    async def _predict(
+        self,
+        form=None,
+        num_sentences=None,
+        reason=False,
+        confidence_score=None,
+        instruction=None,
+        context=None,
+        branch=None,
+        template=None,
+        **kwargs,
+    ):
+        """
+        Predicts a response based on the provided parameters.
+
+        Args:
+            form: Form data.
+            num_sentences: Number of sentences for the prediction.
+            reason: Flag indicating if reason should be included.
+            confidence_score: Confidence score for the prediction.
+            instruction: Instruction for the prediction.
+            context: Context to perform the prediction on.
+            branch: Branch instance.
+            template: Template for the prediction.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The prediction response.
+        """
+        branch = branch or self.branch
+
+        if not form:
+            form = template(
+                instruction=instruction,
+                context=context,
+                num_sentences=num_sentences,
+                confidence_score=confidence_score,
+                reason=reason,
+            )
+
+        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+
+    async def _score(
+        self,
+        form=None,
+        score_range=None,
+        include_endpoints=None,
+        num_digit=None,
+        reason=False,
+        confidence_score=None,
+        instruction=None,
+        context=None,
+        branch=None,
+        template=None,
+        **kwargs,
+    ):
+        """
+        Scores a response based on the provided parameters.
+
+        Args:
+            form: Form data.
+            score_range: Range for score.
+            include_endpoints: Flag indicating if endpoints should be included.
+            num_digit: Number of digits for score.
+            reason: Flag indicating if reason should be included.
+            confidence_score: Confidence score for the score.
+            instruction: Instruction for the score.
+            context: Context to perform the score on.
+            branch: Branch instance.
+            template: Template for the score.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The score response.
+        """
+        branch = branch or self.branch
+        if not form:
+            form = template(
+                score_range=score_range,
+                include_endpoints=include_endpoints,
+                num_digit=num_digit,
+                reason=reason,
+                confidence_score=confidence_score,
+                instruction=instruction,
+                context=context,
+            )
+
+        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+
+    async def _plan(
+        self,
+        form=None,
+        num_step=None,
+        reason=False,
+        confidence_score=None,
+        instruction=None,
+        context=None,
+        branch=None,
+        template=None,
+        **kwargs,
+    ):
+        """
+        Plans a response based on the provided parameters.
+
+        Args:
+            form: Form data.
+            num_step: Number of steps for the plan.
+            reason: Flag indicating if reason should be included.
+            confidence_score: Confidence score for the plan.
+            instruction: Instruction for the plan.
+            context: Context to perform the plan on.
+            branch: Branch instance.
+            template: Template for the plan.
+            kwargs: Additional keyword arguments.
+
+        Returns:
+            Any: The plan response.
+        """
+        branch = branch or self.branch
+        template = template or self.default_template
+
+        if not form:
+            form = template(
+                instruction=instruction,
+                context=context,
+                num_step=num_step,
+                reason=reason,
+                confidence_score=confidence_score,
+            )
+
+        return await self._chat(form=form, **kwargs)
+
+    @staticmethod
+    def _process_model_response(content_, requested_fields):
+        """
+        Processes the model response content.
+
+        Args:
+            content_: The content data.
+            requested_fields: Fields requested in the response.
+
+        Returns:
+            Any: The processed response.
+        """
+        out_ = content_.get("content", "")
+        if out_ == "":
+            out_ = content_
+
+        if requested_fields:
+            with contextlib.suppress(Exception):
+                return StringMatch.force_validate_dict(out_, requested_fields)
+
+        if isinstance(out_, str):
+            with contextlib.suppress(Exception):
+                return ParseUtil.fuzzy_parse_json(out_)
+
+            with contextlib.suppress(Exception):
+                return ParseUtil.extract_json_block(out_)
+
+            with contextlib.suppress(Exception):
+                match = re.search(r"```json\n({.*?})\n```", out_, re.DOTALL)
+                if match:
+                    return ParseUtil.fuzzy_parse_json(match.group(1))
+
+        return out_