feat: add ai provider adapters
This commit is contained in:
QiuSW
@@ -0,0 +1,26 @@
|
||||
from .base import (
|
||||
AiCapabilityError,
|
||||
AiProviderConfigError,
|
||||
AiProviderError,
|
||||
AiResponseParseError,
|
||||
ImageGenerationResult,
|
||||
Provider,
|
||||
ResolvedModel,
|
||||
TextGenerationResult,
|
||||
)
|
||||
from .registry import default_registry, get_provider, register_provider, resolve_api_type
|
||||
|
||||
__all__ = [
|
||||
"AiCapabilityError",
|
||||
"AiProviderConfigError",
|
||||
"AiProviderError",
|
||||
"AiResponseParseError",
|
||||
"ImageGenerationResult",
|
||||
"Provider",
|
||||
"ResolvedModel",
|
||||
"TextGenerationResult",
|
||||
"default_registry",
|
||||
"get_provider",
|
||||
"register_provider",
|
||||
"resolve_api_type",
|
||||
]
|
||||
@@ -0,0 +1,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Mapping, Protocol
|
||||
|
||||
|
||||
class AiProviderError(RuntimeError):
|
||||
"""Base error for AI provider failures."""
|
||||
|
||||
|
||||
class AiProviderConfigError(AiProviderError, ValueError):
|
||||
"""Raised when model/provider configuration is invalid."""
|
||||
|
||||
|
||||
class AiCapabilityError(AiProviderError):
|
||||
"""Raised when a provider cannot perform the requested capability."""
|
||||
|
||||
|
||||
class AiResponseParseError(AiProviderError):
|
||||
"""Raised when a provider response cannot be parsed into the expected result."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedModel:
|
||||
"""Runtime model config resolved from future AiModel/ModelAlias records."""
|
||||
|
||||
name: str
|
||||
url: str
|
||||
model: str
|
||||
api_key: str
|
||||
api_type: str = "auto"
|
||||
timeout_seconds: int = 0
|
||||
connect_timeout_seconds: int = 30
|
||||
extra_body: dict[str, Any] = field(default_factory=dict)
|
||||
capabilities: frozenset[str] = field(default_factory=frozenset)
|
||||
|
||||
@classmethod
|
||||
def from_mapping(cls, data: Mapping[str, Any]) -> "ResolvedModel":
|
||||
raw_capabilities = data.get("capabilities") or ()
|
||||
if isinstance(raw_capabilities, str):
|
||||
capabilities = frozenset(
|
||||
item.strip() for item in raw_capabilities.split(",") if item.strip()
|
||||
)
|
||||
else:
|
||||
capabilities = frozenset(str(item) for item in raw_capabilities)
|
||||
|
||||
return cls(
|
||||
name=str(data.get("name") or data.get("model") or ""),
|
||||
url=str(data.get("url") or ""),
|
||||
model=str(data.get("model") or ""),
|
||||
api_key=str(data.get("api_key") or ""),
|
||||
api_type=str(data.get("api_type") or "auto"),
|
||||
timeout_seconds=_to_int(data.get("timeout_seconds"), 0),
|
||||
connect_timeout_seconds=_to_int(data.get("connect_timeout_seconds"), 30),
|
||||
extra_body=dict(data.get("extra_body") or {}),
|
||||
capabilities=capabilities,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TextGenerationResult:
|
||||
text: str
|
||||
titles: tuple[str, ...]
|
||||
model_used: str
|
||||
raw: Mapping[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageGenerationResult:
|
||||
image: bytes
|
||||
model_used: str
|
||||
raw: Mapping[str, Any]
|
||||
|
||||
|
||||
class Provider(Protocol):
|
||||
def capabilities(self) -> set[str]:
|
||||
...
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
resolution: str = "1K",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> TextGenerationResult:
|
||||
...
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
image_filename: str = "image.png",
|
||||
resolution: str = "1K",
|
||||
aspect_ratio: str = "1:1",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> ImageGenerationResult:
|
||||
...
|
||||
|
||||
|
||||
def validate_model_config(model: ResolvedModel) -> None:
|
||||
errors = []
|
||||
if not model.url.strip():
|
||||
errors.append("missing url")
|
||||
if not model.model.strip():
|
||||
errors.append("missing model")
|
||||
if not model.api_key.strip():
|
||||
errors.append("missing api_key")
|
||||
if model.timeout_seconds < 0:
|
||||
errors.append("timeout_seconds must be >= 0")
|
||||
if model.connect_timeout_seconds <= 0:
|
||||
errors.append("connect_timeout_seconds must be > 0")
|
||||
if errors:
|
||||
raise AiProviderConfigError("; ".join(errors))
|
||||
|
||||
|
||||
def _to_int(value: Any, default: int) -> int:
|
||||
if value is None or value == "":
|
||||
return default
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
@@ -0,0 +1,376 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
import requests
|
||||
|
||||
from .base import (
|
||||
AiCapabilityError,
|
||||
AiResponseParseError,
|
||||
ImageGenerationResult,
|
||||
ResolvedModel,
|
||||
TextGenerationResult,
|
||||
validate_model_config,
|
||||
)
|
||||
from .utils import (
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
extract_image_from_response,
|
||||
extract_text_from_response,
|
||||
extract_titles_from_response,
|
||||
image_bytes_to_data_url,
|
||||
normalize_api_url,
|
||||
request_timeout,
|
||||
resolution_to_size,
|
||||
split_data_url,
|
||||
)
|
||||
|
||||
|
||||
class BaseHttpProvider:
|
||||
def __init__(self, session: requests.Session | None = None):
|
||||
self.session = session or requests.Session()
|
||||
if hasattr(self.session, "trust_env"):
|
||||
self.session.trust_env = False
|
||||
|
||||
def _headers(self, model: ResolvedModel, *, json: bool = False) -> dict[str, str]:
|
||||
headers = {"Authorization": f"Bearer {model.api_key}"}
|
||||
if json:
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
def _timeout(self, model: ResolvedModel, resolution: str) -> tuple[int, int]:
|
||||
return request_timeout(
|
||||
model.connect_timeout_seconds,
|
||||
model.timeout_seconds,
|
||||
resolution,
|
||||
)
|
||||
|
||||
def _read_timeout(self, model: ResolvedModel, resolution: str) -> int:
|
||||
return self._timeout(model, resolution)[1]
|
||||
|
||||
|
||||
class ChatCompletionsProvider(BaseHttpProvider):
|
||||
def capabilities(self) -> set[str]:
|
||||
return {"text", "image", "vision"}
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
resolution: str = "1K",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> TextGenerationResult:
|
||||
validate_model_config(model)
|
||||
url = normalize_api_url(model.url, API_CHAT)
|
||||
payload = build_chat_text_payload(
|
||||
model,
|
||||
prompt,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
parameters=parameters,
|
||||
)
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model, json=True),
|
||||
json=payload,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
titles = extract_titles_from_response(raw)
|
||||
text = extract_text_from_response(raw)
|
||||
if not text:
|
||||
raise AiResponseParseError("AI response did not contain text")
|
||||
return TextGenerationResult(text=text, titles=titles, model_used=model.model, raw=raw)
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
image_filename: str = "image.png",
|
||||
resolution: str = "1K",
|
||||
aspect_ratio: str = "1:1",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> ImageGenerationResult:
|
||||
validate_model_config(model)
|
||||
url = normalize_api_url(model.url, API_CHAT)
|
||||
payload = build_chat_image_payload(
|
||||
model,
|
||||
prompt,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
parameters=parameters,
|
||||
)
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model, json=True),
|
||||
json=payload,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
raw,
|
||||
session=self.session,
|
||||
timeout=self._read_timeout(model, resolution),
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiResponseParseError("AI response did not contain an image")
|
||||
return ImageGenerationResult(image=image_bytes, model_used=model.model, raw=raw)
|
||||
|
||||
|
||||
class GeminiProvider(ChatCompletionsProvider):
|
||||
def generate_text(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
resolution: str = "1K",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> TextGenerationResult:
|
||||
validate_model_config(model)
|
||||
url = normalize_api_url(model.url, API_GEMINI).replace("{model}", model.model)
|
||||
payload = build_gemini_payload(
|
||||
model,
|
||||
prompt,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
response_modalities=["TEXT"],
|
||||
parameters=parameters,
|
||||
)
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model, json=True),
|
||||
json=payload,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
titles = extract_titles_from_response(raw)
|
||||
text = extract_text_from_response(raw)
|
||||
if not text:
|
||||
raise AiResponseParseError("AI response did not contain text")
|
||||
return TextGenerationResult(text=text, titles=titles, model_used=model.model, raw=raw)
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
image_filename: str = "image.png",
|
||||
resolution: str = "1K",
|
||||
aspect_ratio: str = "1:1",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> ImageGenerationResult:
|
||||
validate_model_config(model)
|
||||
url = normalize_api_url(model.url, API_GEMINI).replace("{model}", model.model)
|
||||
payload = build_gemini_payload(
|
||||
model,
|
||||
prompt,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
response_modalities=["TEXT", "IMAGE"],
|
||||
parameters=parameters,
|
||||
)
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model, json=True),
|
||||
json=payload,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
raw,
|
||||
session=self.session,
|
||||
timeout=self._read_timeout(model, resolution),
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiResponseParseError("AI response did not contain an image")
|
||||
return ImageGenerationResult(image=image_bytes, model_used=model.model, raw=raw)
|
||||
|
||||
|
||||
class ImagesGenerationProvider(BaseHttpProvider):
|
||||
def capabilities(self) -> set[str]:
|
||||
return {"image"}
|
||||
|
||||
def generate_text(self, *args: Any, **kwargs: Any) -> TextGenerationResult:
|
||||
raise AiCapabilityError("images generation provider cannot generate text")
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
image_filename: str = "image.png",
|
||||
resolution: str = "1K",
|
||||
aspect_ratio: str = "1:1",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> ImageGenerationResult:
|
||||
validate_model_config(model)
|
||||
url = normalize_api_url(model.url, API_IMAGES)
|
||||
payload: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"resolution": resolution,
|
||||
"n": 1,
|
||||
}
|
||||
if image is not None:
|
||||
payload["image_urls"] = [image_bytes_to_data_url(image, image_mime_type)]
|
||||
apply_extra_body(payload, model, parameters)
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model, json=True),
|
||||
json=payload,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
raw,
|
||||
session=self.session,
|
||||
timeout=self._read_timeout(model, resolution),
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiResponseParseError("AI response did not contain an image")
|
||||
return ImageGenerationResult(image=image_bytes, model_used=model.model, raw=raw)
|
||||
|
||||
|
||||
class ImagesEditsProvider(BaseHttpProvider):
|
||||
def capabilities(self) -> set[str]:
|
||||
return {"image", "vision"}
|
||||
|
||||
def generate_text(self, *args: Any, **kwargs: Any) -> TextGenerationResult:
|
||||
raise AiCapabilityError("images edits provider cannot generate text")
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
prompt: str,
|
||||
model: ResolvedModel,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
image_filename: str = "image.png",
|
||||
resolution: str = "1K",
|
||||
aspect_ratio: str = "1:1",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> ImageGenerationResult:
|
||||
validate_model_config(model)
|
||||
if image is None:
|
||||
raise AiCapabilityError("images edits provider requires an input image")
|
||||
|
||||
url = normalize_api_url(model.url, API_IMAGES_EDITS)
|
||||
data: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"prompt": prompt,
|
||||
"n": "1",
|
||||
"size": resolution_to_size(resolution),
|
||||
}
|
||||
apply_extra_body(data, model, parameters)
|
||||
files = {"image": (image_filename, image, image_mime_type)}
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=self._headers(model),
|
||||
data=data,
|
||||
files=files,
|
||||
timeout=self._timeout(model, resolution),
|
||||
)
|
||||
response.raise_for_status()
|
||||
raw = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
raw,
|
||||
session=self.session,
|
||||
timeout=self._read_timeout(model, resolution),
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiResponseParseError("AI response did not contain an image")
|
||||
return ImageGenerationResult(image=image_bytes, model_used=model.model, raw=raw)
|
||||
|
||||
|
||||
def build_chat_text_payload(
|
||||
model: ResolvedModel,
|
||||
prompt: str,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
if image is not None:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_bytes_to_data_url(image, image_mime_type)},
|
||||
}
|
||||
)
|
||||
payload: dict[str, Any] = {
|
||||
"model": model.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
}
|
||||
apply_extra_body(payload, model, parameters)
|
||||
return payload
|
||||
|
||||
|
||||
def build_chat_image_payload(
|
||||
model: ResolvedModel,
|
||||
prompt: str,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return build_chat_text_payload(
|
||||
model,
|
||||
prompt,
|
||||
image=image,
|
||||
image_mime_type=image_mime_type,
|
||||
parameters=parameters,
|
||||
)
|
||||
|
||||
|
||||
def build_gemini_payload(
|
||||
model: ResolvedModel,
|
||||
prompt: str,
|
||||
*,
|
||||
image: bytes | None = None,
|
||||
image_mime_type: str = "image/png",
|
||||
response_modalities: list[str],
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
parts: list[dict[str, Any]] = [{"text": prompt}]
|
||||
if image is not None:
|
||||
data_url = image_bytes_to_data_url(image, image_mime_type)
|
||||
mime_type, data = split_data_url(data_url)
|
||||
parts.append({"inlineData": {"mimeType": mime_type, "data": data}})
|
||||
payload: dict[str, Any] = {
|
||||
"contents": [{"parts": parts}],
|
||||
"generationConfig": {"responseModalities": response_modalities},
|
||||
}
|
||||
apply_extra_body(payload, model, parameters)
|
||||
return payload
|
||||
|
||||
|
||||
def apply_extra_body(
|
||||
payload: dict[str, Any],
|
||||
model: ResolvedModel,
|
||||
parameters: Mapping[str, Any] | None = None,
|
||||
) -> None:
|
||||
payload.update(model.extra_body)
|
||||
if parameters:
|
||||
payload.update(dict(parameters))
|
||||
@@ -0,0 +1,56 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .base import AiProviderConfigError, Provider
|
||||
from .openai_compatible import (
|
||||
ChatCompletionsProvider,
|
||||
GeminiProvider,
|
||||
ImagesEditsProvider,
|
||||
ImagesGenerationProvider,
|
||||
)
|
||||
from .utils import (
|
||||
API_AUTO,
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
detect_api_type,
|
||||
)
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
def __init__(self) -> None:
|
||||
self._providers: dict[str, Provider] = {}
|
||||
|
||||
def register(self, api_type: str, provider: Provider) -> None:
|
||||
self._providers[api_type] = provider
|
||||
|
||||
def resolve_api_type(self, api_type: str, url: str = "") -> str:
|
||||
return detect_api_type(url, api_type) if api_type == API_AUTO else api_type
|
||||
|
||||
def get(self, api_type: str, url: str = "") -> Provider:
|
||||
resolved_api_type = self.resolve_api_type(api_type, url)
|
||||
try:
|
||||
return self._providers[resolved_api_type]
|
||||
except KeyError as exc:
|
||||
raise AiProviderConfigError(
|
||||
f"no provider registered for api_type: {resolved_api_type}"
|
||||
) from exc
|
||||
|
||||
|
||||
default_registry = ProviderRegistry()
|
||||
default_registry.register(API_CHAT, ChatCompletionsProvider())
|
||||
default_registry.register(API_GEMINI, GeminiProvider())
|
||||
default_registry.register(API_IMAGES, ImagesGenerationProvider())
|
||||
default_registry.register(API_IMAGES_EDITS, ImagesEditsProvider())
|
||||
|
||||
|
||||
def register_provider(api_type: str, provider: Provider) -> None:
|
||||
default_registry.register(api_type, provider)
|
||||
|
||||
|
||||
def resolve_api_type(api_type: str, url: str = "") -> str:
|
||||
return default_registry.resolve_api_type(api_type, url)
|
||||
|
||||
|
||||
def get_provider(api_type: str, url: str = "") -> Provider:
|
||||
return default_registry.get(api_type, url)
|
||||
@@ -0,0 +1,262 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
import requests
|
||||
|
||||
from .base import AiProviderConfigError, AiResponseParseError
|
||||
|
||||
API_AUTO = "auto"
|
||||
API_CHAT = "chat"
|
||||
API_GEMINI = "gemini"
|
||||
API_IMAGES = "images"
|
||||
API_IMAGES_EDITS = "images_edits"
|
||||
|
||||
SUPPORTED_API_TYPES = {
|
||||
API_AUTO,
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
}
|
||||
|
||||
RESOLUTION_TIMEOUTS = {"512": 180, "1K": 240, "2K": 360, "4K": 600}
|
||||
BASE64_KEYS = {"image_base64", "base64", "b64_json", "data"}
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
|
||||
|
||||
TITLE_LEAD = re.compile(r"^\s*(?:\d+\s*[\.\)、:：]|[-*•])\s*")
|
||||
TITLE_CIRCLED = "①②③④⑤⑥⑦⑧⑨⑩"
|
||||
TITLE_QUOTES = "\"'「」『』“”‘’"
|
||||
TITLE_SPLIT = re.compile(r"[,，\r\n]+")
|
||||
|
||||
|
||||
def resolution_timeout(resolution: str, default: int = 240) -> int:
|
||||
return RESOLUTION_TIMEOUTS.get(str(resolution).strip().upper(), default)
|
||||
|
||||
|
||||
def request_timeout(connect_timeout: int, read_timeout: int, resolution: str) -> tuple[int, int]:
|
||||
resolved_read_timeout = read_timeout if read_timeout > 0 else resolution_timeout(resolution)
|
||||
return connect_timeout, resolved_read_timeout
|
||||
|
||||
|
||||
def detect_api_type(url: str, api_type: str = API_AUTO) -> str:
|
||||
if api_type and api_type != API_AUTO:
|
||||
if api_type not in SUPPORTED_API_TYPES:
|
||||
raise AiProviderConfigError(f"unsupported api_type: {api_type}")
|
||||
return api_type
|
||||
|
||||
lower_url = str(url).lower()
|
||||
if "generatecontent" in lower_url or "gemini" in lower_url:
|
||||
return API_GEMINI
|
||||
if "/images/edits" in lower_url:
|
||||
return API_IMAGES_EDITS
|
||||
if "/images" in lower_url:
|
||||
return API_IMAGES
|
||||
return API_CHAT
|
||||
|
||||
|
||||
def normalize_api_url(url: str, api_type: str) -> str:
|
||||
raw = str(url).strip()
|
||||
if not raw:
|
||||
return raw
|
||||
|
||||
endpoint = {
|
||||
API_CHAT: "chat/completions",
|
||||
API_IMAGES: "images/generations",
|
||||
API_IMAGES_EDITS: "images/edits",
|
||||
}.get(api_type)
|
||||
|
||||
lower_path = urlparse(raw).path.lower().rstrip("/")
|
||||
if api_type == API_GEMINI:
|
||||
if "generatecontent" in lower_path:
|
||||
return raw
|
||||
return join_url(raw, "v1beta/models/{model}:generateContent")
|
||||
|
||||
if endpoint is None:
|
||||
return raw
|
||||
if lower_path.endswith("/" + endpoint):
|
||||
return raw
|
||||
if lower_path.endswith("/v1"):
|
||||
return join_url(raw, endpoint)
|
||||
return join_url(raw, "v1/" + endpoint)
|
||||
|
||||
|
||||
def join_url(base_url: str, suffix: str) -> str:
|
||||
base = str(base_url).rstrip("/") + "/"
|
||||
return urljoin(base, suffix)
|
||||
|
||||
|
||||
def image_bytes_to_data_url(image: bytes, mime_type: str = "image/png") -> str:
|
||||
encoded = base64.b64encode(image).decode("ascii")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def split_data_url(data_url: str) -> tuple[str, str]:
|
||||
prefix, encoded = data_url.split(",", 1)
|
||||
mime_type = prefix[len("data:") :].split(";", 1)[0]
|
||||
return mime_type, encoded
|
||||
|
||||
|
||||
def decode_image_data_url(data_url: str) -> bytes:
|
||||
marker = ";base64,"
|
||||
if marker not in data_url:
|
||||
raise AiResponseParseError("unsupported data URL image format")
|
||||
return base64.b64decode(data_url.split(marker, 1)[1])
|
||||
|
||||
|
||||
def resolution_to_size(resolution: str) -> str:
|
||||
mapping = {
|
||||
"512": "512x512",
|
||||
"512px": "512x512",
|
||||
"1K": "1024x1024",
|
||||
"2K": "2048x2048",
|
||||
"4K": "4096x4096",
|
||||
}
|
||||
return mapping.get(str(resolution), str(resolution))
|
||||
|
||||
|
||||
def extract_image_from_response(
|
||||
data: Any,
|
||||
*,
|
||||
session: requests.Session | None = None,
|
||||
timeout: int = 60,
|
||||
) -> bytes | None:
|
||||
for key, value in walk_json_items(data):
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
text = value.strip()
|
||||
if text.startswith("data:image/"):
|
||||
return decode_image_data_url(text)
|
||||
if key and key.lower() in BASE64_KEYS and looks_like_base64(text):
|
||||
try:
|
||||
return base64.b64decode(text)
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
pass
|
||||
|
||||
for _key, value in walk_json_items(data):
|
||||
if isinstance(value, str) and is_image_url(value):
|
||||
return download_image_url(value, session=session, timeout=timeout)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def download_image_url(
|
||||
url: str,
|
||||
*,
|
||||
session: requests.Session | None = None,
|
||||
timeout: int = 60,
|
||||
) -> bytes:
|
||||
client = session or requests.Session()
|
||||
if hasattr(client, "trust_env"):
|
||||
client.trust_env = False
|
||||
response = client.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
|
||||
def extract_titles_from_response(data: Any) -> tuple[str, ...]:
|
||||
return tuple(clean_titles(extract_raw_text(data)))
|
||||
|
||||
|
||||
def extract_text_from_response(data: Any) -> str:
|
||||
titles = extract_titles_from_response(data)
|
||||
return titles[0] if titles else ""
|
||||
|
||||
|
||||
def extract_raw_text(data: Any) -> str:
|
||||
if not isinstance(data, dict):
|
||||
return ""
|
||||
|
||||
choices = data.get("choices")
|
||||
if isinstance(choices, list) and choices and isinstance(choices[0], dict):
|
||||
message = choices[0].get("message")
|
||||
if isinstance(message, dict):
|
||||
text = content_to_text(message.get("content"))
|
||||
if text.strip():
|
||||
return text
|
||||
legacy = choices[0].get("text")
|
||||
if isinstance(legacy, str) and legacy.strip():
|
||||
return legacy
|
||||
|
||||
candidates = data.get("candidates")
|
||||
if isinstance(candidates, list) and candidates and isinstance(candidates[0], dict):
|
||||
content = candidates[0].get("content")
|
||||
if isinstance(content, dict) and isinstance(content.get("parts"), list):
|
||||
texts = [
|
||||
part.get("text")
|
||||
for part in content["parts"]
|
||||
if isinstance(part, dict) and isinstance(part.get("text"), str)
|
||||
]
|
||||
joined = "\n".join(text for text in texts if text)
|
||||
if joined.strip():
|
||||
return joined
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def content_to_text(content: Any) -> str:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
texts = []
|
||||
for part in content:
|
||||
if isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
texts.append(part["text"])
|
||||
elif isinstance(part, str):
|
||||
texts.append(part)
|
||||
return "\n".join(texts)
|
||||
return ""
|
||||
|
||||
|
||||
def clean_titles(text: str) -> list[str]:
|
||||
if not text:
|
||||
return []
|
||||
out = []
|
||||
for piece in TITLE_SPLIT.split(str(text)):
|
||||
cleaned = clean_title_line(piece)
|
||||
if cleaned:
|
||||
out.append(cleaned)
|
||||
return out
|
||||
|
||||
|
||||
def clean_title_line(line: str) -> str:
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
return ""
|
||||
stripped = TITLE_LEAD.sub("", stripped)
|
||||
stripped = stripped.lstrip(TITLE_CIRCLED).strip()
|
||||
stripped = stripped.strip(TITLE_QUOTES).strip()
|
||||
return stripped
|
||||
|
||||
|
||||
def walk_json_items(value: Any, key: str | None = None) -> Iterable[tuple[str | None, Any]]:
|
||||
yield key, value
|
||||
if isinstance(value, dict):
|
||||
for child_key, child_value in value.items():
|
||||
yield from walk_json_items(child_value, str(child_key))
|
||||
elif isinstance(value, list):
|
||||
for child_value in value:
|
||||
yield from walk_json_items(child_value, key)
|
||||
|
||||
|
||||
def looks_like_base64(text: str) -> bool:
|
||||
if len(text) < 8:
|
||||
return False
|
||||
try:
|
||||
base64.b64decode(text, validate=True)
|
||||
return True
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
return False
|
||||
|
||||
|
||||
def is_image_url(text: str) -> bool:
|
||||
parsed = urlparse(text.strip())
|
||||
if parsed.scheme.lower() not in ("http", "https"):
|
||||
return False
|
||||
suffix = Path(parsed.path).suffix.lower()
|
||||
return suffix in IMAGE_EXTENSIONS
|
||||
+176
-2
@@ -1,3 +1,177 @@
|
||||
from django.test import TestCase
|
||||
import base64
|
||||
|
||||
# Create your tests here.
|
||||
from django.test import SimpleTestCase
|
||||
|
||||
from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type
|
||||
from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, payload, content=b""):
|
||||
self.payload = payload
|
||||
self.content = content
|
||||
|
||||
def json(self):
|
||||
return self.payload
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, *responses):
|
||||
self.responses = list(responses)
|
||||
self.posts = []
|
||||
self.gets = []
|
||||
self.trust_env = True
|
||||
|
||||
def post(self, url, **kwargs):
|
||||
self.posts.append({"url": url, **kwargs})
|
||||
return self.responses.pop(0)
|
||||
|
||||
def get(self, url, **kwargs):
|
||||
self.gets.append({"url": url, **kwargs})
|
||||
return self.responses.pop(0)
|
||||
|
||||
|
||||
class ProviderRegistryTests(SimpleTestCase):
|
||||
def test_auto_api_type_resolves_chat_provider_from_url(self):
|
||||
url = "https://api.vectorengine.ai/v1/chat/completions"
|
||||
|
||||
self.assertEqual(resolve_api_type("auto", url), "chat")
|
||||
self.assertIsInstance(get_provider("auto", url), ChatCompletionsProvider)
|
||||
|
||||
|
||||
class ChatCompletionsProviderTests(SimpleTestCase):
|
||||
def test_generate_text_builds_chat_payload_and_cleans_titles(self):
|
||||
session = FakeSession(
|
||||
FakeResponse(
|
||||
{
|
||||
"choices": [
|
||||
{"message": {"content": "1. Red Dress\n2. Blue Coat"}}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
provider = ChatCompletionsProvider(session=session)
|
||||
model = ResolvedModel(
|
||||
name="GPT-5.5 text",
|
||||
url="https://api.vectorengine.ai/v1",
|
||||
model="gpt-5.5",
|
||||
api_key="test-key",
|
||||
api_type="chat",
|
||||
)
|
||||
|
||||
result = provider.generate_text(
|
||||
"Generate titles",
|
||||
model,
|
||||
parameters={"temperature": 0.2},
|
||||
)
|
||||
|
||||
self.assertEqual(result.text, "Red Dress")
|
||||
self.assertEqual(result.titles, ("Red Dress", "Blue Coat"))
|
||||
request = session.posts[0]
|
||||
self.assertEqual(
|
||||
request["url"],
|
||||
"https://api.vectorengine.ai/v1/chat/completions",
|
||||
)
|
||||
self.assertEqual(request["headers"]["Authorization"], "Bearer test-key")
|
||||
self.assertEqual(request["json"]["model"], "gpt-5.5")
|
||||
self.assertFalse(request["json"]["stream"])
|
||||
self.assertEqual(request["json"]["temperature"], 0.2)
|
||||
|
||||
def test_generate_image_parses_chat_multimodal_data_url(self):
|
||||
generated = b"generated-image"
|
||||
encoded = base64.b64encode(generated).decode("ascii")
|
||||
session = FakeSession(
|
||||
FakeResponse(
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": [
|
||||
{"type": "text", "text": "done"},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{encoded}"
|
||||
},
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
provider = ChatCompletionsProvider(session=session)
|
||||
model = ResolvedModel(
|
||||
name="Nano Banana 2",
|
||||
url="https://api.vectorengine.ai/v1/chat/completions",
|
||||
model="gemini-3.1-flash-image-preview",
|
||||
api_key="test-key",
|
||||
api_type="auto",
|
||||
)
|
||||
|
||||
result = provider.generate_image(
|
||||
"Generate product image",
|
||||
model,
|
||||
image=b"input-image",
|
||||
image_mime_type="image/jpeg",
|
||||
)
|
||||
|
||||
self.assertEqual(result.image, generated)
|
||||
content = session.posts[0]["json"]["messages"][0]["content"]
|
||||
self.assertEqual(content[0], {"type": "text", "text": "Generate product image"})
|
||||
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,"))
|
||||
|
||||
|
||||
class ImagesEditsProviderTests(SimpleTestCase):
|
||||
def test_generate_image_builds_multipart_request_and_parses_base64(self):
|
||||
generated = b"edited-image"
|
||||
encoded = base64.b64encode(generated).decode("ascii")
|
||||
session = FakeSession(FakeResponse({"data": [{"b64_json": encoded}]}))
|
||||
provider = ImagesEditsProvider(session=session)
|
||||
model = ResolvedModel(
|
||||
name="GPT Image 2",
|
||||
url="https://api.vectorengine.ai/v1/images/edits",
|
||||
model="gpt-image-2",
|
||||
api_key="test-key",
|
||||
api_type="images_edits",
|
||||
)
|
||||
|
||||
result = provider.generate_image(
|
||||
"Replace background",
|
||||
model,
|
||||
image=b"source-image",
|
||||
image_mime_type="image/png",
|
||||
image_filename="source.png",
|
||||
resolution="1K",
|
||||
)
|
||||
|
||||
self.assertEqual(result.image, generated)
|
||||
request = session.posts[0]
|
||||
self.assertEqual(request["url"], "https://api.vectorengine.ai/v1/images/edits")
|
||||
self.assertEqual(request["headers"]["Authorization"], "Bearer test-key")
|
||||
self.assertEqual(request["data"]["model"], "gpt-image-2")
|
||||
self.assertEqual(request["data"]["size"], "1024x1024")
|
||||
self.assertEqual(
|
||||
request["files"]["image"],
|
||||
("source.png", b"source-image", "image/png"),
|
||||
)
|
||||
|
||||
def test_images_edits_requires_input_image(self):
|
||||
provider = ImagesEditsProvider(session=FakeSession())
|
||||
model = ResolvedModel(
|
||||
name="GPT Image 2",
|
||||
url="https://api.vectorengine.ai/v1/images/edits",
|
||||
model="gpt-image-2",
|
||||
api_key="test-key",
|
||||
api_type="images_edits",
|
||||
)
|
||||
|
||||
with self.assertRaises(AiCapabilityError):
|
||||
provider.generate_image("Replace background", model)
|
||||
|
||||
with self.assertRaises(AiCapabilityError):
|
||||
provider.generate_text("Generate title", model)
|
||||
|
||||
Reference in New Issue
Block a user