ellements 0.2.0__py3-none-any.whl

ellements/execution/collaborative.py

@@ -0,0 +1,191 @@
+"""Collaborative editing strategy — LLM generate → user edit → LLM continue.
+
+The LLM produces an initial draft, the user edits it through any
+:class:`EditCallback` implementation, and the LLM continues from the
+edited version. The loop ends when the user signals completion
+(via *done_signal*), leaves the content unchanged, returns empty
+content, or ``max_rounds`` is reached.
+
+Required prompts:
+    - ``generate``: produce the initial draft
+    - ``continue``: receives ``{{edited_content}}`` and ``{{original_content}}``
+
+Optional prompts:
+    - ``system``: persona/context forwarded to every LLM call
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+from ellements.core import LLMClient
+
+from .callbacks import EditCallback
+from .config import CollaborativeEditingConfig, StrategyConfigInput
+from .strategies import BaseStrategy, StepRecord, StrategyResult
+
+_logger = logging.getLogger(__name__)
+
+
+class CollaborativeEditingStrategy(BaseStrategy):
+    """LLM generate → user edit → LLM continue loop.
+
+    Config keys:
+        edit_callback (EditCallback): How to present content and receive edits.
+        max_rounds (int): Maximum edit rounds after the initial generation.
+        done_signal (str): If the user adds this string on its own line,
+            the strategy finishes immediately.
+        continue_on_unchanged (bool): If True, the LLM continues even
+            when the user returns content identical to the LLM's
+            previous output.
+        generate_temperature (float): Temperature for initial generation.
+        continue_temperature (float): Temperature for continuation.
+    """
+
+    REQUIRED_PROMPTS = ("generate", "continue")
+
+    async def execute(
+        self,
+        prompts: Mapping[str, str],
+        client: LLMClient,
+        tools: list[dict[str, Any]] | Mapping[str, Any] | None = None,
+        config: StrategyConfigInput = None,
+    ) -> StrategyResult:
+        typed_config = self._normalize_config(config, CollaborativeEditingConfig)
+        max_rounds = typed_config.max_rounds
+        done_signal = typed_config.done_signal
+        continue_on_unchanged = typed_config.continue_on_unchanged
+        callback: EditCallback = typed_config.edit_callback
+        system = prompts.get("system")
+        steps: list[StepRecord] = []
+
+        generate_temp = self._stage_temperature(
+            typed_config, "generate", default=typed_config.generate_temperature
+        )
+        continue_temp = self._stage_temperature(
+            typed_config, "continue", default=typed_config.continue_temperature
+        )
+
+        generate_prompt = self._get_prompt(prompts, "generate")
+        current = await self._complete(
+            client,
+            generate_prompt,
+            config=typed_config,
+            temperature=generate_temp,
+            system=system,
+            tools=tools,
+        )
+        steps.append(
+            StepRecord(
+                name="generate",
+                prompt_key="generate",
+                response=current,
+                metadata={"round": 0, "source": "llm"},
+            )
+        )
+        self._notify_step(typed_config, steps[-1])
+
+        continue_template = self._get_prompt(prompts, "continue")
+
+        for round_num in range(max_rounds):
+            original = current
+            context = f"Round {round_num + 1} of {max_rounds}"
+            edited = await callback.request_edit(current, context)
+
+            if not edited.strip():
+                _logger.info(
+                    "Collaborative editing aborted by user (empty content)"
+                )
+                steps.append(
+                    StepRecord(
+                        name=f"abort_{round_num}",
+                        prompt_key="continue",
+                        response="",
+                        metadata={"round": round_num + 1, "reason": "user_abort"},
+                    )
+                )
+                self._notify_step(typed_config, steps[-1])
+                break
+
+            if _has_done_signal(edited, done_signal):
+                edited = _strip_done_signal(edited, done_signal)
+                steps.append(
+                    StepRecord(
+                        name=f"user_edit_{round_num}",
+                        prompt_key="continue",
+                        response=edited,
+                        metadata={
+                            "round": round_num + 1,
+                            "source": "user",
+                            "done_signal": True,
+                        },
+                    )
+                )
+                self._notify_step(typed_config, steps[-1])
+                current = edited
+                break
+
+            steps.append(
+                StepRecord(
+                    name=f"user_edit_{round_num}",
+                    prompt_key="continue",
+                    response=edited,
+                    metadata={"round": round_num + 1, "source": "user"},
+                )
+            )
+            self._notify_step(typed_config, steps[-1])
+
+            if not continue_on_unchanged and edited.strip() == original.strip():
+                _logger.info(
+                    "Collaborative editing done: user approved without changes"
+                )
+                current = edited
+                break
+
+            continue_prompt = self._render_template(
+                continue_template,
+                {"edited_content": edited, "original_content": original},
+            )
+            current = await self._complete(
+                client,
+                continue_prompt,
+                config=typed_config,
+                temperature=continue_temp,
+                system=system,
+                tools=tools,
+            )
+            steps.append(
+                StepRecord(
+                    name=f"continue_{round_num}",
+                    prompt_key="continue",
+                    response=current,
+                    metadata={"round": round_num + 1, "source": "llm"},
+                )
+            )
+            self._notify_step(typed_config, steps[-1])
+
+        rounds_completed = len(
+            [s for s in steps if s.name.startswith("user_edit_")]
+        )
+        return StrategyResult(
+            output=current,
+            steps=steps,
+            metadata={
+                "rounds_completed": rounds_completed,
+                "max_rounds": max_rounds,
+            },
+        )
+
+
+def _has_done_signal(content: str, signal: str) -> bool:
+    return any(line.strip() == signal for line in content.splitlines())
+
+
+def _strip_done_signal(content: str, signal: str) -> str:
+    lines = [line for line in content.splitlines() if line.strip() != signal]
+    return "\n".join(lines).strip()
+
+
+__all__ = ["CollaborativeEditingStrategy"]

ellements/execution/config.py

@@ -0,0 +1,135 @@
+"""Typed config models for execution strategies.
+
+All built-in strategies accept a :class:`ExecutionStrategyConfig` (or a
+plain dict / pydantic BaseModel that maps onto one). The shared base
+exposes ``model``, ``on_step``, and the ``stage_temperatures`` override
+map; individual strategy configs add their own knobs.
+
+Unknown keys are ignored so callers can mix in extra metadata without
+breaking validation.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from typing import Any, Literal, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .callbacks import EditCallback, PassthroughEditCallback
+
+OnStepCallback = Callable[[Any], None]
+StrategyConfigInput = Mapping[str, Any] | BaseModel | None
+
+
+class ExecutionStrategyConfig(BaseModel):
+    """Shared config fields accepted by every built-in execution strategy."""
+
+    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)
+
+    model: str | None = None
+    """Override the LLMClient's default model for every call in this strategy."""
+
+    on_step: OnStepCallback | None = None
+    """Callback invoked after each StepRecord is produced."""
+
+    stage_temperatures: dict[str, float] | None = None
+    """Per-stage temperature overrides keyed by stage name.
+
+    When a stage is listed here, its value wins over the matching
+    ``{stage}_temperature`` field on the typed config. Unknown stage
+    names are ignored.
+    """
+
+
+class SingleCallConfig(ExecutionStrategyConfig):
+    """Typed config for :class:`SingleCallStrategy`."""
+
+    temperature: float | None = None
+
+
+class SelfConsistencyConfig(ExecutionStrategyConfig):
+    """Typed config for :class:`SelfConsistencyStrategy`."""
+
+    samples: int = 5
+    sample_temperature: float = 0.9
+    judge_temperature: float = 0.1
+    aggregation: Literal["majority_vote", "llm_judge"] = "majority_vote"
+    judge_model: str | None = None
+    max_concurrent: int = 3
+
+
+class TreeOfThoughtConfig(ExecutionStrategyConfig):
+    """Typed config for :class:`TreeOfThoughtStrategy`.
+
+    ``mode`` selects the search algorithm:
+
+    - ``"beam"``: BFS with step-level beam search (canonical Yao et al. 2023).
+    - ``"dfs"``: depth-first with backtracking + early exit on max score.
+    - ``"simple"``: simplified flat search (generate N → evaluate → synthesize).
+    """
+
+    mode: Literal["beam", "dfs", "simple"] = "beam"
+    max_depth: int = 3
+    branching_factor: int = 3
+    beam_width: int = 3
+    thought_temperature: float = 0.9
+    evaluate_temperature: float = 0.2
+    synthesize_temperature: float = 0.3
+    max_concurrent: int = 2
+
+
+class ReflectionConfig(ExecutionStrategyConfig):
+    """Typed config for :class:`ReflectionStrategy`."""
+
+    max_rounds: int = 3
+    generate_temperature: float = 0.7
+    critique_temperature: float = 0.3
+    revise_temperature: float = 0.4
+
+
+class CollaborativeEditingConfig(ExecutionStrategyConfig):
+    """Typed config for :class:`CollaborativeEditingStrategy`."""
+
+    max_rounds: int = 5
+    done_signal: str = "DONE"
+    continue_on_unchanged: bool = False
+    edit_callback: EditCallback = Field(default_factory=PassthroughEditCallback)
+    generate_temperature: float = 0.7
+    continue_temperature: float = 0.7
+
+
+ConfigModelT = TypeVar("ConfigModelT", bound=ExecutionStrategyConfig)
+
+
+def normalize_strategy_config(
+    config: StrategyConfigInput,
+    config_model: type[ConfigModelT],
+) -> ConfigModelT:
+    """Normalize a mixed config input to the typed *config_model*."""
+    if config is None:
+        return config_model()
+    if isinstance(config, config_model):
+        return config
+    if isinstance(config, BaseModel):
+        raw_config = config.model_dump(mode="python", exclude_none=False)
+    elif isinstance(config, Mapping):
+        raw_config = dict(config)
+    else:
+        raise TypeError(
+            "Strategy config must be a mapping, pydantic BaseModel, or None."
+        )
+    return config_model.model_validate(raw_config)
+
+
+__all__ = [
+    "CollaborativeEditingConfig",
+    "ExecutionStrategyConfig",
+    "OnStepCallback",
+    "ReflectionConfig",
+    "SelfConsistencyConfig",
+    "SingleCallConfig",
+    "StrategyConfigInput",
+    "TreeOfThoughtConfig",
+    "normalize_strategy_config",
+]

