mindagent 0.1.0__py3-none-any.whl

mindagent/__init__.py

@@ -0,0 +1 @@
+"""mindagent public package."""

mindagent/context/__init__.py

@@ -0,0 +1,25 @@
+from .manager import ContextConfig, ContextManager
+from .models import (
+    ContextBundle,
+    ContextPackResult,
+    ContextPressure,
+    ContextPressureLevel,
+    ContextRecord,
+    ContextScope,
+)
+from .packer import CacheAwarePacker, ContextOverflowError
+from .store import ContextStore
+
+__all__ = [
+    "CacheAwarePacker",
+    "ContextBundle",
+    "ContextConfig",
+    "ContextManager",
+    "ContextOverflowError",
+    "ContextPackResult",
+    "ContextPressure",
+    "ContextPressureLevel",
+    "ContextRecord",
+    "ContextScope",
+    "ContextStore",
+]

mindagent/context/manager.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from mindagent.core.context import (
+    AgentContext,
+    AgentDecision,
+    DecisionType,
+    Observation,
+)
+
+from .models import ContextPackResult, ContextRecord, ContextScope
+from .packer import CacheAwarePacker
+from .store import ContextStore
+
+
+@dataclass
+class ContextConfig:
+    system_prompt: str | None = None
+    max_tokens: int = 120000
+    response_reserve_tokens: int = 4096
+    safety_margin_tokens: int = 256
+    warning_ratio: float = 0.70
+    compact_trigger_ratio: float = 0.85
+    compact_target_ratio: float = 0.60
+    model: str = "gpt-4o"
+    image_token_cost: int = 1024
+    use_tiktoken: bool = False
+
+    def __post_init__(self) -> None:
+        if self.max_tokens < 1:
+            raise ValueError("max_tokens 必须大于 0")
+        if self.response_reserve_tokens < 0:
+            raise ValueError("response_reserve_tokens 不能小于 0")
+        if self.safety_margin_tokens < 0:
+            raise ValueError("safety_margin_tokens 不能小于 0")
+        if self.image_token_cost < 0:
+            raise ValueError("image_token_cost 不能小于 0")
+        if not (
+            0 < self.warning_ratio
+            < self.compact_trigger_ratio
+            <= 1
+        ):
+            raise ValueError("warning/compact trigger 比例配置无效")
+        if not 0 < self.compact_target_ratio < self.compact_trigger_ratio:
+            raise ValueError("compact_target_ratio 必须小于触发比例")
+
+
+@dataclass
+class _RunContextState:
+    store: ContextStore
+    epoch: int = 0
+    packed_once: bool = False
+    dropped_signature: tuple[str, ...] = ()
+
+
+class ContextManager:
+    def __init__(self, config: ContextConfig | None = None):
+        self.config = config or ContextConfig()
+        self._encoder = (
+            self._load_encoder(self.config.model)
+            if self.config.use_tiktoken
+            else None
+        )
+        self._runs: dict[str, _RunContextState] = {}
+
+    async def build_context(self, context: AgentContext) -> None:
+        if context.metadata.get("_context_built"):
+            return
+
+        state = self._state(context)
+        self._ingest_existing_messages(state.store, context.messages)
+        if self.config.system_prompt and not any(
+            record.kind == "system"
+            for record in state.store.records
+        ):
+            state.store.append(
+                ContextRecord(
+                    kind="system",
+                    message={
+                        "role": "system",
+                        "content": self.config.system_prompt,
+                    },
+                    scope=ContextScope.GLOBAL,
+                    required=True,
+                    priority=100,
+                    salience=1.0,
+                    stable_prefix=True,
+                )
+            )
+
+        state.store.append(
+            ContextRecord(
+                kind="user",
+                message={
+                    "role": "user",
+                    "content": self._build_user_content(context),
+                },
+                scope=ContextScope.TURN,
+                required=True,
+                priority=100,
+                salience=1.0,
+                group_id=f"turn:{context.run_id}",
+            )
+        )
+        context.metadata["_context_built"] = True
+        await self.prepare_context(context)
+
+    async def prepare_context(
+        self,
+        context: AgentContext,
+        *,
+        tools: list[dict[str, Any]] | None = None,
+    ) -> ContextPackResult:
+        state = self._state(context)
+        result = self._packer().pack(
+            state.store.bundles(),
+            tools=tools or (),
+            epoch=state.epoch,
+            allow_compaction=not state.packed_once,
+        )
+        dropped_signature = tuple(sorted(result.dropped_record_ids))
+        if (
+            dropped_signature
+            and dropped_signature != state.dropped_signature
+            and not state.packed_once
+        ):
+            state.epoch += 1
+            result.epoch = state.epoch
+        state.dropped_signature = dropped_signature
+        state.packed_once = True
+        context.messages[:] = result.messages
+        context.metadata["_context_pack"] = {
+            "epoch": result.epoch,
+            "prefix_hash": result.prefix_hash,
+            "input_tokens": result.input_tokens,
+            "available_tokens": result.available_tokens,
+            "pressure": result.pressure.level.value,
+            "pressure_ratio": result.pressure.ratio,
+            "dropped_record_ids": result.dropped_record_ids,
+        }
+        return result
+
+    async def record_decision(
+        self,
+        context: AgentContext,
+        decision: AgentDecision,
+    ) -> None:
+        state = self._state(context)
+        if decision.decision_type == DecisionType.FINAL:
+            message = {
+                "role": "assistant",
+                "content": decision.final_answer or "",
+            }
+            group_id = f"turn:{context.run_id}"
+        else:
+            action = decision.action
+            if action is None:
+                return
+            call_id = decision.metadata.get("provider_call_id")
+            group_id = f"tool:{call_id or context.step_index}"
+            message = {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "id": call_id or f"call-{context.step_index}",
+                        "type": "function",
+                        "function": {
+                            "name": action.name,
+                            "arguments": action.arguments,
+                        },
+                    }
+                ],
+            }
+
+        state.store.append(
+            ContextRecord(
+                kind="assistant",
+                message=message,
+                scope=ContextScope.ITERATION,
+                required=True,
+                priority=95,
+                salience=1.0,
+                group_id=group_id,
+            )
+        )
+        context.messages.append(message)
+
+    async def update_after_observation(
+        self,
+        context: AgentContext,
+        observation: Observation,
+    ) -> None:
+        call_id = None
+        if context.current_decision:
+            call_id = context.current_decision.metadata.get(
+                "provider_call_id"
+            )
+        message: dict[str, Any] = {
+            "role": "tool",
+            "name": observation.action.name,
+            "content": json.dumps(
+                {
+                    "ok": observation.ok,
+                    "result": observation.result,
+                    "error": observation.error,
+                },
+                ensure_ascii=False,
+                default=str,
+            ),
+        }
+        if call_id:
+            message["tool_call_id"] = call_id
+
+        self._state(context).store.append(
+            ContextRecord(
+                kind="tool",
+                message=message,
+                scope=ContextScope.ITERATION,
+                required=True,
+                priority=95,
+                salience=1.0,
+                group_id=f"tool:{call_id or context.step_index}",
+            )
+        )
+        context.messages.append(message)
+
+    def estimate_tokens(self, messages: list[dict[str, Any]]) -> int:
+        total = 0
+        for message in messages:
+            total += 4
+            content = message.get("content")
+            if isinstance(content, list):
+                for item in content:
+                    if item.get("type") in {"image_url", "input_image"}:
+                        total += self.config.image_token_cost
+                    else:
+                        total += self._count_text(
+                            json.dumps(item, ensure_ascii=False)
+                        )
+            else:
+                total += self._count_text(
+                    json.dumps(message, ensure_ascii=False, default=str)
+                )
+        return total
+
+    def trim(self, context: AgentContext) -> None:
+        store = ContextStore()
+        self._ingest_existing_messages(store, context.messages)
+        bundles = store.bundles()
+        stable = [bundle for bundle in bundles if bundle.stable_prefix]
+        optional = [bundle for bundle in bundles if not bundle.stable_prefix]
+        selected = stable + (optional[-1:] if optional else [])
+        context.messages[:] = [
+            record.message
+            for bundle in selected
+            for record in bundle.records
+        ]
+
+    def get_store(self, run_id: str) -> ContextStore | None:
+        state = self._runs.get(run_id)
+        return state.store if state else None
+
+    def release_context(self, run_id: str) -> None:
+        self._runs.pop(run_id, None)
+
+    def _state(self, context: AgentContext) -> _RunContextState:
+        return self._runs.setdefault(
+            context.run_id,
+            _RunContextState(store=ContextStore()),
+        )
+
+    def _packer(self) -> CacheAwarePacker:
+        reserve = min(
+            self.config.response_reserve_tokens,
+            max(0, self.config.max_tokens // 4),
+        )
+        margin = min(
+            self.config.safety_margin_tokens,
+            max(0, self.config.max_tokens // 10),
+        )
+        return CacheAwarePacker(
+            self.estimate_tokens,
+            max_tokens=self.config.max_tokens,
+            response_reserve=reserve,
+            safety_margin=margin,
+            warning_ratio=self.config.warning_ratio,
+            compact_trigger_ratio=self.config.compact_trigger_ratio,
+            compact_target_ratio=self.config.compact_target_ratio,
+        )
+
+    def _ingest_existing_messages(
+        self,
+        store: ContextStore,
+        messages: list[dict[str, Any]],
+    ) -> None:
+        turn_index = 0
+        current_turn = "history:0"
+        for message in messages:
+            role = message.get("role", "unknown")
+            if role == "system":
+                group_id = None
+            elif role == "user":
+                turn_index += 1
+                current_turn = f"history:{turn_index}"
+                group_id = current_turn
+            elif role == "assistant" and message.get("tool_calls"):
+                calls = message["tool_calls"]
+                call_id = calls[0].get("id") if calls else None
+                group_id = f"tool:{call_id or current_turn}"
+            elif role == "tool":
+                call_id = message.get("tool_call_id")
+                group_id = f"tool:{call_id or current_turn}"
+            else:
+                group_id = current_turn
+
+            store.append(
+                ContextRecord(
+                    kind=role,
+                    message=message,
+                    scope=(
+                        ContextScope.GLOBAL
+                        if role == "system"
+                        else ContextScope.SESSION
+                    ),
+                    required=role == "system",
+                    priority=100 if role == "system" else 50,
+                    salience=1.0 if role == "system" else 0.5,
+                    group_id=group_id,
+                    stable_prefix=role == "system",
+                )
+            )
+
+    def _build_user_content(self, context: AgentContext) -> Any:
+        images = context.artifacts.get("images") or []
+        if not images:
+            return context.user_input
+
+        content: list[dict[str, Any]] = [
+            {"type": "text", "text": context.user_input}
+        ]
+        for image in images:
+            if isinstance(image, str):
+                image_url = {"url": image}
+            elif isinstance(image, dict) and "url" in image:
+                image_url = {
+                    key: value
+                    for key, value in image.items()
+                    if key in {"url", "detail"}
+                }
+            else:
+                raise ValueError(
+                    "images 仅支持 URL/data URL 字符串或包含 url 的字典"
+                )
+            content.append({"type": "image_url", "image_url": image_url})
+        return content
+
+    def _count_text(self, text: str) -> int:
+        if self._encoder is not None:
+            return len(self._encoder.encode(text))
+        return max(1, len(text) // 4)
+
+    @staticmethod
+    def _load_encoder(model: str):
+        try:
+            import tiktoken
+
+            try:
+                return tiktoken.encoding_for_model(model)
+            except KeyError:
+                return tiktoken.get_encoding("cl100k_base")
+        except Exception:
+            return None

mindagent/context/models.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class ContextScope(str, Enum):
+    GLOBAL = "global"
+    SESSION = "session"
+    TURN = "turn"
+    ITERATION = "iteration"
+
+
+class ContextPressureLevel(str, Enum):
+    NORMAL = "normal"
+    WARNING = "warning"
+    CRITICAL = "critical"
+
+
+@dataclass(frozen=True)
+class ContextRecord:
+    kind: str
+    message: dict[str, Any]
+    scope: ContextScope
+    required: bool = False
+    priority: int = 50
+    salience: float = 0.5
+    group_id: str | None = None
+    stable_prefix: bool = False
+    record_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
+@dataclass(frozen=True)
+class ContextBundle:
+    bundle_id: str
+    records: tuple[ContextRecord, ...]
+    required: bool
+    priority: int
+    salience: float
+    stable_prefix: bool
+
+
+@dataclass(frozen=True)
+class ContextPressure:
+    level: ContextPressureLevel
+    total_tokens: int
+    available_tokens: int
+    ratio: float
+
+
+@dataclass
+class ContextPackResult:
+    messages: list[dict[str, Any]]
+    input_tokens: int
+    available_tokens: int
+    pressure: ContextPressure
+    prefix_hash: str
+    epoch: int
+    selected_record_ids: list[str] = field(default_factory=list)
+    dropped_record_ids: list[str] = field(default_factory=list)
+    decisions: list[dict[str, Any]] = field(default_factory=list)

mindagent/context/packer.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from .models import (
+    ContextBundle,
+    ContextPackResult,
+    ContextPressure,
+    ContextPressureLevel,
+)
+
+
+TokenCounter = Callable[[list[dict[str, Any]]], int]
+
+
+class ContextOverflowError(RuntimeError):
+    pass
+
+
+class CacheAwarePacker:
+    def __init__(
+        self,
+        token_counter: TokenCounter,
+        *,
+        max_tokens: int,
+        response_reserve: int,
+        safety_margin: int,
+        warning_ratio: float,
+        compact_trigger_ratio: float,
+        compact_target_ratio: float,
+    ):
+        self.token_counter = token_counter
+        self.max_tokens = max_tokens
+        self.response_reserve = response_reserve
+        self.safety_margin = safety_margin
+        self.warning_ratio = warning_ratio
+        self.compact_trigger_ratio = compact_trigger_ratio
+        self.compact_target_ratio = compact_target_ratio
+
+    def pack(
+        self,
+        bundles: Sequence[ContextBundle],
+        *,
+        tools: Sequence[dict[str, Any]] = (),
+        epoch: int = 0,
+        allow_compaction: bool = True,
+    ) -> ContextPackResult:
+        available = (
+            self.max_tokens
+            - self.response_reserve
+            - self.safety_margin
+        )
+        if available < 1:
+            raise ValueError("Context 可用输入预算必须大于 0")
+
+        total = sum(self._bundle_tokens(bundle) for bundle in bundles)
+        ratio = total / available
+        level = self._pressure_level(ratio)
+        pressure = ContextPressure(level, total, available, ratio)
+        prefix_hash = self._prefix_hash(bundles, tools)
+
+        if ratio < self.compact_trigger_ratio or not allow_compaction:
+            selected = list(bundles)
+            if total > available:
+                selected = self._select(bundles, available)
+        else:
+            target = max(
+                self._required_tokens(bundles),
+                int(available * self.compact_target_ratio),
+            )
+            selected = self._select(bundles, min(target, available))
+
+        messages = [
+            record.message
+            for bundle in selected
+            for record in bundle.records
+        ]
+        input_tokens = self.token_counter(messages)
+        if input_tokens > available:
+            raise ContextOverflowError(
+                "required context 超过模型可用输入预算"
+            )
+
+        selected_ids = {
+            record.record_id
+            for bundle in selected
+            for record in bundle.records
+        }
+        selected_record_ids = [
+            record.record_id
+            for bundle in bundles
+            for record in bundle.records
+            if record.record_id in selected_ids
+        ]
+        dropped_ids = [
+            record.record_id
+            for bundle in bundles
+            for record in bundle.records
+            if record.record_id not in selected_ids
+        ]
+        decisions = [
+            {
+                "bundle_id": bundle.bundle_id,
+                "action": (
+                    "kept"
+                    if all(
+                        record.record_id in selected_ids
+                        for record in bundle.records
+                    )
+                    else "dropped"
+                ),
+                "reason": (
+                    "required_or_selected"
+                    if any(
+                        record.record_id in selected_ids
+                        for record in bundle.records
+                    )
+                    else "context_pressure"
+                ),
+            }
+            for bundle in bundles
+        ]
+        return ContextPackResult(
+            messages=messages,
+            input_tokens=input_tokens,
+            available_tokens=available,
+            pressure=pressure,
+            prefix_hash=prefix_hash,
+            epoch=epoch,
+            selected_record_ids=selected_record_ids,
+            dropped_record_ids=dropped_ids,
+            decisions=decisions,
+        )
+
+    def _select(
+        self,
+        bundles: Sequence[ContextBundle],
+        budget: int,
+    ) -> list[ContextBundle]:
+        selected_ids: set[str] = set()
+        used = 0
+
+        for bundle in bundles:
+            if not (bundle.stable_prefix or bundle.required):
+                continue
+            cost = self._bundle_tokens(bundle)
+            selected_ids.add(bundle.bundle_id)
+            used += cost
+
+        optional = [
+            (index, bundle)
+            for index, bundle in enumerate(bundles)
+            if bundle.bundle_id not in selected_ids
+        ]
+        optional.sort(
+            key=lambda item: (
+                -item[1].priority,
+                -item[1].salience,
+                -item[0],
+            )
+        )
+        for _, bundle in optional:
+            cost = self._bundle_tokens(bundle)
+            if used + cost <= budget:
+                selected_ids.add(bundle.bundle_id)
+                used += cost
+
+        return [
+            bundle
+            for bundle in bundles
+            if bundle.bundle_id in selected_ids
+        ]
+
+    def _bundle_tokens(self, bundle: ContextBundle) -> int:
+        return self.token_counter(
+            [record.message for record in bundle.records]
+        )
+
+    def _required_tokens(
+        self,
+        bundles: Sequence[ContextBundle],
+    ) -> int:
+        return sum(
+            self._bundle_tokens(bundle)
+            for bundle in bundles
+            if bundle.stable_prefix or bundle.required
+        )
+
+    def _pressure_level(self, ratio: float) -> ContextPressureLevel:
+        if ratio >= self.compact_trigger_ratio:
+            return ContextPressureLevel.CRITICAL
+        if ratio >= self.warning_ratio:
+            return ContextPressureLevel.WARNING
+        return ContextPressureLevel.NORMAL
+
+    @staticmethod
+    def _prefix_hash(
+        bundles: Sequence[ContextBundle],
+        tools: Sequence[dict[str, Any]],
+    ) -> str:
+        prefix = [
+            record.message
+            for bundle in bundles
+            if bundle.stable_prefix
+            for record in bundle.records
+        ]
+        payload = json.dumps(
+            {"messages": prefix, "tools": list(tools)},
+            ensure_ascii=False,
+            sort_keys=True,
+            separators=(",", ":"),
+            default=str,
+        )
+        return hashlib.sha256(payload.encode("utf-8")).hexdigest()