python-basekit 0.0.11__py3-none-any.whl

basekit/ai/clients/anthropic.py

@@ -0,0 +1,192 @@
+from collections.abc import Callable
+from typing import ClassVar, override
+
+import logfire
+from httpx import Response
+from pydantic import BaseModel
+
+from basekit.ai.schema import AIClient, Tool, Usage, ValidationError
+from basekit.ai.utils import transform_schema
+from basekit.utils.misc import recursive_merge
+
+
+class AnthropicClient(AIClient):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ("max_tokens",)
+
+    api_key: str
+    base_url: str = "https://api.anthropic.com"
+    max_tokens: int = 16000
+
+    def _headers(self) -> dict:
+        return {
+            "X-Api-Key": self.api_key,
+            "Anthropic-Version": "2023-06-01",
+            "Content-Type": "application/json",
+        }
+
+    def _message_url(self) -> str:
+        return f"{self.base_url}/v1/messages"
+
+    def _message_params(self, system: str, prompt: str, /) -> dict:
+        return {
+            "model": self.model.request_name,
+            "system": system,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": self.max_tokens,
+        }
+
+    def _structured_message_params(self, schema: dict, /) -> dict:
+        return {"output_config": {"format": {"type": "json_schema", "schema": schema}}}
+
+    def _tool_use_params(self, tool: dict, /) -> dict:
+        return {"tools": [tool], "tool_choice": {"type": "any"}}
+
+    def _get_block(
+        self, data: dict, stop_reason: str, filter_: Callable[[dict], bool], /
+    ) -> dict:
+        if data["stop_reason"] != stop_reason:
+            raise ValueError(f"STOP REASON：{data['stop_reason']}")
+
+        blocks: list[dict] = [c for c in data["content"] if filter_(c)]
+        if len(blocks) != 1:
+            raise ValueError(f"LEN(BLOCKS)：{len(blocks)}")
+
+        return blocks[0]
+
+    def _get_usage(self, data: dict, /) -> Usage:
+        usage: dict = data["usage"]
+        cache_creation_tokens: int = usage.get("cache_creation_input_tokens", 0)
+        cache_read_tokens: int = usage.get("cache_read_input_tokens", 0)
+        return Usage(
+            input_tokens=usage["input_tokens"]
+            + cache_creation_tokens
+            + cache_read_tokens,
+            cache_creation_tokens=cache_creation_tokens,
+            cache_read_tokens=cache_read_tokens,
+            output_tokens=usage["output_tokens"],
+        )
+
+    @override
+    async def _create_message_once(
+        self, system: str, prompt: str, extra: dict | None
+    ) -> str:
+        with logfire.span("anthropic client | create message") as span:
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(
+                    data, "end_turn", lambda x: x["type"] == "text"
+                )
+                message: str = block["text"]
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=message)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return message
+            except Exception as exc:
+                raise ValidationError("生成消息失败") from exc
+
+    @override
+    async def _create_structured_message_once[T: BaseModel](
+        self, system: str, prompt: str, schema: type[T], extra: dict | None
+    ) -> T:
+        schema_name: str = schema.__name__
+        with logfire.span(
+            f"anthropic client | create structured message | {schema_name}"
+        ) as span:
+            schema_data: dict = transform_schema(schema.model_json_schema())
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="schema", value=schema_data)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self._structured_message_params(schema_data),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(
+                    data, "end_turn", lambda x: x["type"] == "text"
+                )
+                result: T = schema.model_validate_json(json_data=block["text"])
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=result)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return result
+            except Exception as exc:
+                raise ValidationError("生成结构化消息失败") from exc
+
+    @override
+    async def _create_tool_use_once[T: BaseModel](
+        self, system: str, prompt: str, tool: Tool[T], extra: dict | None
+    ) -> T:
+        tool_name: str = tool.name
+        with logfire.span(f"anthropic client | create tool use | {tool_name}") as span:
+            tool_: dict = {
+                "name": tool.name,
+                "description": tool.description,
+                "input_schema": transform_schema(tool.input_schema.model_json_schema()),
+                "strict": True,
+            }
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="tool", value=tool_)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self._tool_use_params(tool_),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(
+                    data, "tool_use", lambda x: x["type"] == "tool_use"
+                )
+                result: T = tool.input_schema.model_validate(obj=block["input"])
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=result)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return result
+            except Exception as exc:
+                raise ValidationError("生成工具调用失败") from exc

