argus-cloud-optimizer 0.2.0__py3-none-any.whl

ai/azure_openai.py

@@ -0,0 +1,241 @@
+"""
+AI provider backed by Azure OpenAI (GPT-4o).
+
+Authentication uses DefaultAzureCredential — run:
+    az login
+
+No API key needed when running on Azure Functions with a managed identity.
+
+Environment variables:
+    AZURE_OPENAI_ENDPOINT    Azure OpenAI resource endpoint (required)
+                             e.g. https://my-resource.openai.azure.com/
+    AZURE_OPENAI_DEPLOYMENT  Deployment name (default: gpt-4o)
+    AZURE_OPENAI_API_VERSION API version (default: 2024-10-21)
+    AZURE_OPENAI_API_KEY     Optional — use only for local dev without az login
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+import openai
+import structlog
+
+from ai.base import AIProvider, AIResponse, Message, Tool, ToolCall
+
+logger = structlog.get_logger(__name__)
+
+MAX_RETRIES = 3
+_BASE_DELAY = 1.0
+
+
+class AzureOpenAIProvider(AIProvider):
+    """
+    AI provider backed by Azure OpenAI GPT-4o.
+    Uses DefaultAzureCredential (managed identity / az login) — no API key needed
+    when running on Azure infrastructure.
+
+    Falls back to AZURE_OPENAI_API_KEY for local dev without az login.
+    """
+
+    DEFAULT_DEPLOYMENT = "gpt-4o"
+    DEFAULT_API_VERSION = "2024-10-21"
+    DEFAULT_MAX_TOKENS = 4096
+    DEFAULT_TEMPERATURE = 0.0
+
+    def __init__(
+        self,
+        endpoint: str | None = None,
+        deployment: str | None = None,
+        api_version: str | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        api_key: str | None = None,
+        temperature: float | None = None,
+    ) -> None:
+        from core.config import get_settings
+
+        cfg = get_settings().ai
+        self._endpoint = endpoint or cfg.azure_openai_endpoint
+        if not self._endpoint:
+            raise EnvironmentError(
+                "AZURE_OPENAI_ENDPOINT is not set. "
+                "Set it in .env or pass endpoint= explicitly. "
+                "Example: https://my-resource.openai.azure.com/"
+            )
+        self._deployment = deployment or cfg.resolved_model("azure_openai")
+        self._api_version = api_version or cfg.azure_openai_api_version
+        self._max_tokens = max_tokens
+        self._temperature = temperature if temperature is not None else cfg.temperature
+
+        resolved_key = api_key or cfg.azure_openai_api_key
+        if resolved_key:
+            self._client = openai.AzureOpenAI(
+                azure_endpoint=self._endpoint,
+                api_key=resolved_key,
+                api_version=self._api_version,
+            )
+            self._credential = None
+        else:
+            # DefaultAzureCredential: works with managed identity, az login, env vars
+            from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+            credential = DefaultAzureCredential()
+            token_provider = get_bearer_token_provider(
+                credential,
+                "https://cognitiveservices.azure.com/.default",
+            )
+            self._client = openai.AzureOpenAI(
+                azure_endpoint=self._endpoint,
+                azure_ad_token_provider=token_provider,
+                api_version=self._api_version,
+            )
+            self._credential = credential
+
+    @classmethod
+    def from_env(cls) -> "AzureOpenAIProvider":
+        return cls()
+
+    def chat(
+        self,
+        messages: list[Message],
+        tools: list[Tool],
+        system_prompt: str | None = None,
+    ) -> AIResponse:
+        openai_messages = self._build_messages(messages, system_prompt)
+        openai_tools = [self._to_openai_tool(t) for t in tools] if tools else None
+
+        kwargs: dict[str, Any] = {
+            "model": self._deployment,
+            "messages": openai_messages,
+            "max_tokens": self._max_tokens,
+            "temperature": self._temperature,
+        }
+        if openai_tools:
+            kwargs["tools"] = openai_tools
+            kwargs["tool_choice"] = "auto"
+
+        response = self._call_with_retry(kwargs)
+        return self._parse_response(response)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _call_with_retry(self, kwargs: dict[str, Any]) -> Any:
+        delay = _BASE_DELAY
+        for attempt in range(MAX_RETRIES):
+            try:
+                return self._client.chat.completions.create(**kwargs)
+            except openai.RateLimitError:
+                if attempt < MAX_RETRIES - 1:
+                    logger.warning(
+                        "Azure OpenAI rate limited (attempt %d/%d), retrying in %.1fs",
+                        attempt + 1,
+                        MAX_RETRIES,
+                        delay,
+                    )
+                    time.sleep(delay)
+                    delay *= 2
+                else:
+                    raise
+            except openai.AuthenticationError as exc:
+                raise EnvironmentError(
+                    "Azure OpenAI authentication failed. "
+                    "Run 'az login' or set AZURE_OPENAI_API_KEY."
+                ) from exc
+        raise RuntimeError("Unreachable")  # pragma: no cover
+
+    def _build_messages(
+        self,
+        messages: list[Message],
+        system_prompt: str | None,
+    ) -> list[dict[str, Any]]:
+        result: list[dict[str, Any]] = []
+
+        if system_prompt:
+            result.append({"role": "system", "content": system_prompt})
+
+        for msg in messages:
+            if msg.role == "user":
+                if msg.tool_results:
+                    for tr in msg.tool_results:
+                        result.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": tr.tool_call_id,
+                                "content": tr.content,
+                            }
+                        )
+                else:
+                    result.append({"role": "user", "content": msg.text or ""})
+
+            else:
+                # assistant
+                tool_calls_out = []
+                for tc in msg.tool_calls:
+                    tool_calls_out.append(
+                        {
+                            "id": tc.id,
+                            "type": "function",
+                            "function": {
+                                "name": tc.name,
+                                "arguments": json.dumps(tc.arguments),
+                            },
+                        }
+                    )
+                assistant_msg: dict[str, Any] = {
+                    "role": "assistant",
+                    "content": msg.text or "",
+                }
+                if tool_calls_out:
+                    assistant_msg["tool_calls"] = tool_calls_out
+                result.append(assistant_msg)
+
+        return result
+
+    def _to_openai_tool(self, tool: Tool) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": tool.input_schema,
+            },
+        }
+
+    def _parse_response(self, response: Any) -> AIResponse:
+        choice = response.choices[0]
+        message = choice.message
+        stop_reason = choice.finish_reason  # "stop" | "tool_calls" | "length"
+
+        text: str | None = message.content or None
+        tool_calls: list[ToolCall] = []
+
+        if message.tool_calls:
+            for tc in message.tool_calls:
+                tool_calls.append(
+                    ToolCall(
+                        id=tc.id,
+                        name=tc.function.name,
+                        arguments=json.loads(tc.function.arguments),
+                    )
+                )
+
+        # Normalise to internal vocabulary
+        if tool_calls:
+            stop_reason = "tool_use"
+        elif stop_reason == "stop":
+            stop_reason = "end_turn"
+        elif stop_reason == "length":
+            stop_reason = "max_tokens"
+
+        usage = getattr(response, "usage", None)
+        return AIResponse(
+            stop_reason=stop_reason,
+            text=text,
+            tool_calls=tool_calls,
+            input_tokens=getattr(usage, "prompt_tokens", 0) if usage else 0,
+            output_tokens=getattr(usage, "completion_tokens", 0) if usage else 0,
+        )

ai/base.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ToolCall:
+    """A tool invocation requested by the AI."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+@dataclass
+class ToolResult:
+    """The result of executing a tool call."""
+
+    tool_call_id: str
+    content: str
+    is_error: bool = False
+
+
+@dataclass
+class Message:
+    """
+    A single turn in the agent conversation.
+    role: "user" | "assistant"
+    Exactly one of text, tool_calls, or tool_results will be populated.
+    """
+
+    role: str
+    text: str | None = None
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    tool_results: list[ToolResult] = field(default_factory=list)
+
+
+@dataclass
+class Tool:
+    """Definition of a tool the AI can call."""
+
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+@dataclass
+class AIResponse:
+    """Parsed response from an AI provider."""
+
+    stop_reason: str  # "tool_use" | "end_turn" | "max_tokens"
+    text: str | None
+    tool_calls: list[ToolCall]
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+
+class AIProvider(ABC):
+    """
+    Abstract AI provider. One implementation per model family.
+    The agent loop only ever calls chat() — never raw SDK methods.
+    """
+
+    @abstractmethod
+    def chat(
+        self,
+        messages: list[Message],
+        tools: list[Tool],
+        system_prompt: str | None = None,
+    ) -> AIResponse:
+        """
+        Send the conversation to the AI and get a response.
+        system_prompt is passed separately so each provider can handle it
+        in the way their API expects (e.g. Anthropic has a dedicated system param).
+        """
+        ...

