ace-framework 0.1.0__py3-none-any.whl

ace/__init__.py

@@ -0,0 +1,53 @@
+"""Agentic Context Engineering (ACE) reproduction framework."""
+
+from .playbook import Bullet, Playbook
+from .delta import DeltaOperation, DeltaBatch
+from .llm import LLMClient, DummyLLMClient, TransformersLLMClient
+from .roles import (
+    Generator,
+    Reflector,
+    Curator,
+    GeneratorOutput,
+    ReflectorOutput,
+    CuratorOutput,
+)
+from .adaptation import (
+    OfflineAdapter,
+    OnlineAdapter,
+    Sample,
+    TaskEnvironment,
+    EnvironmentResult,
+    AdapterStepResult,
+)
+
+# Import production LLM clients if available
+try:
+    from .llm_providers import LiteLLMClient
+    LITELLM_AVAILABLE = True
+except ImportError:
+    LiteLLMClient = None
+    LITELLM_AVAILABLE = False
+
+__all__ = [
+    "Bullet",
+    "Playbook",
+    "DeltaOperation",
+    "DeltaBatch",
+    "LLMClient",
+    "DummyLLMClient",
+    "TransformersLLMClient",
+    "LiteLLMClient",
+    "Generator",
+    "Reflector",
+    "Curator",
+    "GeneratorOutput",
+    "ReflectorOutput",
+    "CuratorOutput",
+    "OfflineAdapter",
+    "OnlineAdapter",
+    "Sample",
+    "TaskEnvironment",
+    "EnvironmentResult",
+    "AdapterStepResult",
+    "LITELLM_AVAILABLE",
+]

ace/adaptation.py