basekit/ai/clients/dashscope.py

@@ -0,0 +1,133 @@
+from collections.abc import Callable
+from typing import ClassVar, override
+
+import logfire
+from httpx import Response
+from pydantic import BaseModel
+
+from basekit.ai.schema import (
+    AIClient,
+    ImageRatio,
+    ImageSize,
+    NotSupportedError,
+    ValidationError,
+)
+from basekit.utils.misc import recursive_merge
+
+
+class ImageUsage(BaseModel):
+    width: int
+    height: int
+
+    def format(self) -> str:
+        return f"{self.width}x{self.height}"
+
+
+class DashScopeClient(AIClient):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ()
+
+    api_key: str
+    base_url: str = "https://dashscope.aliyuncs.com/api/v1"
+
+    def _headers(self) -> dict:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def _image_url(self) -> str:
+        return f"{self.base_url}/services/aigc/multimodal-generation/generation"
+
+    def _image_resolution(self, ratio: ImageRatio, size: ImageSize, /) -> str:
+        base_resolution_map: dict[ImageRatio, tuple[int, int]] = {
+            "9:16": (768, 1344),
+            "2:3": (832, 1248),
+            "3:4": (864, 1184),
+            "4:5": (896, 1152),
+            "1:1": (1024, 1024),
+            "5:4": (1152, 896),
+            "4:3": (1184, 864),
+            "3:2": (1248, 832),
+            "16:9": (1344, 768),
+        }
+        resolution_ratio_map: dict[ImageSize, int] = {"1K": 1, "2K": 2, "4K": 4}
+
+        width: int = base_resolution_map[ratio][0] * resolution_ratio_map[size]
+        height: int = base_resolution_map[ratio][1] * resolution_ratio_map[size]
+        resolution: str = f"{width}*{height}"
+
+        match self.model.name:
+            case "qwen-image-2.0" | "qwen-image-2.0-pro":
+                if size not in {"1K", "2K"}:
+                    raise NotSupportedError(
+                        f"模型 {self.model.name} 仅支持 1K 和 2K 分辨率"
+                    )
+                return resolution
+            case _:
+                raise NotSupportedError(f"模型 {self.model.name} 不合法")
+
+    def _image_params(self, prompt: str, resolution: str, /) -> dict:
+        return {
+            "model": self.model.request_name,
+            "input": {"messages": [{"role": "user", "content": [{"text": prompt}]}]},
+            "parameters": {"size": resolution},
+        }
+
+    def _get_block(self, data: dict, filter_: Callable[[dict], bool], /) -> dict:
+        choices: list = data["output"]["choices"]
+        if len(choices) != 1:
+            raise ValueError(f"LEN(CHOICES)：{len(choices)}")
+
+        choice: dict = choices[0]
+        if choice["finish_reason"] != "stop":
+            raise ValueError(f"FINISH REASON：{choice['finish_reason']}")
+
+        blocks: list[dict] = [c for c in choice["message"]["content"] if filter_(c)]
+        if not blocks:
+            raise ValueError(f"LEN(BLOCKS)：{len(blocks)}")
+
+        return blocks[0]
+
+    def _get_usage(self, data: dict, /) -> ImageUsage:
+        usage: dict = data["usage"]
+        return ImageUsage(width=usage["width"], height=usage["height"])
+
+    @override
+    async def _create_image_once(
+        self, prompt: str, ratio: ImageRatio, size: ImageSize, extra: dict | None
+    ) -> bytes:
+        with logfire.span("dashscope client | create image") as span:
+            resolution: str = self._image_resolution(ratio, size)
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="ratio", value=ratio)
+            span.set_attribute(key="size", value=size)
+            span.set_attribute(key="extra", value=extra)
+            span.set_attribute(key="resolution", value=resolution)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._image_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._image_params(prompt, resolution),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda c: "image" in c)
+                image_url: str = block["image"]
+                image: bytes = (await self._http_client.get(image_url)).content
+                usage: ImageUsage = self._get_usage(data)
+
+                span.set_attribute(key="image_url", value=image_url)
+                span.set_attribute(key="filesize", value=len(image))
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} pixels"
+                return image
+            except Exception as exc:
+                raise ValidationError("生成图像失败") from exc