ai/bedrock.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import boto3
+import structlog
+from botocore.exceptions import ClientError
+
+from ai.base import AIProvider, AIResponse, Message, Tool, ToolCall
+
+logger = structlog.get_logger(__name__)
+
+MAX_RETRIES = 3
+_BASE_DELAY = 1.0  # seconds; doubles each retry
+
+
+class BedrockProvider(AIProvider):
+    """
+    AI provider backed by Amazon Bedrock Converse API.
+    Uses the execution role when running inside Lambda — no API keys needed.
+    Configure via BEDROCK_MODEL_ID and BEDROCK_REGION env vars.
+    """
+
+    DEFAULT_MODEL = "anthropic.claude-sonnet-4-6"
+    DEFAULT_REGION = "us-east-1"
+    DEFAULT_MAX_TOKENS = 4096
+    DEFAULT_TEMPERATURE = 0.0
+
+    def __init__(
+        self,
+        model_id: str | None = None,
+        region: str | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        temperature: float | None = None,
+        session: Any = None,
+    ) -> None:
+        from core.config import get_settings
+
+        cfg = get_settings().ai
+        self._model_id = model_id or cfg.resolved_model("bedrock")
+        resolved_region = region or cfg.bedrock_region
+        self._max_tokens = max_tokens
+        self._temperature = temperature if temperature is not None else cfg.temperature
+
+        boto_session = session or boto3.Session(region_name=resolved_region)
+        self._client = boto_session.client(
+            "bedrock-runtime", region_name=resolved_region
+        )
+
+    def chat(
+        self,
+        messages: list[Message],
+        tools: list[Tool],
+        system_prompt: str | None = None,
+    ) -> AIResponse:
+        kwargs: dict[str, Any] = {
+            "modelId": self._model_id,
+            "messages": [self._to_bedrock_message(m) for m in messages],
+            "inferenceConfig": {
+                "maxTokens": self._max_tokens,
+                "temperature": self._temperature,
+            },
+        }
+        if system_prompt:
+            kwargs["system"] = [{"text": system_prompt}]
+        if tools:
+            kwargs["toolConfig"] = {"tools": [self._to_bedrock_tool(t) for t in tools]}
+
+        response = self._call_with_retry(kwargs)
+        return self._parse_response(response)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _call_with_retry(self, kwargs: dict[str, Any]) -> Any:
+        delay = _BASE_DELAY
+        for attempt in range(MAX_RETRIES):
+            try:
+                return self._client.converse(**kwargs)
+            except ClientError as exc:
+                code = exc.response["Error"]["Code"]
+                if code == "ThrottlingException" and attempt < MAX_RETRIES - 1:
+                    logger.warning(
+                        "Bedrock throttled (attempt %d/%d), retrying in %.1fs",
+                        attempt + 1,
+                        MAX_RETRIES,
+                        delay,
+                    )
+                    time.sleep(delay)
+                    delay *= 2
+                else:
+                    raise
+        raise RuntimeError("Unreachable")  # pragma: no cover
+
+    def _to_bedrock_message(self, msg: Message) -> dict[str, Any]:
+        if msg.role == "user":
+            if msg.tool_results:
+                return {
+                    "role": "user",
+                    "content": [
+                        {
+                            "toolResult": {
+                                "toolUseId": tr.tool_call_id,
+                                "content": [{"text": tr.content}],
+                                **({"status": "error"} if tr.is_error else {}),
+                            }
+                        }
+                        for tr in msg.tool_results
+                    ],
+                }
+            return {"role": "user", "content": [{"text": msg.text or ""}]}
+
+        # assistant
+        content: list[dict[str, Any]] = []
+        if msg.text:
+            content.append({"text": msg.text})
+        for tc in msg.tool_calls:
+            content.append(
+                {
+                    "toolUse": {
+                        "toolUseId": tc.id,
+                        "name": tc.name,
+                        "input": tc.arguments,
+                    }
+                }
+            )
+        return {"role": "assistant", "content": content}
+
+    def _to_bedrock_tool(self, tool: Tool) -> dict[str, Any]:
+        return {
+            "toolSpec": {
+                "name": tool.name,
+                "description": tool.description,
+                "inputSchema": {"json": tool.input_schema},
+            }
+        }
+
+    def _parse_response(self, response: dict[str, Any]) -> AIResponse:
+        content_blocks: list[dict[str, Any]] = (
+            response.get("output", {}).get("message", {}).get("content", [])
+        )
+        stop_reason: str = response.get("stopReason", "end_turn")
+
+        tool_calls: list[ToolCall] = []
+        text: str | None = None
+
+        for block in content_blocks:
+            if "toolUse" in block:
+                tu = block["toolUse"]
+                tool_calls.append(
+                    ToolCall(
+                        id=tu["toolUseId"],
+                        name=tu["name"],
+                        arguments=dict(tu.get("input", {})),
+                    )
+                )
+            elif "text" in block:
+                text = block["text"]
+
+        usage = response.get("usage", {})
+        return AIResponse(
+            stop_reason=stop_reason,
+            text=text,
+            tool_calls=tool_calls,
+            input_tokens=usage.get("inputTokens", 0),
+            output_tokens=usage.get("outputTokens", 0),
+        )