@@ -0,0 +1,193 @@
+"""Adaptation loops for offline and online ACE training."""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, List, Optional, Sequence
+
+from .playbook import Playbook
+from .roles import Curator, CuratorOutput, Generator, GeneratorOutput, Reflector, ReflectorOutput
+
+
+@dataclass
+class Sample:
+    """Single task instance presented to ACE."""
+
+    question: str
+    context: str = ""
+    ground_truth: Optional[str] = None
+    metadata: Dict[str, object] = field(default_factory=dict)
+
+
+@dataclass
+class EnvironmentResult:
+    """Feedback returned by the task environment after executing the generator output."""
+
+    feedback: str
+    ground_truth: Optional[str]
+    metrics: Dict[str, float] = field(default_factory=dict)
+
+
+class TaskEnvironment(ABC):
+    """Defines how to evaluate generator outputs for a sample."""
+
+    @abstractmethod
+    def evaluate(
+        self, sample: Sample, generator_output: GeneratorOutput
+    ) -> EnvironmentResult:
+        """Return environment feedback plus optional ground truth information."""
+
+
+@dataclass
+class AdapterStepResult:
+    sample: Sample
+    generator_output: GeneratorOutput
+    environment_result: EnvironmentResult
+    reflection: ReflectorOutput
+    curator_output: CuratorOutput
+    playbook_snapshot: str
+
+
+class AdapterBase:
+    """Shared orchestration logic for offline and online ACE adaptation."""
+
+    def __init__(
+        self,
+        *,
+        playbook: Optional[Playbook] = None,
+        generator: Generator,
+        reflector: Reflector,
+        curator: Curator,
+        max_refinement_rounds: int = 1,
+        reflection_window: int = 3,
+    ) -> None:
+        self.playbook = playbook or Playbook()
+        self.generator = generator
+        self.reflector = reflector
+        self.curator = curator
+        self.max_refinement_rounds = max_refinement_rounds
+        self.reflection_window = reflection_window
+        self._recent_reflections: List[str] = []
+
+    # ------------------------------------------------------------------ #
+    def _reflection_context(self) -> str:
+        return "\n---\n".join(self._recent_reflections)
+
+    def _update_recent_reflections(self, reflection: ReflectorOutput) -> None:
+        serialized = json.dumps(reflection.raw, ensure_ascii=False)
+        self._recent_reflections.append(serialized)
+        if len(self._recent_reflections) > self.reflection_window:
+            self._recent_reflections = self._recent_reflections[-self.reflection_window :]
+
+    def _apply_bullet_tags(self, reflection: ReflectorOutput) -> None:
+        for tag in reflection.bullet_tags:
+            try:
+                self.playbook.tag_bullet(tag.id, tag.tag)
+            except ValueError:
+                continue
+
+    def _question_context(self, sample: Sample, environment_result: EnvironmentResult) -> str:
+        parts = [
+            f"question: {sample.question}",
+            f"context: {sample.context}",
+            f"metadata: {json.dumps(sample.metadata)}",
+            f"feedback: {environment_result.feedback}",
+            f"ground_truth: {environment_result.ground_truth}",
+        ]
+        return "\n".join(parts)
+
+    def _progress_string(self, epoch: int, total_epochs: int, step: int, total_steps: int) -> str:
+        return f"epoch {epoch}/{total_epochs} · sample {step}/{total_steps}"
+
+    def _process_sample(
+        self,
+        sample: Sample,
+        environment: TaskEnvironment,
+        *,
+        epoch: int,
+        total_epochs: int,
+        step_index: int,
+        total_steps: int,
+    ) -> AdapterStepResult:
+        generator_output = self.generator.generate(
+            question=sample.question,
+            context=sample.context,
+            playbook=self.playbook,
+            reflection=self._reflection_context(),
+        )
+        env_result = environment.evaluate(sample, generator_output)
+        reflection = self.reflector.reflect(
+            question=sample.question,
+            generator_output=generator_output,
+            playbook=self.playbook,
+            ground_truth=env_result.ground_truth,
+            feedback=env_result.feedback,
+            max_refinement_rounds=self.max_refinement_rounds,
+        )
+        self._apply_bullet_tags(reflection)
+        self._update_recent_reflections(reflection)
+        curator_output = self.curator.curate(
+            reflection=reflection,
+            playbook=self.playbook,
+            question_context=self._question_context(sample, env_result),
+            progress=self._progress_string(epoch, total_epochs, step_index, total_steps),
+        )
+        self.playbook.apply_delta(curator_output.delta)
+        return AdapterStepResult(
+            sample=sample,
+            generator_output=generator_output,
+            environment_result=env_result,
+            reflection=reflection,
+            curator_output=curator_output,
+            playbook_snapshot=self.playbook.as_prompt(),
+        )
+
+
+class OfflineAdapter(AdapterBase):
+    """Runs multi-epoch offline adaptation on a training split."""
+
+    def run(
+        self,
+        samples: Sequence[Sample],
+        environment: TaskEnvironment,
+        epochs: int = 1,
+    ) -> List[AdapterStepResult]:
+        results: List[AdapterStepResult] = []
+        total_steps = len(samples)
+        for epoch_idx in range(1, epochs + 1):
+            for step_idx, sample in enumerate(samples, start=1):
+                result = self._process_sample(
+                    sample,
+                    environment,
+                    epoch=epoch_idx,
+                    total_epochs=epochs,
+                    step_index=step_idx,
+                    total_steps=total_steps,
+                )
+                results.append(result)
+        return results
+
+
+class OnlineAdapter(AdapterBase):
+    """Processes a stream of samples sequentially, updating the playbook in-place."""
+
+    def run(
+        self,
+        samples: Iterable[Sample],
+        environment: TaskEnvironment,
+    ) -> List[AdapterStepResult]:
+        results: List[AdapterStepResult] = []
+        step_idx = 0
+        for step_idx, sample in enumerate(samples, start=1):
+            result = self._process_sample(
+                sample,
+                environment,
+                epoch=1,
+                total_epochs=1,
+                step_index=step_idx,
+                total_steps=step_idx,
+            )
+            results.append(result)
+        return results

ace/delta.py