basekit/ai/clients/gemini.py

@@ -0,0 +1,277 @@
+import base64
+from collections.abc import Callable
+from typing import ClassVar, override
+
+import logfire
+from httpx import Response
+from pydantic import BaseModel
+
+from basekit.ai.schema import (
+    AIClient,
+    ImageRatio,
+    ImageSize,
+    NotSupportedError,
+    Tool,
+    Usage,
+    ValidationError,
+)
+from basekit.ai.utils import transform_schema
+from basekit.utils.misc import recursive_merge
+
+
+class GeminiClient(AIClient):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ()
+
+    api_key: str
+    base_url: str = "https://generativelanguage.googleapis.com"
+
+    def _headers(self) -> dict:
+        return {"X-Goog-Api-Key": self.api_key, "Content-Type": "application/json"}
+
+    def _message_url(self) -> str:
+        return (
+            f"{self.base_url}/v1beta/models/{self.model.request_name}:generateContent"
+        )
+
+    def _message_params(self, system: str, prompt: str, /) -> dict:
+        return {
+            "systemInstruction": {"parts": [{"text": system}]},
+            "contents": [{"parts": [{"text": prompt}]}],
+            "generationConfig": {
+                "responseModalities": ["TEXT"],
+                "thinkingConfig": {"includeThoughts": True},
+            },
+        }
+
+    def _structured_message_params(self, schema: dict, /) -> dict:
+        return {
+            "generationConfig": {
+                "responseMimeType": "application/json",
+                "responseSchema": schema,
+            },
+        }
+
+    def _tool_use_params(self, tool: dict, /) -> dict:
+        return {
+            "tools": [{"functionDeclarations": [tool]}],
+            "toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
+        }
+
+    def _image_url(self) -> str:
+        return (
+            f"{self.base_url}/v1beta/models/{self.model.request_name}:generateContent"
+        )
+
+    def _image_config(self, ratio: ImageRatio, size: ImageSize, /) -> dict:
+        match self.model.name:
+            case "gemini-3.1-flash-image-preview":
+                return {"aspectRatio": ratio, "imageSize": size}
+            case "gemini-3-pro-image-preview":
+                return {"aspectRatio": ratio, "imageSize": size}
+            case "gemini-2.5-flash-image":
+                if size != "1K":
+                    raise NotSupportedError(f"模型 {self.model.name} 仅支持 1K 分辨率")
+                return {"aspectRatio": ratio}
+            case _:
+                raise NotSupportedError(f"模型 {self.model.name} 不合法")
+
+    def _image_params(self, prompt: str, config: dict, /) -> dict:
+        return {
+            "contents": [{"parts": [{"text": prompt}]}],
+            "generationConfig": {
+                "responseModalities": ["IMAGE"],
+                "thinkingConfig": {"includeThoughts": True},
+                "imageConfig": config,
+            },
+        }
+
+    def _get_block(self, data: dict, filter_: Callable[[dict], bool], /) -> dict:
+        candidates: list = data["candidates"]
+        if len(candidates) != 1:
+            raise ValueError(f"LEN(CANDIDATES)：{len(candidates)}")
+
+        candidate: dict = candidates[0]
+        if candidate["finishReason"] != "STOP":
+            raise ValueError(f"FINISH REASON：{candidate['finishReason']}")
+
+        blocks: list[dict] = [
+            p
+            for p in candidate["content"]["parts"]
+            if filter_(p) and "thought" not in p
+        ]
+        if len(blocks) != 1:
+            raise ValueError(f"LEN(BLOCKS)：{len(blocks)}")
+
+        return blocks[0]
+
+    def _get_usage(self, data: dict, /) -> Usage:
+        metadata: dict = data["usageMetadata"]
+        return Usage(
+            input_tokens=metadata["promptTokenCount"],
+            cache_creation_tokens=0,
+            cache_read_tokens=metadata.get("cachedContentTokenCount", 0),
+            output_tokens=metadata.get("thoughtsTokenCount", 0)
+            + metadata["candidatesTokenCount"],
+        )
+
+    @override
+    async def _create_message_once(
+        self, system: str, prompt: str, extra: dict | None
+    ) -> str:
+        with logfire.span("gemini client | create message") as span:
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda x: "text" in x)
+                message: str = block["text"]
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=message)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return message
+            except Exception as exc:
+                raise ValidationError("生成消息失败") from exc
+
+    @override
+    async def _create_structured_message_once[T: BaseModel](
+        self, system: str, prompt: str, schema: type[T], extra: dict | None
+    ) -> T:
+        schema_name: str = schema.__name__
+        with logfire.span(
+            f"gemini client | create structured message | {schema_name}"
+        ) as span:
+            schema_data: dict = transform_schema(
+                schema.model_json_schema(), gemini=True
+            )
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="schema", value=schema_data)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self._structured_message_params(schema_data),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda x: "text" in x)
+                result: T = schema.model_validate_json(json_data=block["text"])
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=result)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return result
+            except Exception as exc:
+                raise ValidationError("生成结构化消息失败") from exc
+
+    @override
+    async def _create_tool_use_once[T: BaseModel](
+        self, system: str, prompt: str, tool: Tool[T], extra: dict | None
+    ) -> T:
+        tool_name: str = tool.name
+        with logfire.span(f"gemini client | create tool use | {tool_name}") as span:
+            tool_: dict = {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": transform_schema(
+                    tool.input_schema.model_json_schema(), gemini=True
+                ),
+            }
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="system", value=system)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="tool", value=tool_)
+            span.set_attribute(key="extra", value=extra)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._message_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._message_params(system, prompt),
+                        self._tool_use_params(tool_),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda x: "functionCall" in x)
+                result: T = tool.input_schema.model_validate(
+                    obj=block["functionCall"]["args"]
+                )
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="return", value=result)
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return result
+            except Exception as exc:
+                raise ValidationError("生成工具调用失败") from exc
+
+    @override
+    async def _create_image_once(
+        self, prompt: str, ratio: ImageRatio, size: ImageSize, extra: dict | None
+    ) -> bytes:
+        with logfire.span("gemini client | create image") as span:
+            config: dict = self._image_config(ratio, size)
+
+            span.set_attribute(key="model", value=self.model)
+            span.set_attribute(key="prompt", value=prompt)
+            span.set_attribute(key="ratio", value=ratio)
+            span.set_attribute(key="size", value=size)
+            span.set_attribute(key="extra", value=extra)
+            span.set_attribute(key="config", value=config)
+
+            async with self._limiter:
+                response: Response = await self._http_client.post(
+                    self._image_url(),
+                    headers=self._headers(),
+                    data=recursive_merge(
+                        self._image_params(prompt, config),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda x: "inlineData" in x)
+                image: bytes = base64.b64decode(s=block["inlineData"]["data"])
+                usage: Usage = self._get_usage(data)
+
+                span.set_attribute(key="filesize", value=len(image))
+                span.set_attribute(key="usage", value=usage)
+                span.message = f"{span.message} -> {usage.format()} tokens"
+                return image
+            except Exception as exc:
+                raise ValidationError("生成图像失败") from exc