turingpulse-sdk-bedrock 1.0.0__tar.gz

turingpulse_sdk_bedrock-1.0.0/.gitignore

@@ -0,0 +1,42 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+
+# Distribution / packaging
+dist/
+build/
+*.egg-info/
+
+# Database files
+*.db
+*.sqlite3
+
+# Environment variables
+.env
+.env.local
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+
+# Logs
+*.log
+logs/
+
+# OS files
+.DS_Store
+Thumbs.db

turingpulse_sdk_bedrock-1.0.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: turingpulse-sdk-bedrock
+Version: 1.0.0
+Summary: TuringPulse SDK integration for AWS Bedrock
+License-Expression: Apache-2.0
+Requires-Python: >=3.11
+Requires-Dist: boto3>=1.35.0
+Requires-Dist: turingpulse-sdk>=1.0.0
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'

turingpulse_sdk_bedrock-1.0.0/pyproject.toml

@@ -0,0 +1,17 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "turingpulse-sdk-bedrock"
+version = "1.0.0"
+description = "TuringPulse SDK integration for AWS Bedrock"
+requires-python = ">=3.11"
+license = "Apache-2.0"
+dependencies = [
+    "turingpulse-sdk>=1.0.0",
+    "boto3>=1.35.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.0", "pytest-asyncio>=0.23"]

turingpulse_sdk_bedrock-1.0.0/turingpulse_sdk_bedrock/__init__.py

@@ -0,0 +1,6 @@
+"""TuringPulse SDK integration for AWS Bedrock."""
+
+from ._wrapper import patch_bedrock, unpatch_bedrock
+
+__version__ = "0.1.0"
+__all__ = ["patch_bedrock", "unpatch_bedrock"]

turingpulse_sdk_bedrock-1.0.0/turingpulse_sdk_bedrock/_wrapper.py

@@ -0,0 +1,139 @@
+"""AWS Bedrock monkey-patch instrumentation for TuringPulse."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from contextvars import ContextVar
+from typing import Any, Optional
+
+from turingpulse_sdk import instrument, GovernanceDirective
+from turingpulse_sdk.config import MAX_FIELD_SIZE
+from turingpulse_sdk.context import current_context
+from turingpulse_sdk.exceptions import ConfigurationError
+
+logger = logging.getLogger("turingpulse.sdk.bedrock")
+
+_INSTRUMENTING: ContextVar[bool] = ContextVar("_tp_bedrock_instrumenting", default=False)
+
+_ORIGINAL_CONVERSE: Any = None
+
+
+def patch_bedrock(
+    *,
+    name: str | None = None,
+    governance: Optional[GovernanceDirective] = None,
+) -> None:
+    """Monkey-patch ``botocore`` to intercept ``bedrock-runtime.converse`` calls."""
+    global _ORIGINAL_CONVERSE
+
+    effective_name = name or os.getenv("TP_WORKFLOW_NAME", "")
+    if not effective_name:
+        raise ConfigurationError(
+            "patch_bedrock() requires name= or TP_WORKFLOW_NAME to be set."
+        )
+
+    try:
+        import botocore.client
+    except ImportError as exc:
+        raise ImportError("boto3 package is required: pip install boto3") from exc
+
+    if _ORIGINAL_CONVERSE is not None:
+        logger.warning("Bedrock is already patched — skipping")
+        return
+
+    _ORIGINAL_CONVERSE = botocore.client.ClientCreator.create_client
+
+    original_create = _ORIGINAL_CONVERSE
+
+    def _patched_create_client(self_creator, *args: Any, **kwargs: Any) -> Any:
+        client = original_create(self_creator, *args, **kwargs)
+        service_name = args[0] if args else kwargs.get("service_name", "")
+        if service_name == "bedrock-runtime":
+            _wrap_converse(client, effective_name, governance)
+        return client
+
+    botocore.client.ClientCreator.create_client = _patched_create_client
+    logger.info("Bedrock patched for TuringPulse instrumentation")
+
+
+def unpatch_bedrock() -> None:
+    """Restore original botocore client creator."""
+    global _ORIGINAL_CONVERSE
+    if _ORIGINAL_CONVERSE is None:
+        return
+    try:
+        import botocore.client
+        botocore.client.ClientCreator.create_client = _ORIGINAL_CONVERSE
+    except ImportError:
+        pass
+    _ORIGINAL_CONVERSE = None
+    logger.info("Bedrock unpatched — original methods restored")
+
+
+def _wrap_converse(client: Any, name: str, governance: Optional[GovernanceDirective]) -> None:
+    """Wrap the converse method of a bedrock-runtime client."""
+    original_converse = client.converse
+
+    @instrument(name=name, governance=governance)
+    def _patched_converse(**kwargs: Any) -> Any:
+        if _INSTRUMENTING.get(False):
+            return original_converse(**kwargs)
+        token = _INSTRUMENTING.set(True)
+        try:
+            response = original_converse(**kwargs)
+            _record_bedrock_span(response, kwargs)
+            return response
+        finally:
+            _INSTRUMENTING.reset(token)
+
+    client.converse = _patched_converse
+
+
+def _record_bedrock_span(response: Any, kwargs: dict) -> None:
+    ctx = current_context()
+    if not ctx:
+        return
+    ctx.framework = "bedrock"
+    ctx.node_type = "llm"
+
+    usage = response.get("usage", {}) if isinstance(response, dict) else {}
+    ctx.set_tokens(
+        usage.get("inputTokens", 0),
+        usage.get("outputTokens", 0),
+    )
+
+    model_id = kwargs.get("modelId", "unknown")
+    ctx.set_model(model_id, "bedrock")
+
+    messages = kwargs.get("messages", [])
+    if messages:
+        last_user = next(
+            (m for m in reversed(messages) if m.get("role") == "user"),
+            None,
+        )
+        if last_user:
+            content = last_user.get("content", [])
+            if isinstance(content, list):
+                texts = [c.get("text", "") for c in content if "text" in c]
+                ctx.set_prompt("\n".join(texts)[:MAX_FIELD_SIZE])
+            elif isinstance(content, str):
+                ctx.set_prompt(content[:MAX_FIELD_SIZE])
+
+    output = response.get("output", {}) if isinstance(response, dict) else {}
+    message = output.get("message", {})
+    content = message.get("content", [])
+    if content:
+        texts = [c.get("text", "") for c in content if "text" in c]
+        if texts:
+            ctx.set_io(output_data="\n".join(texts)[:MAX_FIELD_SIZE])
+
+        for c in content:
+            if "toolUse" in c:
+                tool = c["toolUse"]
+                ctx.add_tool_call(
+                    tool_name=tool.get("name", "unknown"),
+                    tool_args=tool.get("input", {}),
+                    tool_id=tool.get("toolUseId"),
+                )