@@ -0,0 +1,67 @@
+"""Delta operations produced by the ACE Curator."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, List, Literal, Optional
+
+
+OperationType = Literal["ADD", "UPDATE", "TAG", "REMOVE"]
+
+
+@dataclass
+class DeltaOperation:
+    """Single mutation to apply to the playbook."""
+
+    type: OperationType
+    section: str
+    content: Optional[str] = None
+    bullet_id: Optional[str] = None
+    metadata: Dict[str, int] = field(default_factory=dict)
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, object]) -> "DeltaOperation":
+        return cls(
+            type=str(payload["type"]),
+            section=str(payload.get("section", "")),
+            content=payload.get("content") and str(payload["content"]),
+            bullet_id=payload.get("bullet_id")
+            and str(payload.get("bullet_id")),  # type: ignore[arg-type]
+            metadata={
+                str(k): int(v) for k, v in (payload.get("metadata") or {}).items()
+            },
+        )
+
+    def to_json(self) -> Dict[str, object]:
+        data: Dict[str, object] = {"type": self.type, "section": self.section}
+        if self.content is not None:
+            data["content"] = self.content
+        if self.bullet_id is not None:
+            data["bullet_id"] = self.bullet_id
+        if self.metadata:
+            data["metadata"] = self.metadata
+        return data
+
+
+@dataclass
+class DeltaBatch:
+    """Bundle of curator reasoning and operations."""
+
+    reasoning: str
+    operations: List[DeltaOperation] = field(default_factory=list)
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, object]) -> "DeltaBatch":
+        ops_payload = payload.get("operations")
+        operations = []
+        if isinstance(ops_payload, Iterable):
+            for item in ops_payload:
+                if isinstance(item, dict):
+                    operations.append(DeltaOperation.from_json(item))
+        return cls(reasoning=str(payload.get("reasoning", "")), operations=operations)
+
+    def to_json(self) -> Dict[str, object]:
+        return {
+            "reasoning": self.reasoning,
+            "operations": [op.to_json() for op in self.operations],
+        }

ace/llm.py

@@ -0,0 +1,169 @@
+"""LLM client abstractions used by ACE components."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import json
+from collections import deque
+from dataclasses import dataclass
+from typing import Any, Deque, Dict, Optional, Union
+
+
+@dataclass
+class LLMResponse:
+    """Container for LLM outputs."""
+
+    text: str
+    raw: Optional[Dict[str, Any]] = None
+
+
+class LLMClient(ABC):
+    """Abstract interface so ACE can plug into any chat/completions API."""
+
+    def __init__(self, model: Optional[str] = None) -> None:
+        self.model = model
+
+    @abstractmethod
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        """Return the model text for a given prompt."""
+
+
+class DummyLLMClient(LLMClient):
+    """Deterministic LLM stub for testing and dry runs."""
+
+    def __init__(self, responses: Optional[Deque[str]] = None) -> None:
+        super().__init__(model="dummy")
+        self._responses: Deque[str] = responses or deque()
+
+    def queue(self, text: str) -> None:
+        """Enqueue a response to be used on the next completion call."""
+        self._responses.append(text)
+
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        if not self._responses:
+            raise RuntimeError("DummyLLMClient ran out of queued responses.")
+        return LLMResponse(text=self._responses.popleft())
+
+
+class TransformersLLMClient(LLMClient):
+    """LLM client powered by `transformers` pipelines for chat-style models."""
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        max_new_tokens: int = 512,
+        temperature: float = 0.0,
+        top_p: float = 0.9,
+        device_map: Union[str, Dict[str, int]] = "auto",
+        torch_dtype: Union[str, "torch.dtype"] = "auto",
+        trust_remote_code: bool = True,
+        system_prompt: Optional[str] = None,
+        generation_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        super().__init__(model=model_path)
+
+        # Import transformers lazily to avoid mandatory dependency for all users.
+        from transformers import AutoTokenizer, pipeline  # type: ignore[import-untyped]
+
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_path, trust_remote_code=trust_remote_code
+        )
+        self._pipeline = pipeline(
+            "text-generation",
+            model=model_path,
+            tokenizer=self._tokenizer,
+            torch_dtype=torch_dtype,
+            device_map=device_map,
+            trust_remote_code=trust_remote_code,
+        )
+        self._system_prompt = system_prompt or (
+            "You are a JSON-only assistant that MUST reply with a single valid JSON object without extra text.\n"
+            "Reasoning: low\n"
+            "Do not expose analysis or chain-of-thought. Respond using the final JSON only."
+        )
+        self._defaults: Dict[str, Any] = {
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "do_sample": temperature > 0.0,
+            "return_full_text": False,
+        }
+        if generation_kwargs:
+            self._defaults.update(generation_kwargs)
+
+    def complete(self, prompt: str, **kwargs: Any) -> LLMResponse:
+        call_kwargs = dict(self._defaults)
+        kwargs = dict(kwargs)
+        kwargs.pop("refinement_round", None)
+        call_kwargs.update(kwargs)
+
+        # Build chat-formatted messages to leverage harmony template.
+        messages = [
+            {"role": "system", "content": self._system_prompt},
+            {"role": "user", "content": prompt},
+        ]
+
+        outputs = self._pipeline(messages, **call_kwargs)
+        text = self._postprocess_text(self._extract_text(outputs))
+        return LLMResponse(text=text, raw={"outputs": outputs})
+
+    def _extract_text(self, outputs: Any) -> str:
+        """Normalize pipeline outputs into a single string response."""
+        if not outputs:
+            return ""
+        candidate = outputs[0]
+
+        # Newer transformers versions return {"generated_text": [{"role": ..., "content": ...}, ...]}
+        if isinstance(candidate, dict) and "generated_text" in candidate:
+            generated = candidate["generated_text"]
+            if isinstance(generated, list):
+                # Grab the assistant role content if present.
+                for message in generated:
+                    if isinstance(message, dict) and message.get("role") == "assistant":
+                        content = message.get("content")
+                        if isinstance(content, str):
+                            return content.strip()
+                # Fallback to last item's content/text.
+                last = generated[-1]
+                if isinstance(last, dict):
+                    return str(last.get("content") or last.get("text") or "")
+                return str(last)
+            if isinstance(generated, dict):
+                return str(generated.get("content") or generated.get("text") or "")
+            return str(generated)
+
+        # Older versions might return {"generated_text": "..."}
+        if isinstance(candidate, dict) and isinstance(candidate.get("generated_text"), str):
+            return candidate["generated_text"].strip()
+
+        # Ultimate fallback: string representation.
+        return str(candidate).strip()
+
+    def _postprocess_text(self, text: str) -> str:
+        """Trim analyzer prefixes and isolate JSON payloads when present."""
+        trimmed = text.strip()
+        if not trimmed:
+            return trimmed
+
+        marker = "assistantfinal"
+        if marker in trimmed:
+            trimmed = trimmed.split(marker, 1)[1].strip()
+
+        if trimmed.startswith(marker):
+            trimmed = trimmed[len(marker) :].strip()
+
+        # Attempt to extract the first JSON object substring.
+        if trimmed and trimmed[0] != "{":
+            start = trimmed.find("{")
+            end = trimmed.rfind("}")
+            if start != -1 and end != -1 and end > start:
+                candidate = trimmed[start : end + 1].strip()
+                candidate_clean = candidate.replace("\r", " ").replace("\n", " ")
+                try:
+                    json.loads(candidate_clean)
+                    return candidate_clean
+                except json.JSONDecodeError:
+                    pass
+
+        return trimmed.replace("\r", " ").replace("\n", " ")

ace/llm_providers/__init__.py

@@ -0,0 +1,13 @@
+"""Production LLM client implementations for ACE."""
+
+from .litellm_client import LiteLLMClient
+
+try:
+    from .langchain_client import LangChainLiteLLMClient
+except ImportError:
+    LangChainLiteLLMClient = None  # Optional dependency
+
+__all__ = [
+    "LiteLLMClient",
+    "LangChainLiteLLMClient",
+]