devtorch-core 3.0.1__py3-none-any.whl

devtorch_core/wrapper/bedrock.py

@@ -0,0 +1,342 @@
+"""
+devtorch_core.wrapper.bedrock
+==============================
+Stub wrapper around the AWS Bedrock ``boto3`` client.
+
+Intercepts ``invoke_model`` for Anthropic Claude models served via Bedrock
+(body format: ``{"anthropic_version": ..., "messages": [...], "system": ...}``).
+
+Other model families (Titan, Llama, Mistral, etc.) pass through unchanged
+with a debug log.
+
+Usage
+-----
+::
+
+    from devtorch_core.wrapper.bedrock import DevTorchBedrock
+
+    client = DevTorchBedrock(region_name="us-east-1")
+    response = client.invoke_model(
+        modelId="anthropic.claude-opus-4-5-20240229-v1:0",
+        body=json.dumps({
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 1024,
+            "messages": [{"role": "user", "content": "Hello"}],
+        }),
+        contentType="application/json",
+        accept="application/json",
+    )
+"""
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from typing import Any
+
+from .base import CaptureOrchestrator, RACPInjector, _reasoning_text, find_gcc_repo, is_disabled
+
+logger = logging.getLogger("devtorch.wrapper.bedrock")
+
+# ---------------------------------------------------------------------------
+# Optional import
+# ---------------------------------------------------------------------------
+
+try:
+    import boto3 as _boto3  # type: ignore[import]
+
+    HAS_BOTO3 = True
+except ImportError:
+    HAS_BOTO3 = False
+    _boto3 = None  # type: ignore[assignment]
+
+# Model ID prefixes that use the Anthropic Claude Bedrock body format
+_ANTHROPIC_BEDROCK_PREFIXES = (
+    "anthropic.claude",
+    "us.anthropic.claude",
+    "eu.anthropic.claude",
+)
+
+
+# ---------------------------------------------------------------------------
+# Wrapper
+# ---------------------------------------------------------------------------
+
+
+class DevTorchBedrock:
+    """
+    Wrapper around a ``boto3`` Bedrock runtime client.
+
+    Intercepts ``invoke_model`` for Anthropic Claude models to inject RACP
+    context and capture reasoning.  All other calls are passed through
+    unchanged.  Degrades gracefully if boto3 is not installed or .GCC/ is
+    missing.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        gcc_repo: Any = None,
+        service_name: str = "bedrock-runtime",
+        **kwargs: Any,
+    ) -> None:
+        if not HAS_BOTO3:
+            raise ImportError(
+                "boto3 package not installed. "
+                "pip install devtorch[wrapper]"
+            )
+        self._client = _boto3.client(service_name, *args, **kwargs)  # type: ignore[union-attr]
+        self._gcc = gcc_repo or find_gcc_repo()
+        self._injector = RACPInjector(self._gcc) if self._gcc else None
+        self._capture = CaptureOrchestrator(self._gcc) if self._gcc else None
+
+    # ------------------------------------------------------------------
+    # Intercepted method
+    # ------------------------------------------------------------------
+
+    def invoke_model(self, *, modelId: str, body: Any, **kwargs: Any) -> Any:
+        """
+        Intercept invoke_model for Claude-via-Bedrock calls.
+
+        1. Decode body JSON.
+        2. If Anthropic Claude model: inject RACP into ``system`` field.
+        3. Call real client.
+        4. Decode response body and capture.
+        5. Return raw boto3 response (stream-compatible).
+        """
+        session_id = str(uuid.uuid4())[:8]
+        call_id = str(uuid.uuid4())
+        is_anthropic = any(
+            modelId.startswith(pfx) for pfx in _ANTHROPIC_BEDROCK_PREFIXES
+        )
+
+        # Parse body for prompt text and possible injection
+        body_dict = self._parse_body(body)
+        prompt_text = _extract_prompt_text(body_dict.get("messages", []))
+
+        if is_anthropic and not is_disabled():
+            body = self._inject_anthropic(body_dict, prompt_text, session_id, call_id, modelId)
+        elif not is_anthropic:
+            logger.debug(
+                "devtorch: Bedrock model %s is not Anthropic Claude — "
+                "passing through without RACP injection",
+                modelId,
+            )
+
+        # Call the real Bedrock client
+        try:
+            response = self._client.invoke_model(modelId=modelId, body=body, **kwargs)
+        except Exception:
+            self._record_llm_outcome(
+                prompt_text=prompt_text,
+                response_text="",
+                thinking_blocks=[],
+                session_id=session_id,
+                call_id=call_id,
+                model_name=modelId,
+                outcome="failure",
+            )
+            raise
+
+        # Capture response for Anthropic models
+        if is_anthropic:
+            text, thinking_blocks = self._capture_response(response, session_id)
+            self._record_llm_outcome(
+                prompt_text=prompt_text,
+                response_text=text,
+                thinking_blocks=thinking_blocks,
+                session_id=session_id,
+                call_id=call_id,
+                model_name=modelId,
+                outcome="success",
+            )
+
+        return response
+
+    # ------------------------------------------------------------------
+    # Injection
+    # ------------------------------------------------------------------
+
+    def _parse_body(self, body: Any) -> dict:
+        """Best-effort parse of the Bedrock request body into a dict."""
+        try:
+            if isinstance(body, (bytes, bytearray)):
+                return json.loads(body.decode("utf-8"))
+            elif isinstance(body, str):
+                return json.loads(body)
+            elif isinstance(body, dict):
+                return dict(body)
+        except Exception:
+            pass
+        return {}
+
+    def _inject_anthropic(
+        self,
+        body_dict: dict,
+        prompt_text: str = "",
+        session_id: str = "",
+        call_id: str = "",
+        model_name: str = "",
+    ) -> str:
+        """Inject RACP prefix into the ``system`` field of an Anthropic body."""
+        if not self._injector:
+            return json.dumps(body_dict)
+
+        try:
+            prefix, _ = self._injector.build_system_prefix_with_provenance(
+                prompt_text=prompt_text or None,
+                session_id=session_id,
+                call_id=call_id,
+                model_name=model_name,
+            )
+            if prefix:
+                existing_system = body_dict.get("system", "")
+                body_dict["system"] = (
+                    (prefix + "\n\n" + existing_system).strip()
+                    if existing_system
+                    else prefix
+                )
+
+            return json.dumps(body_dict)
+        except Exception as exc:
+            logger.warning("devtorch: Bedrock RACP injection failed — %s", exc)
+            return json.dumps(body_dict)
+
+    # ------------------------------------------------------------------
+    # Capture
+    # ------------------------------------------------------------------
+
+    def _capture_response(self, response: Any, session_id: str) -> tuple[str, list]:
+        text = ""
+        thinking: list = []
+        if not self._capture or is_disabled():
+            return text, thinking
+        try:
+            # Read and re-inject body stream
+            body_stream = response.get("body")
+            if body_stream is None:
+                return text, thinking
+            raw = body_stream.read()
+            response_dict = json.loads(raw.decode("utf-8"))
+
+            # Extract text from Anthropic Bedrock response format
+            for block in response_dict.get("content", []):
+                if isinstance(block, dict) and block.get("type") == "text":
+                    text += block.get("text", "")
+
+            # Extract thinking blocks
+            try:
+                from devtorch_core.parser.thinking import extract_thinking_blocks
+
+                thinking = extract_thinking_blocks(response_dict)
+            except Exception:
+                pass
+
+            self._capture.capture(text, thinking, session_id)
+
+            # Strip inline <thinking> blocks from the response body before re-wrapping
+            try:
+                from devtorch_core.reasoning_plus.capture import strip_inline_thinking
+
+                for block in response_dict.get("content", []):
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        block["text"] = strip_inline_thinking(block.get("text", ""))
+                raw = json.dumps(response_dict).encode("utf-8")
+            except Exception as exc:
+                logger.warning("devtorch: Bedrock strip thinking failed — %s", exc)
+
+            # Wrap body bytes back into a StreamingBody-compatible object
+            # so the caller can still read the response
+            import io
+
+            response["body"] = _BytesBodyWrapper(raw)
+        except Exception as exc:
+            logger.warning("devtorch: Bedrock _capture_response failed — %s", exc)
+        return text, thinking
+
+    def _record_llm_outcome(
+        self,
+        *,
+        prompt_text: str,
+        response_text: str,
+        thinking_blocks: list,
+        session_id: str,
+        call_id: str,
+        model_name: str,
+        outcome: str,
+    ) -> None:
+        if not self._capture or is_disabled():
+            return
+        try:
+            reasoning = _reasoning_text(thinking_blocks)
+            self._capture.record_learning_call(
+                call_type="llm",
+                reasoning=reasoning,
+                session_id=session_id,
+                call_id=call_id,
+                name=model_name,
+                input=prompt_text,
+                output=response_text,
+                outcome=outcome,
+            )
+        except Exception as exc:
+            logger.warning("devtorch: _record_llm_outcome failed — %s", exc)
+
+    # ------------------------------------------------------------------
+    # Proxy everything else to the underlying boto3 client
+    # ------------------------------------------------------------------
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._client, name)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _extract_prompt_text(messages: list) -> str:
+    """Best-effort concatenation of user-visible message text for smart context."""
+    parts: list[str] = []
+    for m in messages or []:
+        if isinstance(m, dict):
+            role = m.get("role", "")
+            content = m.get("content", "")
+            if role in ("user", "system") and content:
+                if isinstance(content, str):
+                    parts.append(content)
+                elif isinstance(content, list):
+                    for item in content:
+                        if isinstance(item, dict) and item.get("type") == "text":
+                            parts.append(item.get("text", ""))
+                        elif isinstance(item, str):
+                            parts.append(item)
+        elif hasattr(m, "role") and hasattr(m, "content"):
+            role = getattr(m, "role", "")
+            content = getattr(m, "content", "")
+            if role in ("user", "system") and content:
+                parts.append(str(content))
+    return "\n".join(parts)[:4000]
+
+
+# ---------------------------------------------------------------------------
+# Helper: re-wrap consumed body bytes
+# ---------------------------------------------------------------------------
+
+
+class _BytesBodyWrapper:
+    """
+    Minimal wrapper that restores a boto3-style StreamingBody after we've
+    read all bytes for capture.
+    """
+
+    def __init__(self, data: bytes) -> None:
+        import io
+
+        self._stream = io.BytesIO(data)
+
+    def read(self, amt: int = -1) -> bytes:  # type: ignore[override]
+        return self._stream.read() if amt == -1 else self._stream.read(amt)
+
+    def __iter__(self) -> Any:
+        return iter(self._stream)