ellements/execution/reflection.py

@@ -0,0 +1,156 @@
+"""Reflection strategy — generate → critique → revise loop.
+
+Produces an initial response, then asks the model to critique it via a
+structured :class:`CritiqueResult`. When ``is_satisfied`` becomes true
+(or ``max_rounds`` is exhausted), the most recent revision is returned.
+
+Required prompts:
+    - ``generate``: produce the initial response
+    - ``critique``: receives ``response``; must return a CritiqueResult
+    - ``revise``: receives ``response`` and ``issues``
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+from ellements.core import LLMClient
+from pydantic import BaseModel, Field
+
+from .config import ReflectionConfig, StrategyConfigInput
+from .strategies import BaseStrategy, StepRecord, StrategyResult
+
+_logger = logging.getLogger(__name__)
+
+
+class CritiqueResult(BaseModel):
+    """Structured critique output produced by the ``critique`` stage."""
+
+    is_satisfied: bool = Field(
+        description="True if the response is acceptable as-is (no further revision needed)"
+    )
+    issues: list[str] = Field(
+        default_factory=list,
+        description="Specific issues the next revision should address",
+    )
+
+
+class ReflectionStrategy(BaseStrategy):
+    """Generate → critique → revise loop using a structured CritiqueResult."""
+
+    REQUIRED_PROMPTS = ("generate", "critique", "revise")
+
+    async def execute(
+        self,
+        prompts: Mapping[str, str],
+        client: LLMClient,
+        tools: list[dict[str, Any]] | Mapping[str, Any] | None = None,
+        config: StrategyConfigInput = None,
+    ) -> StrategyResult:
+        typed_config = self._normalize_config(config, ReflectionConfig)
+        system = prompts.get("system")
+        steps: list[StepRecord] = []
+
+        generate_temp = self._stage_temperature(
+            typed_config, "generate", default=typed_config.generate_temperature
+        )
+        critique_temp = self._stage_temperature(
+            typed_config, "critique", default=typed_config.critique_temperature
+        )
+        revise_temp = self._stage_temperature(
+            typed_config, "revise", default=typed_config.revise_temperature
+        )
+
+        generate_prompt = self._get_prompt(prompts, "generate")
+        current = await self._complete(
+            client,
+            generate_prompt,
+            config=typed_config,
+            temperature=generate_temp,
+            system=system,
+            tools=tools,
+        )
+        steps.append(
+            StepRecord(name="generate", prompt_key="generate", response=current)
+        )
+        self._notify_step(typed_config, steps[-1])
+
+        critique_template = self._get_prompt(prompts, "critique")
+        revise_template = self._get_prompt(prompts, "revise")
+
+        for round_num in range(typed_config.max_rounds):
+            critique_prompt = self._render_template(
+                critique_template, {"response": current}
+            )
+            critique = await self._complete_structured(
+                client,
+                critique_prompt,
+                response_model=CritiqueResult,
+                config=typed_config,
+                temperature=critique_temp,
+                system=system,
+            )
+            steps.append(
+                StepRecord(
+                    name=f"critique_{round_num}",
+                    prompt_key="critique",
+                    response=critique.model_dump_json(),
+                    metadata={
+                        "round": round_num,
+                        "is_satisfied": critique.is_satisfied,
+                        "issue_count": len(critique.issues),
+                    },
+                )
+            )
+            self._notify_step(typed_config, steps[-1])
+
+            if critique.is_satisfied:
+                stop_step = StepRecord(
+                    name="stop",
+                    prompt_key="critique",
+                    response=f"Stopped at round {round_num}: critique satisfied.",
+                    metadata={"round": round_num, "reason": "satisfied"},
+                )
+                steps.append(stop_step)
+                self._notify_step(typed_config, stop_step)
+                break
+
+            revise_prompt = self._render_template(
+                revise_template,
+                {
+                    "response": current,
+                    "issues": "\n".join(f"- {issue}" for issue in critique.issues),
+                },
+            )
+            current = await self._complete(
+                client,
+                revise_prompt,
+                config=typed_config,
+                temperature=revise_temp,
+                system=system,
+                tools=tools,
+            )
+            steps.append(
+                StepRecord(
+                    name=f"revise_{round_num}",
+                    prompt_key="revise",
+                    response=current,
+                    metadata={"round": round_num},
+                )
+            )
+            self._notify_step(typed_config, steps[-1])
+
+        revise_steps = [s for s in steps if s.name.startswith("revise_")]
+        return StrategyResult(
+            output=current,
+            steps=steps,
+            metadata={
+                "rounds_used": len(revise_steps) + 1,
+                "satisfied": any(s.name == "stop" for s in steps),
+            },
+        )
+
+
+__all__ = ["CritiqueResult", "ReflectionStrategy"]

ellements/execution/self_consistency.py

@@ -0,0 +1,189 @@
+"""Self-consistency strategy — sample N responses and aggregate.
+
+Generates ``samples`` parallel responses to the same prompt (controlled
+concurrency), then either picks the most common answer (majority vote)
+or asks the LLM to synthesize the best one (llm-judge).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from collections import Counter
+from collections.abc import Mapping
+from typing import Any
+
+from ellements.core import LLMClient
+
+from .config import SelfConsistencyConfig, StrategyConfigInput
+from .strategies import BaseStrategy, StepRecord, StrategyResult
+
+_logger = logging.getLogger(__name__)
+
+
+class SelfConsistencyStrategy(BaseStrategy):
+    """Run a prompt N times and aggregate the results.
+
+    Config keys:
+        samples (int): Number of samples. Default 5.
+        sample_temperature (float): Sampling temperature. Default 0.9.
+        aggregation (str): ``"majority_vote"`` or ``"llm_judge"``.
+        judge_temperature (float): Temperature for the judge call.
+        judge_model (str | None): Model override for the judge call.
+        max_concurrent (int): Max in-flight LLM calls.
+    """
+
+    REQUIRED_PROMPTS = ("default",)
+
+    async def execute(
+        self,
+        prompts: Mapping[str, str],
+        client: LLMClient,
+        tools: list[dict[str, Any]] | Mapping[str, Any] | None = None,
+        config: StrategyConfigInput = None,
+    ) -> StrategyResult:
+        typed_config = self._normalize_config(config, SelfConsistencyConfig)
+        prompt = self._get_prompt(prompts, "default")
+        system = prompts.get("system")
+
+        sample_temp = self._stage_temperature(
+            typed_config, "sample", default=typed_config.sample_temperature
+        )
+        judge_temp = self._stage_temperature(
+            typed_config, "judge", default=typed_config.judge_temperature
+        )
+
+        sem = asyncio.Semaphore(typed_config.max_concurrent)
+
+        async def _sample_one() -> str:
+            async with sem:
+                return await self._complete(
+                    client,
+                    prompt,
+                    config=typed_config,
+                    temperature=sample_temp,
+                    system=system,
+                    tools=tools,
+                )
+
+        responses = await asyncio.gather(
+            *(_sample_one() for _ in range(typed_config.samples))
+        )
+
+        steps: list[StepRecord] = [
+            StepRecord(
+                name=f"sample_{i}",
+                prompt_key="default",
+                response=resp,
+                metadata={"temperature": sample_temp},
+            )
+            for i, resp in enumerate(responses)
+        ]
+        for step in steps:
+            self._notify_step(typed_config, step)
+
+        if typed_config.aggregation == "llm_judge":
+            output = await self._llm_judge(
+                client,
+                prompt,
+                list(responses),
+                config=typed_config,
+                temperature=judge_temp,
+                system=system,
+            )
+            aggregation = "llm_judge"
+            agg_step = StepRecord(
+                name="judge",
+                prompt_key="default",
+                response=output,
+                metadata={"aggregation": aggregation, "temperature": judge_temp},
+            )
+        else:
+            output = self._majority_vote(list(responses))
+            aggregation = "majority_vote"
+            agg_step = StepRecord(
+                name="vote",
+                prompt_key="default",
+                response=output,
+                metadata={"aggregation": aggregation},
+            )
+
+        steps.append(agg_step)
+        self._notify_step(typed_config, agg_step)
+
+        return StrategyResult(
+            output=output,
+            steps=steps,
+            metadata={
+                "samples": typed_config.samples,
+                "aggregation": aggregation,
+            },
+        )
+
+    @staticmethod
+    def _majority_vote(responses: list[str]) -> str:
+        """Pick the most common final answer, falling back to full responses."""
+        displays: dict[tuple[str, str], str] = {}
+        counter: Counter[tuple[str, str]] = Counter()
+        for response in responses:
+            answer = SelfConsistencyStrategy._extract_final_answer(response)
+            if answer is None:
+                display = response.strip()
+                key = ("response", display)
+            else:
+                display = f"ANSWER: {answer}"
+                key = ("answer", answer.casefold())
+            displays.setdefault(key, display)
+            counter[key] += 1
+        winner, _ = counter.most_common(1)[0]
+        return displays[winner]
+
+    @staticmethod
+    def _extract_final_answer(response: str) -> str | None:
+        """Extract a final ``ANSWER: ...`` line when a prompt provides one."""
+        matches = re.findall(
+            r"(?im)^\s*(?:\*\*)?answer(?:\*\*)?\s*:\s*(.+?)\s*$",
+            response,
+        )
+        if not matches:
+            return None
+        answer = matches[-1].strip().strip("`")
+        answer = re.sub(r"\s+", " ", answer)
+        return answer.rstrip(".")
+
+    async def _llm_judge(
+        self,
+        client: LLMClient,
+        original_prompt: str,
+        responses: list[str],
+        *,
+        config: SelfConsistencyConfig,
+        temperature: float,
+        system: str | None,
+    ) -> str:
+        numbered = "\n\n".join(
+            f"--- Response {i + 1} ---\n{resp}"
+            for i, resp in enumerate(responses)
+        )
+        judge_prompt = (
+            "You were asked the following question/task:\n\n"
+            f"{original_prompt}\n\n"
+            f"Here are {len(responses)} independent responses:\n\n"
+            f"{numbered}\n\n"
+            "Synthesize the best answer by combining the strongest elements. "
+            "If the responses agree, use the consensus answer. If they "
+            "disagree, reason about which is most likely correct and explain "
+            "why.\n\n"
+            "Provide ONLY the final synthesized answer, nothing else."
+        )
+        judge_config = config.model_copy(
+            update={"model": config.judge_model or config.model}
+        )
+        return await self._complete(
+            client,
+            judge_prompt,
+            config=judge_config,
+            temperature=temperature,
+            system=system,
+        )