@scemoon/cdh 1.0.0

package/onecode/agent/context.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional, Union
+
+from onecode.models.provider import (
+    ContentBlock,
+    ContentBlockType,
+    ImageContent,
+    Message as ProviderMessage,
+    ToolResult as ProviderToolResult,
+    ToolUse,
+)
+
+logger = logging.getLogger("onecode.context")
+
+
+# ── Token estimation (model-aware, with tiktoken fallback) ──
+
+_ENCODING_CACHE: dict[str, Any] = {}
+
+def _get_encoding(model: str):
+    if model in _ENCODING_CACHE:
+        return _ENCODING_CACHE[model]
+    try:
+        import tiktoken
+        try:
+            enc = tiktoken.encoding_for_model(model)
+        except KeyError:
+            enc = tiktoken.get_encoding("cl100k_base")
+        _ENCODING_CACHE[model] = enc
+        return enc
+    except Exception:
+        _ENCODING_CACHE[model] = None
+        return None
+
+
+def _estimate_tokens(text: str, model: str = "gpt-4") -> int:
+    if not text:
+        return 0
+    enc = _get_encoding(model)
+    if enc is not None:
+        try:
+            return len(enc.encode(text, disallowed_special=()))
+        except Exception:
+            pass
+    return max(1, len(text) // 4)
+
+
+# ── Context config ──
+
+class ContextConfig:
+    max_tokens: int = 100000
+    max_messages: int = 1000
+    compact_threshold: float = 0.85
+    model: str = "gpt-4"
+
+
+class Message:
+    def __init__(self, role: str, content: Union[str, list], name: Optional[str] = None):
+        self.role = role
+        self.content = content
+        self.name = name
+
+    def to_dict(self) -> dict:
+        d: dict[str, Any] = {"role": self.role, "content": self.content}
+        if self.name:
+            d["name"] = self.name
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict) -> Message:
+        return cls(role=data["role"], content=data["content"], name=data.get("name"))
+
+
+def _block_text(block: Any) -> str:
+    if isinstance(block, dict):
+        btype = block.get("type", "")
+        if btype == "resource":
+            resource = block.get("resource", {})
+            return resource.get("text", "") or f"[{btype}: {resource.get('mimeType', '')} {len(resource.get('blob', ''))} bytes]"
+        if btype == "image":
+            return f"[image: {block.get('mimeType', block.get('source', {}).get('media_type', 'unknown'))}]"
+        return str(block.get("text", "") or block.get("content", ""))
+    return str(block)
+
+
+class ContextManager:
+    def __init__(self, config: Optional[ContextConfig] = None):
+        self.config = config or ContextConfig()
+        self.messages: list[Message] = []
+        self._token_count = 0
+
+    # -- model-aware token estimation (incremental O(1) per message) --
+
+    def _estimate_message_tokens(self, msg: Message) -> int:
+        model = self.config.model
+        if isinstance(msg.content, str):
+            return _estimate_tokens(msg.content, model)
+        elif isinstance(msg.content, list):
+            return sum(_estimate_tokens(_block_text(b), model) for b in msg.content)
+        return 0
+
+    # -- message management (incremental token counters) --
+
+    def add_message(self, role: str, content: Union[str, list], name: Optional[str] = None) -> None:
+        msg = Message(role=role, content=content, name=name)
+        self.messages.append(msg)
+        self._token_count += self._estimate_message_tokens(msg)
+
+    def add_system(self, content: str) -> None:
+        self.add_message("system", content)
+
+    def replace_system_section(self, marker: str, new_content: str) -> bool:
+        for m in self.messages:
+            if m.role == "system" and isinstance(m.content, str) and marker in m.content:
+                old_tokens = self._estimate_message_tokens(m)
+                m.content = new_content
+                self._token_count += self._estimate_message_tokens(m) - old_tokens
+                return True
+        return False
+
+    def remove_system_by_marker(self, marker: str) -> int:
+        before = len(self.messages)
+        removed_tokens = 0
+        kept: list[Message] = []
+        for m in self.messages:
+            if m.role == "system" and isinstance(m.content, str) and marker in m.content:
+                removed_tokens += self._estimate_message_tokens(m)
+            else:
+                kept.append(m)
+        self.messages = kept
+        self._token_count -= removed_tokens
+        return before - len(self.messages)
+
+    def add_user(self, content: str | list) -> None:
+        self.add_message("user", content)
+
+    def add_assistant(self, content: str | list) -> None:
+        self.add_message("assistant", content)
+
+    def add_tool_use(self, tool_use_id: str, name: str, input_data: dict[str, Any]) -> None:
+        self.add_message("assistant", [{"type": "tool_use", "id": tool_use_id, "name": name, "input": input_data}])
+
+    def add_tool_result(self, tool_call_id: str, content: Union[str, dict, list], is_error: bool = False) -> None:
+        if isinstance(content, (dict, list)):
+            self.add_message("tool", content, name=tool_call_id)
+        else:
+            self.add_message("tool", [{"type": "tool_result", "tool_use_id": tool_call_id, "content": content, "is_error": is_error}], name=tool_call_id)
+
+    # -- full recalibration (rare: after bulk ops) --
+
+    def _update_token_count(self) -> None:
+        total = 0
+        for m in self.messages:
+            total += self._estimate_message_tokens(m)
+        self._token_count = total
+
+    def should_compact(self) -> bool:
+        if self.config.max_tokens <= 0:
+            return False
+        return self._token_count >= self.config.max_tokens * self.config.compact_threshold
+
+    # ── Tiered compression pipeline ──
+
+    def compact(self) -> str:
+        """Apply tiered compression.
+
+        Returns the level applied: ``"none"``, ``"light"``, ``"medium"``, or ``"heavy"``.
+        """
+        if len(self.messages) <= 2:
+            return "none"
+
+        system_msgs = [m for m in self.messages if m.role == "system"]
+        other_msgs = [m for m in self.messages if m.role != "system"]
+        if not other_msgs:
+            return "none"
+
+        threshold = self.config.max_tokens * self.config.compact_threshold
+
+        # Tier 1 — Light: compress verbose tool results in-place
+        self._tier_compress_tool_results(other_msgs)
+        if self._token_count_under(system_msgs, other_msgs, threshold):
+            self.messages = system_msgs + other_msgs
+            self._update_token_count()
+            logger.info("compact: light — compressed tool results")
+            return "light"
+
+        # Tier 2 — Medium: truncate older non-system messages
+        self._tier_truncate_old(other_msgs)
+        if self._token_count_under(system_msgs, other_msgs, threshold):
+            self.messages = system_msgs + other_msgs
+            self._update_token_count()
+            logger.info("compact: medium — truncated old messages")
+            return "medium"
+
+        # Tier 3 — Heavy: full summarization
+        summary = self._summarize_messages(other_msgs)
+        self.messages = system_msgs + [
+            Message(role="system", content=f"[Previous context summarized]\n{summary}")
+        ]
+        self._update_token_count()
+        logger.info("compact: heavy — full summarization")
+        return "heavy"
+
+    def _token_count_under(self, system_msgs, other_msgs, threshold) -> bool:
+        total = 0
+        for m in system_msgs:
+            total += self._estimate_message_tokens(m)
+        for m in other_msgs:
+            total += self._estimate_message_tokens(m)
+        return total < threshold
+
+    def _tier_compress_tool_results(self, msgs: list[Message]) -> None:
+        for m in msgs:
+            if m.role == "tool" and isinstance(m.content, list):
+                for b in m.content:
+                    if isinstance(b, dict) and b.get("type") == "tool_result":
+                        content = b.get("content", "")
+                        if isinstance(content, str) and len(content) > 500:
+                            b["content"] = content[:500] + "\n... [truncated]"
+
+    def _tier_truncate_old(self, msgs: list[Message], keep_recent: int = 10) -> None:
+        if len(msgs) <= keep_recent:
+            return
+        for m in msgs[:-keep_recent]:
+            if isinstance(m.content, str) and len(m.content) > 200:
+                m.content = m.content[:200] + "..."
+            elif isinstance(m.content, list):
+                m.content = [b for b in m.content if isinstance(b, dict) and b.get("type") == "tool_use"]
+                if not m.content:
+                    m.content = "[truncated]"
+
+    def _summarize_messages(self, msgs: list[Message]) -> str:
+        content_parts = []
+        for m in msgs[-30:]:
+            prefix = {"user": "User", "assistant": "Assistant", "tool": "Tool"}.get(m.role, m.role)
+            if isinstance(m.content, str):
+                content_parts.append(f"{prefix}: {m.content[:800]}")
+            elif isinstance(m.content, list):
+                texts = []
+                for b in m.content:
+                    if isinstance(b, dict):
+                        btype = b.get("type", "")
+                        if btype == "tool_use":
+                            texts.append(f"[call {b.get('name', '?')}]")
+                        elif btype == "tool_result":
+                            texts.append(f"[result: {str(b.get('content', ''))[:300]}]")
+                        else:
+                            texts.append(_block_text(b)[:300])
+                    else:
+                        texts.append(str(b)[:300])
+                content_parts.append(f"{prefix}: {' | '.join(texts)[:800]}")
+        return "\n".join(content_parts[-20:])
+
+    def get_context(self) -> list:
+        def _is_image_mime(mime: str) -> bool:
+            return mime and mime.startswith("image/")
+
+        def to_provider_content(msg: Message) -> Union[str, list]:
+            if isinstance(msg.content, str):
+                return msg.content
+            blocks = []
+            for block in msg.content:
+                if isinstance(block, dict):
+                    btype = block.get("type", "text")
+                    if btype == "text":
+                        blocks.append(ContentBlock(type=ContentBlockType.TEXT, text=block.get("text", "")))
+                    elif btype == "thinking":
+                        blocks.append(ContentBlock(type=ContentBlockType.THINKING, thinking=block.get("thinking", "")))
+                    elif btype == "tool_use":
+                        blocks.append(ContentBlock(
+                            type=ContentBlockType.TOOL_USE,
+                            tool_use=ToolUse(
+                                id=block.get("id", ""),
+                                name=block.get("name", ""),
+                                input=block.get("input", {}),
+                                caller=block.get("caller", "agent"),
+                            ),
+                        ))
+                    elif btype == "tool_result":
+                        blocks.append(ContentBlock(
+                            type=ContentBlockType.TOOL_RESULT,
+                            tool_result=ProviderToolResult(
+                                tool_use_id=block.get("tool_use_id", ""),
+                                content=block.get("content", ""),
+                                is_error=block.get("is_error", False),
+                            ),
+                        ))
+                    elif btype == "resource":
+                        resource = block.get("resource", {})
+                        uri = resource.get("uri", "")
+                        mime = resource.get("mimeType", "")
+                        if resource.get("text") is not None:
+                            blocks.append(ContentBlock(
+                                type=ContentBlockType.TEXT,
+                                text=f"\n\n[File: {uri}]({mime})\n```\n{resource['text']}\n```\n",
+                            ))
+                        elif resource.get("blob") and _is_image_mime(mime):
+                            blocks.append(ContentBlock(
+                                type=ContentBlockType.IMAGE,
+                                image=ImageContent(
+                                    data=resource["blob"],
+                                    mime_type=mime,
+                                ),
+                            ))
+                        elif resource.get("blob"):
+                            blocks.append(ContentBlock(
+                                type=ContentBlockType.TEXT,
+                                text=f"\n\n[Attachment: {uri}]({mime})\n[Binary data, {len(resource['blob'])} base64 bytes]\n",
+                            ))
+                    elif btype == "image":
+                        data = block.get("data", block.get("source", {}).get("data", ""))
+                        mime = block.get("mimeType", block.get("source", {}).get("media_type", "image/png"))
+                        blocks.append(ContentBlock(
+                            type=ContentBlockType.IMAGE,
+                            image=ImageContent(data=data, mime_type=mime),
+                        ))
+            return blocks
+
+        return [
+            ProviderMessage(role=m.role, content=to_provider_content(m), name=m.name)
+            for m in self.messages
+        ]
+
+    def reset(self) -> None:
+        self.messages = []
+        self._token_count = 0
+
+    def load_from_session(self, data: list[dict]) -> None:
+        self.messages = [Message.from_dict(d) for d in data]
+        self._update_token_count()
+
+    def to_session_format(self) -> list[dict]:
+        return [m.to_dict() for m in self.messages]
+
+    def info(self) -> str:
+        return f"Messages: {len(self.messages)}, Tokens: ~{self._token_count}"