axonflow 0.4.0__py3-none-any.whl

axonflow/exceptions.py

@@ -0,0 +1,103 @@
+"""AxonFlow SDK Exceptions.
+
+Custom exception hierarchy for clear error handling.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class AxonFlowError(Exception):
+    """Base exception for all AxonFlow errors."""
+
+    def __init__(self, message: str, details: dict[str, Any] | None = None) -> None:
+        self.message = message
+        self.details = details or {}
+        super().__init__(message)
+
+
+class ConfigurationError(AxonFlowError):
+    """Invalid configuration."""
+
+
+class AuthenticationError(AxonFlowError):
+    """Authentication failed."""
+
+
+class PolicyViolationError(AxonFlowError):
+    """Request blocked by policy."""
+
+    def __init__(
+        self,
+        message: str,
+        policy: str | None = None,
+        block_reason: str | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            details={"policy": policy, "block_reason": block_reason},
+        )
+        self.policy = policy
+        self.block_reason = block_reason
+
+
+class RateLimitError(AxonFlowError):
+    """Rate limit exceeded."""
+
+    def __init__(
+        self,
+        message: str,
+        limit: int,
+        remaining: int,
+        reset_at: str | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            details={"limit": limit, "remaining": remaining, "reset_at": reset_at},
+        )
+        self.limit = limit
+        self.remaining = remaining
+        self.reset_at = reset_at
+
+
+class ConnectionError(AxonFlowError):
+    """Connection to AxonFlow Agent failed."""
+
+
+class TimeoutError(AxonFlowError):
+    """Request timed out."""
+
+
+class ConnectorError(AxonFlowError):
+    """MCP connector error."""
+
+    def __init__(
+        self,
+        message: str,
+        connector: str | None = None,
+        operation: str | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            details={"connector": connector, "operation": operation},
+        )
+        self.connector = connector
+        self.operation = operation
+
+
+class PlanExecutionError(AxonFlowError):
+    """Multi-agent plan execution failed."""
+
+    def __init__(
+        self,
+        message: str,
+        plan_id: str | None = None,
+        step: str | None = None,
+    ) -> None:
+        super().__init__(
+            message,
+            details={"plan_id": plan_id, "step": step},
+        )
+        self.plan_id = plan_id
+        self.step = step

axonflow/interceptors/__init__.py

@@ -0,0 +1,20 @@
+"""AxonFlow LLM Provider Interceptors.
+
+Interceptors allow transparent governance integration with popular LLM providers.
+"""
+
+from axonflow.interceptors.anthropic import wrap_anthropic_client
+from axonflow.interceptors.base import BaseInterceptor
+from axonflow.interceptors.bedrock import wrap_bedrock_client
+from axonflow.interceptors.gemini import wrap_gemini_model
+from axonflow.interceptors.ollama import wrap_ollama_client
+from axonflow.interceptors.openai import wrap_openai_client
+
+__all__ = [
+    "BaseInterceptor",
+    "wrap_openai_client",
+    "wrap_anthropic_client",
+    "wrap_gemini_model",
+    "wrap_ollama_client",
+    "wrap_bedrock_client",
+]

axonflow/interceptors/anthropic.py

@@ -0,0 +1,184 @@
+"""Anthropic Interceptor for transparent governance.
+
+Wraps Anthropic client to automatically apply AxonFlow governance
+without changing application code.
+
+Example:
+    >>> from anthropic import Anthropic
+    >>> from axonflow import AxonFlow
+    >>> from axonflow.interceptors.anthropic import wrap_anthropic_client
+    >>>
+    >>> anthropic = Anthropic()
+    >>> axonflow = AxonFlow(...)
+    >>>
+    >>> # Wrap the client - governance is now automatic
+    >>> wrapped = wrap_anthropic_client(anthropic, axonflow)
+    >>>
+    >>> # Use as normal - governance happens invisibly
+    >>> response = wrapped.messages.create(
+    ...     model="claude-3-sonnet-20240229",
+    ...     max_tokens=1024,
+    ...     messages=[{"role": "user", "content": "Hello!"}]
+    ... )
+"""
+
+from __future__ import annotations
+
+import asyncio
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from axonflow.exceptions import PolicyViolationError
+from axonflow.interceptors.base import BaseInterceptor
+
+if TYPE_CHECKING:
+    from axonflow.client import AxonFlow
+
+T = TypeVar("T")
+
+
+class AnthropicInterceptor(BaseInterceptor):
+    """Interceptor for Anthropic client."""
+
+    def get_provider_name(self) -> str:
+        """Get the provider name."""
+        return "anthropic"
+
+    def extract_prompt(self, *_args: Any, **kwargs: Any) -> str:
+        """Extract prompt from messages arguments."""
+        messages = kwargs.get("messages", [])
+        parts = []
+        for m in messages:
+            if isinstance(m, dict):
+                content = m.get("content", "")
+                if isinstance(content, str):
+                    parts.append(content)
+                elif isinstance(content, list):
+                    # Handle content blocks
+                    for block in content:
+                        if isinstance(block, dict) and block.get("type") == "text":
+                            parts.append(block.get("text", ""))
+        return " ".join(parts)
+
+    def wrap(self, client: Any) -> Any:
+        """Wrap Anthropic client with governance."""
+        return wrap_anthropic_client(client, self.axonflow, user_token=self.user_token)
+
+
+def wrap_anthropic_client(
+    anthropic_client: Any,
+    axonflow: AxonFlow,
+    *,
+    user_token: str = "",
+) -> Any:
+    """Wrap Anthropic client with AxonFlow governance.
+
+    Args:
+        anthropic_client: Anthropic client to wrap
+        axonflow: AxonFlow client for governance
+        user_token: User token for policy evaluation
+
+    Returns:
+        Wrapped Anthropic client with automatic governance
+    """
+    original_create = anthropic_client.messages.create
+
+    def _extract_prompt(kwargs: dict[str, Any]) -> str:
+        """Extract prompt from messages."""
+        messages = kwargs.get("messages", [])
+        parts = []
+        for m in messages:
+            if isinstance(m, dict):
+                content = m.get("content", "")
+                if isinstance(content, str):
+                    parts.append(content)
+                elif isinstance(content, list):
+                    for block in content:
+                        if isinstance(block, dict) and block.get("type") == "text":
+                            parts.append(block.get("text", ""))
+        return " ".join(parts)
+
+    def _get_loop() -> asyncio.AbstractEventLoop:
+        """Get or create event loop."""
+        try:
+            return asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
+
+    if asyncio.iscoroutinefunction(original_create):
+
+        @wraps(original_create)
+        async def async_wrapped_create(*args: Any, **kwargs: Any) -> Any:
+            prompt = _extract_prompt(kwargs)
+
+            # Check with AxonFlow
+            response = await axonflow.execute_query(
+                user_token=user_token,
+                query=prompt,
+                request_type="llm_chat",
+                context={
+                    "provider": "anthropic",
+                    "model": kwargs.get("model", "claude-3-sonnet"),
+                    "parameters": {
+                        k: v
+                        for k, v in kwargs.items()
+                        if k not in ("messages", "model", "max_tokens")
+                    },
+                },
+            )
+
+            if response.blocked:
+                raise PolicyViolationError(response.block_reason or "Request blocked by policy")
+
+            # Call original
+            return await original_create(*args, **kwargs)
+
+        anthropic_client.messages.create = async_wrapped_create
+    else:
+
+        @wraps(original_create)
+        def sync_wrapped_create(*args: Any, **kwargs: Any) -> Any:
+            prompt = _extract_prompt(kwargs)
+
+            # Check with AxonFlow (sync)
+            loop = _get_loop()
+            response = loop.run_until_complete(
+                axonflow.execute_query(
+                    user_token=user_token,
+                    query=prompt,
+                    request_type="llm_chat",
+                    context={
+                        "provider": "anthropic",
+                        "model": kwargs.get("model", "claude-3-sonnet"),
+                    },
+                )
+            )
+
+            if response.blocked:
+                raise PolicyViolationError(response.block_reason or "Request blocked by policy")
+
+            # Call original
+            return original_create(*args, **kwargs)
+
+        anthropic_client.messages.create = sync_wrapped_create
+
+    return anthropic_client
+
+
+def create_anthropic_wrapper(axonflow: AxonFlow, user_token: str = "") -> Callable[[Any], Any]:
+    """Create a wrapper function for Anthropic client.
+
+    Args:
+        axonflow: AxonFlow client for governance
+        user_token: User token for policy evaluation
+
+    Returns:
+        Wrapper function that takes an Anthropic client
+    """
+
+    def wrapper(anthropic_client: Any) -> Any:
+        return wrap_anthropic_client(anthropic_client, axonflow, user_token=user_token)
+
+    return wrapper

axonflow/interceptors/base.py

@@ -0,0 +1,58 @@
+"""Base Interceptor class for LLM providers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from axonflow.client import AxonFlow
+
+
+class BaseInterceptor(ABC):
+    """Base class for LLM provider interceptors.
+
+    Interceptors wrap LLM client methods to automatically apply
+    AxonFlow governance without changing application code.
+    """
+
+    def __init__(self, axonflow: AxonFlow, user_token: str = "") -> None:
+        """Initialize interceptor.
+
+        Args:
+            axonflow: AxonFlow client for governance
+            user_token: User token for policy evaluation
+        """
+        self.axonflow = axonflow
+        self.user_token = user_token
+
+    @abstractmethod
+    def wrap(self, client: Any) -> Any:
+        """Wrap an LLM client with governance.
+
+        Args:
+            client: The LLM client to wrap
+
+        Returns:
+            Wrapped client with automatic governance
+        """
+
+    @abstractmethod
+    def extract_prompt(self, *args: Any, **kwargs: Any) -> str:
+        """Extract prompt from method arguments.
+
+        Args:
+            args: Positional arguments
+            kwargs: Keyword arguments
+
+        Returns:
+            Extracted prompt string
+        """
+
+    @abstractmethod
+    def get_provider_name(self) -> str:
+        """Get the provider name.
+
+        Returns:
+            Provider name (e.g., "openai", "anthropic")
+        """

