pyrestkit 1.2.0__py3-none-any.whl

pyrestkit/ai/providers/anthropic.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class AnthropicProvider(BaseAIProvider):
+    """
+    Anthropic provider using the Messages API.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Anthropic API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "max_tokens": config.max_tokens or 1024,
+            "messages": [{"role": "user", "content": prompt}],
+        }
+
+        if system_prompt:
+            payload["system"] = system_prompt
+
+        response = self._session.post(
+            _messages_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Anthropic request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Anthropic response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Anthropic response was not a JSON object.")
+
+        content = data.get("content")
+        if isinstance(content, list):
+            parts: list[str] = []
+            for item in content:
+                if isinstance(item, dict):
+                    text = item.get("text")
+                    if isinstance(text, str):
+                        parts.append(text)
+            if parts:
+                return "\n".join(parts)
+
+        raise AIProviderError("Anthropic response did not include text content.")
+
+
+def _messages_url(base_url: str | None) -> str:
+    base = (base_url or "https://api.anthropic.com/v1").rstrip("/")
+    return f"{base}/messages"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers = BaseAIProvider.json_headers(config)
+
+    headers["anthropic-version"] = "2023-06-01"
+
+    if config.api_key is not None:
+        headers["x-api-key"] = config.api_key
+
+    return headers

pyrestkit/ai/providers/azure_openai.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class AzureOpenAIProvider(BaseAIProvider):
+    """
+    Azure OpenAI provider using the chat/completions endpoint.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Azure OpenAI API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": system_prompt or "You are a helpful assistant.",
+                },
+                {"role": "user", "content": prompt},
+            ],
+        }
+
+        response = self._session.post(
+            _chat_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Azure OpenAI request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Azure OpenAI response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Azure OpenAI response was not a JSON object.")
+
+        choices = data.get("choices")
+        if isinstance(choices, list) and choices:
+            first = choices[0]
+            if isinstance(first, dict):
+                message = first.get("message")
+                if isinstance(message, dict):
+                    content = message.get("content")
+                    if isinstance(content, str):
+                        return content
+
+        raise AIProviderError("Azure OpenAI response did not include text content.")
+
+
+def _chat_url(base_url: str | None) -> str:
+    base = (base_url or "https://example.openai.azure.com").rstrip("/")
+    return f"{base}/openai/deployments?api-version=2024-02-01"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers = BaseAIProvider.json_headers(config)
+
+    if config.api_key is not None:
+        headers["api-key"] = config.api_key
+
+    return headers

pyrestkit/ai/providers/base.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import requests
+
+from pyrestkit.ai.config import AIConfig
+
+
+class BaseAIProvider(ABC):
+    """
+    Base class for all AI providers.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        self._session = session or requests.Session()
+
+    @staticmethod
+    def json_headers(config: AIConfig) -> dict[str, str]:
+        """
+        Build common JSON headers.
+        """
+        headers: dict[str, str] = {
+            "Content-Type": "application/json",
+        }
+
+        headers.update(config.headers)
+
+        return headers
+
+    @abstractmethod
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str: ...

pyrestkit/ai/providers/bedrock.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class BedrockProvider(BaseAIProvider):
+    """
+    AWS Bedrock provider using a simple text-generation style payload.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Bedrock API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "inputText": prompt,
+        }
+
+        if system_prompt:
+            payload["systemPrompt"] = system_prompt
+
+        response = self._session.post(
+            _invoke_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Bedrock request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Bedrock response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Bedrock response was not a JSON object.")
+
+        output = data.get("outputText")
+        if isinstance(output, str) and output.strip():
+            return output
+
+        raise AIProviderError("Bedrock response did not include text content.")
+
+
+def _invoke_url(base_url: str | None) -> str:
+    base = (base_url or "https://bedrock-runtime.us-east-1.amazonaws.com").rstrip("/")
+    return f"{base}/model:invoke"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    return BaseAIProvider.json_headers(config)

pyrestkit/ai/providers/cohere.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class CohereProvider(BaseAIProvider):
+    """
+    Cohere provider using the chat endpoint.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Cohere API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "message": prompt,
+        }
+
+        if system_prompt:
+            payload["chat_history"] = [{"role": "SYSTEM", "message": system_prompt}]
+
+        response = self._session.post(
+            _chat_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Cohere request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Cohere response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Cohere response was not a JSON object.")
+
+        texts = data.get("text")
+        if isinstance(texts, str) and texts.strip():
+            return texts
+
+        raise AIProviderError("Cohere response did not include text content.")
+
+
+def _chat_url(base_url: str | None) -> str:
+    base = (base_url or "https://api.cohere.com/v1").rstrip("/")
+    return f"{base}/chat"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers: dict[str, str] = {
+        "Content-Type": "application/json",
+    }
+
+    if config.api_key is not None:
+        headers["Authorization"] = f"Bearer {config.api_key}"
+
+    if config.organization is not None:
+        headers["OpenAI-Organization"] = config.organization
+
+    headers.update(config.headers)
+
+    return headers

pyrestkit/ai/providers/gemini.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class GeminiProvider(BaseAIProvider):
+    """
+    Google Gemini provider using the generateContent endpoint.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Gemini API key is required.")
+
+        content: dict[str, Any] = {
+            "parts": [
+                {
+                    "text": prompt,
+                }
+            ]
+        }
+
+        payload: dict[str, Any] = {
+            "contents": [content],
+        }
+
+        if system_prompt:
+            payload["system_instruction"] = {
+                "parts": [
+                    {
+                        "text": system_prompt,
+                    }
+                ]
+            }
+
+        response = self._session.post(
+            _generate_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Gemini request failed with status "
+                f"{response.status_code}: {response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Gemini response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Gemini response was not a JSON object.")
+
+        candidates = data.get("candidates")
+
+        if isinstance(candidates, list):
+            for candidate in candidates:
+                if not isinstance(candidate, dict):
+                    continue
+
+                candidate_content = candidate.get("content")
+                if not isinstance(candidate_content, dict):
+                    continue
+
+                parts = candidate_content.get("parts")
+                if not isinstance(parts, list):
+                    continue
+
+                texts: list[str] = []
+
+                for part in parts:
+                    if not isinstance(part, dict):
+                        continue
+
+                    text = part.get("text")
+                    if isinstance(text, str):
+                        texts.append(text)
+
+                if texts:
+                    return "\n".join(texts)
+
+        raise AIProviderError("Gemini response did not include text content.")
+
+
+def _generate_url(base_url: str | None) -> str:
+    base = (base_url or "https://generativelanguage.googleapis.com/v1beta").rstrip("/")
+
+    return f"{base}/models/gemini-2.0-flash:generateContent"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers = BaseAIProvider.json_headers(config)
+
+    if config.api_key is not None:
+        headers["x-goog-api-key"] = config.api_key
+
+    return headers

pyrestkit/ai/providers/groq.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class GroqProvider(BaseAIProvider):
+    """
+    Groq provider using the chat/completions endpoint.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Groq API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "messages": [{"role": "user", "content": prompt}],
+        }
+
+        if system_prompt:
+            payload["messages"].insert(0, {"role": "system", "content": system_prompt})
+
+        response = self._session.post(
+            _chat_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Groq request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Groq response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Groq response was not a JSON object.")
+
+        choices = data.get("choices")
+        if isinstance(choices, list) and choices:
+            first = choices[0]
+            if isinstance(first, dict):
+                message = first.get("message")
+                if isinstance(message, dict):
+                    content = message.get("content")
+                    if isinstance(content, str):
+                        return content
+
+        raise AIProviderError("Groq response did not include text content.")
+
+
+def _chat_url(base_url: str | None) -> str:
+    base = (base_url or "https://api.groq.com/openai/v1").rstrip("/")
+    return f"{base}/chat/completions"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers = BaseAIProvider.json_headers(config)
+
+    if config.api_key is not None:
+        headers["Authorization"] = f"Bearer {config.api_key}"
+
+    return headers

pyrestkit/ai/providers/mistral.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyrestkit.ai.config import AIConfig
+from pyrestkit.ai.exceptions import AIConfigurationError, AIProviderError
+from pyrestkit.ai.providers.base import BaseAIProvider
+
+
+class MistralProvider(BaseAIProvider):
+    """
+    Mistral provider using the chat/completions endpoint.
+    """
+
+    def __init__(self, session: Any | None = None) -> None:
+        super().__init__(session=session)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        config: AIConfig,
+        system_prompt: str | None = None,
+    ) -> str:
+        if not config.api_key:
+            raise AIConfigurationError("Mistral API key is required.")
+
+        payload: dict[str, Any] = {
+            "model": config.model,
+            "messages": [{"role": "user", "content": prompt}],
+        }
+
+        if system_prompt:
+            payload["messages"].insert(0, {"role": "system", "content": system_prompt})
+
+        response = self._session.post(
+            _chat_url(config.base_url),
+            json=payload,
+            headers=_headers(config),
+            timeout=config.timeout,
+        )
+
+        if response.status_code >= 400:
+            raise AIProviderError(
+                f"Mistral request failed with status {response.status_code}: "
+                f"{response.text}"
+            )
+
+        try:
+            data = response.json()
+        except ValueError as exc:
+            raise AIProviderError("Mistral response was not valid JSON.") from exc
+
+        if not isinstance(data, dict):
+            raise AIProviderError("Mistral response was not a JSON object.")
+
+        choices = data.get("choices")
+        if isinstance(choices, list) and choices:
+            first = choices[0]
+            if isinstance(first, dict):
+                message = first.get("message")
+                if isinstance(message, dict):
+                    content = message.get("content")
+                    if isinstance(content, str):
+                        return content
+
+        raise AIProviderError("Mistral response did not include text content.")
+
+
+def _chat_url(base_url: str | None) -> str:
+    base = (base_url or "https://api.mistral.ai/v1").rstrip("/")
+    return f"{base}/chat/completions"
+
+
+def _headers(config: AIConfig) -> dict[str, str]:
+    headers: dict[str, str] = {
+        "Content-Type": "application/json",
+    }
+
+    if config.api_key is not None:
+        headers["Authorization"] = f"Bearer {config.api_key}"
+
+    if config.organization is not None:
+        headers["OpenAI-Organization"] = config.organization
+
+    headers.update(config.headers)
+
+    return headers