lionagi 0.3.8__py3-none-any.whl → 0.5.0__py3-none-any.whl

lionagi/operations/brainstorm/brainstorm.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from lionagi.core.session.branch import Branch
+from lionagi.core.session.session import Session
+from lionagi.core.typing import ID, Any, BaseModel
+from lionagi.libs.func import alcall
+from lionagi.libs.parse import to_flat_list
+from lionagi.protocols.operatives.instruct import (
+    INSTRUCT_MODEL_FIELD,
+    Instruct,
+    InstructResponse,
+)
+
+from ..utils import prepare_instruct, prepare_session
+from .prompt import PROMPT
+
+
+class BrainstormOperation(BaseModel):
+    initial: Any
+    brainstorm: list[Instruct] | None = None
+    explore: list[InstructResponse] | None = None
+
+
+async def run_instruct(
+    ins: Instruct,
+    session: Session,
+    branch: Branch,
+    auto_run: bool,
+    verbose: bool = True,
+    **kwargs: Any,
+) -> Any:
+    """Execute an instruction within a brainstorming session.
+
+    Args:
+        ins: The instruction model to run.
+        session: The current session.
+        branch: The branch to operate on.
+        auto_run: Whether to automatically run nested instructions.
+        verbose: Whether to enable verbose output.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        The result of the instruction execution.
+    """
+
+    async def run(ins_):
+        if verbose:
+            msg_ = (
+                ins_.guidance[:100] + "..."
+                if len(ins_.guidance) > 100
+                else ins_.guidance
+            )
+            print(f"\n-----Running instruction-----\n{msg_}")
+        b_ = session.split(branch)
+        return await run_instruct(
+            ins_, session, b_, False, verbose=verbose, **kwargs
+        )
+
+    config = {**ins.model_dump(), **kwargs}
+    res = await branch.operate(**config)
+    branch.msgs.logger.dump()
+    instructs = []
+
+    if hasattr(res, "instruct_models"):
+        instructs = res.instruct_models
+
+    if auto_run is True and instructs:
+        ress = await alcall(instructs, run)
+        response_ = []
+        for res in ress:
+            if isinstance(res, list):
+                response_.extend(res)
+            else:
+                response_.append(res)
+        response_.insert(0, res)
+        return response_
+
+    return res
+
+
+async def brainstorm(
+    instruct: Instruct | dict[str, Any],
+    num_instruct: int = 2,
+    session: Session | None = None,
+    branch: Branch | ID.Ref | None = None,
+    auto_run: bool = True,
+    auto_explore: bool = False,
+    explore_kwargs: dict[str, Any] | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    return_session: bool = False,
+    verbose: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """Perform a brainstorming session.
+
+    Args:
+        instruct: Instruction model or dictionary.
+        num_instruct: Number of instructions to generate.
+        session: Existing session or None to create a new one.
+        branch: Existing branch or reference.
+        auto_run: If True, automatically run generated instructions.
+        branch_kwargs: Additional arguments for branch creation.
+        return_session: If True, return the session with results.
+        verbose: Whether to enable verbose output.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        The results of the brainstorming session, optionally with the session.
+    """
+
+    if auto_explore and not auto_run:
+        raise ValueError("auto_explore requires auto_run to be True.")
+
+    if verbose:
+        print(f"Starting brainstorming...")
+
+    field_models: list = kwargs.get("field_models", [])
+    if INSTRUCT_MODEL_FIELD not in field_models:
+        field_models.append(INSTRUCT_MODEL_FIELD)
+
+    kwargs["field_models"] = field_models
+    session, branch = prepare_session(session, branch, branch_kwargs)
+    instruct = prepare_instruct(
+        instruct, PROMPT.format(num_instruct=num_instruct)
+    )
+    res1 = await branch.operate(**instruct, **kwargs)
+    out = BrainstormOperation(initial=res1)
+
+    if verbose:
+        print("Initial brainstorming complete.")
+
+    instructs = None
+
+    async def run(ins_):
+        if verbose:
+            msg_ = (
+                ins_.guidance[:100] + "..."
+                if len(ins_.guidance) > 100
+                else ins_.guidance
+            )
+            print(f"\n-----Running instruction-----\n{msg_}")
+        b_ = session.split(branch)
+        return await run_instruct(
+            ins_, session, b_, auto_run, verbose=verbose, **kwargs
+        )
+
+    if not auto_run:
+        if return_session:
+            return out, session
+        return out
+
+    async with session.branches:
+        response_ = []
+        if hasattr(res1, "instruct_models"):
+            instructs: list[Instruct] = res1.instruct_models
+            ress = await alcall(instructs, run)
+            ress = to_flat_list(ress, dropna=True)
+
+            response_ = [
+                res if not isinstance(res, str | dict) else None
+                for res in ress
+            ]
+            response_ = to_flat_list(response_, unique=True, dropna=True)
+            out.brainstorm = (
+                response_ if isinstance(response_, list) else [response_]
+            )
+            response_.insert(0, res1)
+
+        if response_ and auto_explore:
+
+            async def explore(ins_: Instruct):
+                if verbose:
+                    msg_ = (
+                        ins_.guidance[:100] + "..."
+                        if len(ins_.guidance) > 100
+                        else ins_.guidance
+                    )
+                    print(f"\n-----Exploring Idea-----\n{msg_}")
+                b_ = session.split(branch)
+                res = await b_.instruct(ins_, **(explore_kwargs or {}))
+                return InstructResponse(
+                    instruct=ins_,
+                    response=res,
+                )
+
+            response_ = to_flat_list(
+                [
+                    i.instruct_models
+                    for i in response_
+                    if hasattr(i, "instruct_models")
+                ],
+                dropna=True,
+                unique=True,
+            )
+            res_explore = await alcall(response_, explore)
+            out.explore = res_explore
+
+    if return_session:
+        return out, session
+
+    return out

lionagi/operations/brainstorm/prompt.py

@@ -0,0 +1 @@
+PROMPT = """Perform a brainstorm session. Fill in {num_instruct} Instruct for the appropriate next step, we will run them separately and concurrently with same external context, but you should supplement each idea with certain amount of uniqueness while adhering to the guidelines and standards of the project. The Instruct should be concisely informational. If you think a particular step requries further extension, you should mention it in the instruct"""

lionagi/operations/plan/__init__.py

@@ -0,0 +1,3 @@
+from .plan import plan
+
+__all__ = ["plan"]

lionagi/operations/plan/plan.py

@@ -0,0 +1,172 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from lionagi.core.session.branch import Branch
+from lionagi.core.session.session import Session
+from lionagi.core.typing import ID, Any, BaseModel, Literal
+from lionagi.protocols.operatives.instruct import (
+    INSTRUCT_MODEL_FIELD,
+    Instruct,
+    InstructResponse,
+)
+
+from ..utils import prepare_instruct, prepare_session
+from .prompt import EXPANSION_PROMPT, PLAN_PROMPT
+
+
+class PlanOperation(BaseModel):
+    initial: Any
+    plan: list[Instruct] | None = None
+    execute: list[InstructResponse] | None = None
+
+
+async def run_step(
+    ins: Instruct,
+    session: Session,
+    branch: Branch,
+    verbose: bool = True,
+    **kwargs: Any,
+) -> Any:
+    """Execute a single step of the plan.
+
+    Args:
+        ins: The instruction model for the step.
+        session: The current session.
+        branch: The branch to operate on.
+        verbose: Whether to enable verbose output.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        The result of the branch operation.
+    """
+    if verbose:
+        instruction = (
+            ins.instruction[:100] + "..."
+            if len(ins.instruction) > 100
+            else ins.instruction
+        )
+        print(f"Further planning: {instruction}")
+
+    config = {**ins.model_dump(), **kwargs}
+    guide = config.pop("guidance", "")
+    config["guidance"] = EXPANSION_PROMPT + "\n" + str(guide)
+
+    res = await branch.operate(**config)
+    branch.msgs.logger.dump()
+    return res
+
+
+async def plan(
+    instruct: Instruct | dict[str, Any],
+    num_steps: int = 2,
+    session: Session | None = None,
+    branch: Branch | ID.Ref | None = None,
+    auto_run: bool = True,
+    auto_execute: bool = False,
+    execution_strategy: Literal["sequential"] = "sequential",
+    execution_kwargs: dict[str, Any] | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    return_session: bool = False,
+    verbose: bool = True,
+    **kwargs: Any,
+) -> PlanOperation | tuple[list[InstructResponse], Session]:
+    """Create and execute a multi-step plan.
+
+    Args:
+        instruct: Instruction model or dictionary.
+        num_steps: Number of steps in the plan.
+        session: Existing session or None to create a new one.
+        branch: Existing branch or reference.
+        auto_run: If True, automatically run the steps.
+        branch_kwargs: Additional keyword arguments for branch creation.
+        return_session: If True, return the session along with results.
+        verbose: Whether to enable verbose output.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        Results of the plan execution, optionally with the session.
+    """
+    if num_steps > 5:
+        raise ValueError("Number of steps must be 5 or less")
+
+    if verbose:
+        print(f"Planning execution with {num_steps} steps...")
+
+    field_models: list = kwargs.get("field_models", [])
+    if INSTRUCT_MODEL_FIELD not in field_models:
+        field_models.append(INSTRUCT_MODEL_FIELD)
+    kwargs["field_models"] = field_models
+    session, branch = prepare_session(session, branch, branch_kwargs)
+    execute_branch: Branch = session.split(branch)
+    instruct = prepare_instruct(
+        instruct, PLAN_PROMPT.format(num_steps=num_steps)
+    )
+
+    res1 = await branch.operate(**instruct, **kwargs)
+    out = PlanOperation(initial=res1)
+
+    if verbose:
+        print("Initial planning complete. Starting step planning...")
+
+    if not auto_run:
+        if return_session:
+            return res1, session
+        return res1
+
+    results = []
+    if hasattr(res1, "instruct_models"):
+        instructs: list[Instruct] = res1.instruct_models
+        for i, ins in enumerate(instructs, 1):
+            if verbose:
+                print(f"\n----- Planning step {i}/{len(instructs)} -----")
+            res = await run_step(
+                ins, session, branch, verbose=verbose, **kwargs
+            )
+            results.append(res)
+
+        if verbose:
+            print("\nAll planning completed successfully!")
+
+    all_plans = []
+    for res in results:
+        if hasattr(res, "instruct_models"):
+            for i in res.instruct_models:
+                if i and i not in all_plans:
+                    all_plans.append(i)
+    out.plan = all_plans
+
+    if auto_execute:
+        if verbose:
+            print("\nStarting execution of all steps...")
+        results = []
+        match execution_strategy:
+            case "sequential":
+                for i, ins in enumerate(all_plans, 1):
+                    if verbose:
+                        print(
+                            f"\n------ Executing step {i}/{len(all_plans)} ------"
+                        )
+                        msg = (
+                            ins.instruction[:100] + "..."
+                            if len(ins.instruction) > 100
+                            else ins.instruction
+                        )
+                        print(f"Instruction: {msg}")
+                    res = await execute_branch.instruct(
+                        ins, **(execution_kwargs or {})
+                    )
+                    res_ = InstructResponse(instruct=ins, response=res)
+                    results.append(res_)
+                out.execute = results
+                if verbose:
+                    print("\nAll steps executed successfully!")
+            case _:
+                raise ValueError(
+                    f"Invalid execution strategy: {execution_strategy}"
+                )
+
+    if return_session:
+        return out, session
+    return out

lionagi/operations/plan/prompt.py

@@ -0,0 +1,21 @@
+PLAN_PROMPT = """
+Develop a high-level plan with {num_steps} distinct steps. Each step should:
+1. Represent a major milestone or phase
+2. Be logically sequenced for dependencies
+3. Be clearly distinct from other steps
+4. Have measurable completion criteria
+5. Be suitable for further decomposition
+"""
+
+EXPANSION_PROMPT = """
+Break down a high-level plan into detailed concrete executable actions. Each step should:
+- Ensure actions are atomic and verifiable
+- Include necessary context and preconditions
+- Specify expected outcomes and validations
+- Maintain sequential dependencies
+- Be self-contained with clear scope
+- Include all required context/parameters
+- Have unambiguous success criteria
+- Specify error handling approach
+- Define expected outputs
+"""

lionagi/operations/select/__init__.py

@@ -0,0 +1,3 @@
+from .select import select
+
+__all__ = ["select"]

lionagi/operations/select/prompt.py

@@ -0,0 +1 @@
+PROMPT = "Please select up to {max_num_selections} items from the following list {choices}. Provide the selection(s) into appropriate field in format required, and no comments from you"

lionagi/operations/select/select.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from lionagi import Branch
+from lionagi.protocols.operatives.instruct import Instruct
+
+from .prompt import PROMPT
+from .utils import parse_selection, parse_to_representation
+
+
+class SelectionModel(BaseModel):
+    """Model representing the selection output."""
+
+    selected: list[Any] = Field(default_factory=list)
+
+
+async def select(
+    instruct: Instruct | dict[str, Any],
+    choices: list[str] | type[Enum] | dict[str, Any],
+    max_num_selections: int = 1,
+    branch: Branch | None = None,
+    branch_kwargs: dict[str, Any] | None = None,
+    return_branch: bool = False,
+    verbose: bool = False,
+    **kwargs: Any,
+) -> SelectionModel | tuple[SelectionModel, Branch]:
+    """Perform a selection operation from given choices.
+
+    Args:
+        instruct: Instruction model or dictionary.
+        choices: Options to select from.
+        max_num_selections: Maximum selections allowed.
+        branch: Existing branch or None to create a new one.
+        branch_kwargs: Additional arguments for branch creation.
+        return_branch: If True, return the branch with the selection.
+        verbose: Whether to enable verbose output.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        A SelectionModel instance, optionally with the branch.
+    """
+    if verbose:
+        print(f"Starting selection with up to {max_num_selections} choices.")
+
+    branch = branch or Branch(**(branch_kwargs or {}))
+    selections, contents = parse_to_representation(choices)
+    prompt = PROMPT.format(
+        max_num_selections=max_num_selections, choices=selections
+    )
+
+    if isinstance(instruct, Instruct):
+        instruct = instruct.to_dict()
+
+    instruct = instruct or {}
+
+    if instruct.get("instruction", None) is not None:
+        instruct["instruction"] = (
+            f"{instruct['instruction']}\n\n{prompt} \n\n "
+        )
+    else:
+        instruct["instruction"] = prompt
+
+    context = instruct.get("context", None) or []
+    context = [context] if not isinstance(context, list) else context
+    context.extend([{k: v} for k, v in zip(selections, contents)])
+    instruct["context"] = context
+
+    response_model: SelectionModel = await branch.operate(
+        operative_model=SelectionModel,
+        **kwargs,
+        **instruct,
+    )
+    if verbose:
+        print(f"Received selection: {response_model.selected}")
+
+    selected = response_model
+    if isinstance(response_model, BaseModel) and hasattr(
+        response_model, "selected"
+    ):
+        selected = response_model.selected
+    selected = [selected] if not isinstance(selected, list) else selected
+
+    corrected_selections = [parse_selection(i, choices) for i in selected]
+
+    if isinstance(response_model, BaseModel):
+        response_model.selected = corrected_selections
+
+    elif isinstance(response_model, dict):
+        response_model["selected"] = corrected_selections
+
+    if return_branch:
+        return response_model, branch
+    return response_model

lionagi/operations/select/utils.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import inspect
+from enum import Enum
+from typing import Any
+
+from lionagi.core.typing import BaseModel, JsonValue
+from lionagi.libs.parse import is_same_dtype
+from lionagi.libs.string_similarity import string_similarity
+
+
+def parse_to_representation(
+    choices: Enum | dict | list | tuple | set,
+) -> tuple[list[str], JsonValue]:
+    """
+    should use
+    1. iterator of string | BaseModel
+    2. dict[str, JsonValue | BaseModel]
+    3. Enum[str, JsonValue | BaseModel]
+    """
+
+    if isinstance(choices, tuple | set | list):
+        choices = list(choices)
+        if is_same_dtype(choices, str):
+            return choices, choices
+
+    if isinstance(choices, list):
+        if is_same_dtype(choices, BaseModel):
+            choices = {i.__class__.__name__: i for i in choices}
+        if all(
+            inspect.isclass(i) and issubclass(i, BaseModel) for i in choices
+        ):
+            choices = {i.__name__: i for i in choices}
+    if isinstance(choices, type) and issubclass(choices, Enum):
+        keys = [i.name for i in choices]
+        contents = [get_choice_representation(i) for i in choices]
+        return keys, contents
+
+    if isinstance(choices, dict):
+        keys = list(choices.keys())
+        contents = list(choices.values())
+        contents = [get_choice_representation(v) for k, v in choices.items()]
+        return keys, contents
+
+    if isinstance(choices, tuple | set | list):
+        choices = list(choices)
+        if is_same_dtype(choices, str):
+            return choices, choices
+
+    raise NotImplementedError
+
+
+def get_choice_representation(choice: Any) -> str:
+
+    if isinstance(choice, str):
+        return choice
+
+    if isinstance(choice, BaseModel):
+        return f"{choice.__class__.__name__}:\n{choice.model_json_schema(indent=2)}"
+
+    if isinstance(choice, Enum):
+        return get_choice_representation(choice.value)
+
+
+def parse_selection(selection_str: str, choices: Any):
+
+    select_from = []
+
+    if isinstance(choices, dict):
+        select_from = list(choices.keys())
+
+    if inspect.isclass(choices) and issubclass(choices, Enum):
+        select_from = [choice.name for choice in choices]
+
+    if isinstance(choices, list | tuple | set):
+        if is_same_dtype(choices, BaseModel):
+            select_from = [i.__class__.__name__ for i in choices]
+        if is_same_dtype(choices, str):
+            select_from = list(choices)
+        if all(
+            inspect.isclass(i) and issubclass(i, BaseModel) for i in choices
+        ):
+            select_from = [i.__name__ for i in choices]
+
+    if not select_from:
+        raise ValueError("The values provided for choice is not valid")
+
+    selected = string_similarity(
+        selection_str, select_from, return_most_similar=True
+    )
+
+    if isinstance(choices, dict) and selected in choices:
+        return choices[selected]
+
+    if inspect.isclass(choices) and issubclass(choices, Enum):
+        for i in choices:
+            if i.name == selected:
+                return i
+
+    if isinstance(choices, list) and is_same_dtype(choices, str):
+        if selected in choices:
+            return selected
+
+    return selected

lionagi/operations/utils.py

@@ -0,0 +1,35 @@
+from lionagi.core.session.branch import Branch
+from lionagi.core.session.session import Session
+from lionagi.protocols.operatives.instruct import Instruct
+
+
+def prepare_session(
+    session=None, branch=None, branch_kwargs=None
+) -> tuple[Session, Branch]:
+    if session is not None:
+        if branch is not None:
+            branch: Branch = session.branches[branch]
+        else:
+            branch = session.new_branch(**(branch_kwargs or {}))
+    else:
+        session = Session()
+        if isinstance(branch, Branch):
+            session.branches.include(branch)
+            session.default_branch = branch
+        if branch is None:
+            branch = session.new_branch(**(branch_kwargs or {}))
+
+    return session, branch
+
+
+def prepare_instruct(instruct: Instruct | dict, prompt: str):
+    if isinstance(instruct, Instruct):
+        instruct = instruct.to_dict()
+    if not isinstance(instruct, dict):
+        raise ValueError(
+            "instruct needs to be an InstructModel object or a dictionary of valid parameters"
+        )
+
+    guidance = instruct.get("guidance", "")
+    instruct["guidance"] = f"\n{prompt}\n{guidance}"
+    return instruct