axonflow/interceptors/bedrock.py

@@ -0,0 +1,231 @@
+"""AWS Bedrock Interceptor for transparent governance.
+
+Wraps AWS Bedrock Runtime client to automatically apply AxonFlow governance
+without changing application code.
+
+Bedrock uses AWS IAM authentication (no API keys required).
+Supports multiple model providers: Anthropic Claude, Amazon Titan, Meta Llama, etc.
+
+Example:
+    >>> import boto3
+    >>> from axonflow import AxonFlow
+    >>> from axonflow.interceptors.bedrock import wrap_bedrock_client
+    >>>
+    >>> bedrock = boto3.client('bedrock-runtime', region_name='us-east-1')
+    >>> axonflow = AxonFlow(...)
+    >>>
+    >>> # Wrap the client - governance is now automatic
+    >>> wrapped = wrap_bedrock_client(bedrock, axonflow, user_token="user-123")
+    >>>
+    >>> # Use as normal - governance happens invisibly
+    >>> response = wrapped.invoke_model(
+    ...     modelId='anthropic.claude-3-sonnet-20240229-v1:0',
+    ...     body=json.dumps({...})
+    ... )
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from axonflow.exceptions import PolicyViolationError
+from axonflow.interceptors.base import BaseInterceptor
+
+if TYPE_CHECKING:
+    from axonflow.client import AxonFlow
+
+T = TypeVar("T")
+
+
+# Common Bedrock model IDs
+class BedrockModels:
+    """Common Bedrock model identifiers."""
+
+    # Anthropic Claude models
+    CLAUDE_3_OPUS = "anthropic.claude-3-opus-20240229-v1:0"
+    CLAUDE_3_SONNET = "anthropic.claude-3-sonnet-20240229-v1:0"
+    CLAUDE_3_HAIKU = "anthropic.claude-3-haiku-20240307-v1:0"
+    CLAUDE_2_1 = "anthropic.claude-v2:1"
+    CLAUDE_2 = "anthropic.claude-v2"
+    CLAUDE_INSTANT = "anthropic.claude-instant-v1"
+
+    # Amazon Titan models
+    TITAN_TEXT_EXPRESS = "amazon.titan-text-express-v1"
+    TITAN_TEXT_LITE = "amazon.titan-text-lite-v1"
+    TITAN_TEXT_PREMIER = "amazon.titan-text-premier-v1:0"
+
+    # Meta Llama models
+    LLAMA2_13B = "meta.llama2-13b-chat-v1"
+    LLAMA2_70B = "meta.llama2-70b-chat-v1"
+    LLAMA3_8B = "meta.llama3-8b-instruct-v1:0"
+    LLAMA3_70B = "meta.llama3-70b-instruct-v1:0"
+
+
+class BedrockInterceptor(BaseInterceptor):
+    """Interceptor for AWS Bedrock Runtime client."""
+
+    def get_provider_name(self) -> str:
+        """Get the provider name."""
+        return "bedrock"
+
+    def extract_prompt(self, *_args: Any, **kwargs: Any) -> str:
+        """Extract prompt from invoke_model arguments."""
+        body = kwargs.get("body", "")
+        if isinstance(body, (bytes, bytearray)):
+            body = body.decode("utf-8")
+        if isinstance(body, str):
+            try:
+                parsed = json.loads(body)
+                # Claude format
+                if "messages" in parsed:
+                    return " ".join(
+                        m.get("content", "") for m in parsed["messages"] if isinstance(m, dict)
+                    )
+                # Titan format
+                if "inputText" in parsed:
+                    return str(parsed["inputText"])
+                # Generic prompt
+                if "prompt" in parsed:
+                    return str(parsed["prompt"])
+            except json.JSONDecodeError:
+                pass
+        return ""
+
+    def wrap(self, client: Any) -> Any:
+        """Wrap Bedrock client with governance."""
+        return wrap_bedrock_client(client, self.axonflow, user_token=self.user_token)
+
+
+def wrap_bedrock_client(
+    bedrock_client: Any,
+    axonflow: AxonFlow,
+    *,
+    user_token: str = "",
+) -> Any:
+    """Wrap AWS Bedrock Runtime client with AxonFlow governance.
+
+    Args:
+        bedrock_client: AWS Bedrock Runtime client (boto3.client('bedrock-runtime'))
+        axonflow: AxonFlow client for governance
+        user_token: User token for policy evaluation
+
+    Returns:
+        Wrapped Bedrock client with automatic governance
+    """
+
+    def _extract_prompt(body: Any, _model_id: str) -> str:
+        """Extract prompt from request body."""
+        if isinstance(body, (bytes, bytearray)):
+            body = body.decode("utf-8")
+        if isinstance(body, str):
+            try:
+                parsed = json.loads(body)
+                # Claude format
+                if "messages" in parsed:
+                    return " ".join(
+                        m.get("content", "") for m in parsed["messages"] if isinstance(m, dict)
+                    )
+                # Titan format
+                if "inputText" in parsed:
+                    return str(parsed["inputText"])
+                # Generic
+                if "prompt" in parsed:
+                    return str(parsed["prompt"])
+            except json.JSONDecodeError:
+                pass
+        return ""
+
+    def _get_loop() -> asyncio.AbstractEventLoop:
+        """Get or create event loop."""
+        try:
+            return asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
+
+    # Wrap invoke_model
+    if hasattr(bedrock_client, "invoke_model"):
+        original_invoke = bedrock_client.invoke_model
+
+        @wraps(original_invoke)
+        def sync_wrapped_invoke(*args: Any, **kwargs: Any) -> Any:
+            model_id = kwargs.get("modelId", "unknown")
+            body = kwargs.get("body", "")
+            prompt = _extract_prompt(body, model_id)
+
+            loop = _get_loop()
+            response = loop.run_until_complete(
+                axonflow.execute_query(
+                    user_token=user_token,
+                    query=prompt,
+                    request_type="llm_chat",
+                    context={
+                        "provider": "bedrock",
+                        "model": model_id,
+                    },
+                )
+            )
+
+            if response.blocked:
+                raise PolicyViolationError(response.block_reason or "Request blocked by policy")
+
+            return original_invoke(*args, **kwargs)
+
+        bedrock_client.invoke_model = sync_wrapped_invoke
+
+    # Wrap invoke_model_with_response_stream
+    if hasattr(bedrock_client, "invoke_model_with_response_stream"):
+        original_stream = bedrock_client.invoke_model_with_response_stream
+
+        @wraps(original_stream)
+        def sync_wrapped_stream(*args: Any, **kwargs: Any) -> Any:
+            model_id = kwargs.get("modelId", "unknown")
+            body = kwargs.get("body", "")
+            prompt = _extract_prompt(body, model_id)
+
+            loop = _get_loop()
+            response = loop.run_until_complete(
+                axonflow.execute_query(
+                    user_token=user_token,
+                    query=prompt,
+                    request_type="llm_chat",
+                    context={
+                        "provider": "bedrock",
+                        "model": model_id,
+                        "streaming": True,
+                    },
+                )
+            )
+
+            if response.blocked:
+                raise PolicyViolationError(response.block_reason or "Request blocked by policy")
+
+            return original_stream(*args, **kwargs)
+
+        bedrock_client.invoke_model_with_response_stream = sync_wrapped_stream
+
+    return bedrock_client
+
+
+def create_bedrock_wrapper(
+    axonflow: AxonFlow,
+    user_token: str = "",
+) -> Callable[[Any], Any]:
+    """Create a wrapper function for Bedrock client.
+
+    Args:
+        axonflow: AxonFlow client for governance
+        user_token: User token for policy evaluation
+
+    Returns:
+        Wrapper function that takes a Bedrock client
+    """
+
+    def wrapper(bedrock_client: Any) -> Any:
+        return wrap_bedrock_client(bedrock_client, axonflow, user_token=user_token)
+
+    return wrapper