python-basekit 0.0.11__py3-none-any.whl

basekit/ai/clients/openai.py

@@ -0,0 +1,384 @@
+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 OpenAIClient(AIClient):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ()
+
+    api_key: str
+    base_url: str = "https://api.openai.com/v1"
+
+    def _headers(self, /) -> dict:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def _message_url(self, /) -> str:
+        return f"{self.base_url}/responses"
+
+    def _message_params(self, system: str, prompt: str, /) -> dict:
+        return {
+            "model": self.model.request_name,
+            "instructions": system,
+            "input": [{"role": "user", "content": prompt}],
+        }
+
+    def _structured_message_params(self, name: str, schema: dict, /) -> dict:
+        config: dict = {"name": name, "schema": schema, "strict": True}
+        return {"text": {"format": {"type": "json_schema", **config}}}
+
+    def _tool_use_params(self, tool: dict, /) -> dict:
+        return {"tools": [tool], "tool_choice": "required"}
+
+    def _check_status(self, data: dict, /) -> bool:
+        return "status" not in data or data["status"] == "completed"
+
+    def _get_block(self, data: dict, filter_: Callable[[dict], bool], /) -> dict:
+        if not self._check_status(data):
+            raise ValueError(f"STATUS：{data['status']}")
+
+        blocks: list[dict] = [
+            o for o in data["output"] if filter_(o) and self._check_status(o)
+        ]
+        if len(blocks) != 1:
+            raise ValueError(f"LEN(BLOCKS)：{len(blocks)}")
+
+        return blocks[0]
+
+    def _get_text(self, block: dict, /) -> str:
+        contents: list[dict] = [
+            c for c in block["content"] if c["type"] == "output_text"
+        ]
+        if len(contents) != 1:
+            raise ValueError(f"LEN(CONTENTS)：{len(contents)}")
+
+        return contents[0]["text"]
+
+    def _get_usage(self, data: dict, /) -> Usage:
+        usage: dict = data["usage"]
+        return Usage(
+            input_tokens=usage["input_tokens"],
+            cache_creation_tokens=0,
+            cache_read_tokens=usage["input_tokens_details"]["cached_tokens"],
+            output_tokens=usage["output_tokens"],
+        )
+
+    @override
+    async def _create_message_once(
+        self, system: str, prompt: str, extra: dict | None
+    ) -> str:
+        with logfire.span("openai 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: x["type"] == "message")
+                message: str = self._get_text(block)
+                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"openai 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_name, schema_data),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, lambda x: x["type"] == "message")
+                message: str = self._get_text(block)
+                result: T = schema.model_validate_json(json_data=message)
+                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"openai client | create tool use | {tool_name}") as span:
+            tool_: dict = {
+                "type": "function",
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": 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, lambda x: x["type"] == "function_call"
+                )
+                result: T = tool.input_schema.model_validate_json(
+                    json_data=block["arguments"]
+                )
+                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
+
+
+class OpenAILegacyClient(AIClient):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ()
+
+    api_key: str
+    base_url: str = "https://api.openai.com/v1"
+
+    def _headers(self, /) -> dict:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def _message_url(self, /) -> str:
+        return f"{self.base_url}/chat/completions"
+
+    def _message_params(self, system: str, prompt: str, /) -> dict:
+        return {
+            "model": self.model.request_name,
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": prompt},
+            ],
+        }
+
+    def _structured_message_params(self, name: str, schema: dict, /) -> dict:
+        config: dict = {"name": name, "schema": schema, "strict": True}
+        return {"response_format": {"type": "json_schema", "json_schema": config}}
+
+    def _tool_use_params(self, tool: dict, /) -> dict:
+        return {"tools": [tool], "tool_choice": "required"}
+
+    def _get_block(self, data: dict, finish_reason: str, /) -> dict:
+        choices: list[dict] = data["choices"]
+        if len(choices) != 1:
+            raise ValueError(f"LEN(BLOCKS)：{len(choices)}")
+
+        choice: dict = choices[0]
+        if choice["finish_reason"] != finish_reason:
+            raise ValueError(f"FINISH REASON：{choice['finish_reason']}")
+
+        return choice["message"]
+
+    def _get_arguments(self, block: dict, /) -> str:
+        arguments: list[str] = [t["function"]["arguments"] for t in block["tool_calls"]]
+        if len(arguments) != 1:
+            raise ValueError(f"LEN(ARGUMENTS)：{len(arguments)}")
+
+        return arguments[0]
+
+    def _get_usage(self, data: dict, /) -> Usage:
+        usage: dict = data["usage"]
+        return Usage(
+            input_tokens=usage["prompt_tokens"],
+            cache_creation_tokens=0,
+            cache_read_tokens=usage["prompt_tokens_details"]["cached_tokens"],
+            output_tokens=usage["completion_tokens"],
+        )
+
+    @override
+    async def _create_message_once(
+        self, system: str, prompt: str, extra: dict | None
+    ) -> str:
+        with logfire.span("openai client (legacy) | 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, "stop")
+                message: str = block["content"]
+                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"openai client (legacy) | 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_name, schema_data),
+                        self.model.extra_params,
+                        extra,
+                    ),
+                )
+
+            try:
+                data: dict = response.json()
+                block: dict = self._get_block(data, "stop")
+                result: T = schema.model_validate_json(json_data=block["content"])
+                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"openai client (legacy) | create tool use | {tool_name}"
+        ) as span:
+            tool_: dict = {
+                "type": "function",
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "parameters": 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_calls")
+                arguments: str = self._get_arguments(block)
+                result: T = tool.input_schema.model_validate_json(json_data=arguments)
+                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/schema.py

@@ -0,0 +1,160 @@
+from collections.abc import AsyncIterator, Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import Any, ClassVar, Literal, override
+
+from httpx import HTTPStatusError, RequestError
+from pydantic import BaseModel, ConfigDict, PrivateAttr
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+from tenacity.stop import stop_base
+from tenacity.wait import wait_base
+
+from basekit.cache.schema import Cache
+from basekit.http.clients.httpx import HttpxClient
+from basekit.http.schema import DEFAULT_LIMIT
+from basekit.limiter import AsyncLimiter
+
+type ImageSize = Literal["1K", "2K", "4K"]
+type ImageRatio = Literal[
+    "9:16", "2:3", "3:4", "4:5", "1:1", "5:4", "4:3", "3:2", "16:9"
+]
+
+
+class NotSupportedError(ValueError):
+    pass
+
+
+class ValidationError(ValueError):
+    pass
+
+
+class Usage(BaseModel):
+    input_tokens: int
+    cache_creation_tokens: int
+    cache_read_tokens: int
+    output_tokens: int
+
+    def format(self) -> str:
+        return f"↓{self.input_tokens}(+{self.cache_creation_tokens},→{self.cache_read_tokens})/↑{self.output_tokens}"  # noqa: E501
+
+
+class Tool[T: BaseModel](BaseModel):
+    name: str
+    description: str
+    input_schema: type[T]
+
+
+class Model(BaseModel):
+    cache_key_fields: ClassVar[tuple[str, ...]] = ("name", "extra_params")
+
+    name: str
+    alias: str | None = None
+    concurrency_limit: int | None = None
+    rate_limit: int | None = None
+    extra_params: dict | None = None
+
+    @property
+    def request_name(self) -> str:
+        return self.alias or self.name
+
+
+class AIClient(BaseModel):
+    model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
+    cache_key_fields: ClassVar[tuple[str, ...]] = ("model",)
+
+    model: Model
+    retry_stop: stop_base = stop_after_attempt(max_attempt_number=3)
+    retry_wait: wait_base = wait_exponential_jitter()
+    cache: Cache | None = None
+
+    _http_client: HttpxClient = PrivateAttr()
+    _limiter: AsyncLimiter = PrivateAttr()
+
+    @override
+    def model_post_init(self, context: Any) -> None:
+        self._http_client = HttpxClient(
+            timeout=600.0, limit=self.model.concurrency_limit or DEFAULT_LIMIT
+        )
+        self._limiter = AsyncLimiter(
+            concurrency_limit=self.model.concurrency_limit,
+            rate_limit=self.model.rate_limit,
+            rate_window=60.0,
+        )
+
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        async with self._http_client.lifespan():
+            yield
+
+    def _should_retry(self, exc: BaseException, /) -> bool:
+        if isinstance(exc, RequestError):
+            return True
+        if isinstance(exc, HTTPStatusError):
+            status_code: int = exc.response.status_code
+            return status_code == 429 or status_code >= 500
+        if isinstance(exc, ValidationError):
+            return True
+        return False
+
+    def _wrap[T, **P](
+        self, func: Callable[P, Awaitable[T]], /
+    ) -> Callable[P, Awaitable[T]]:
+        wrapped_func: Callable[P, Awaitable[T]] = retry(
+            retry=retry_if_exception(predicate=self._should_retry),
+            stop=self.retry_stop,
+            wait=self.retry_wait,
+            reraise=True,
+        )(func)
+        return self.cache(wrapped_func) if self.cache else wrapped_func
+
+    async def _create_message_once(
+        self, system: str, prompt: str, extra: dict | None
+    ) -> str:
+        raise NotImplementedError
+
+    async def create_message(
+        self, system: str, prompt: str, extra: dict | None = None
+    ) -> str:
+        return await self._wrap(self._create_message_once)(
+            system=system, prompt=prompt, extra=extra
+        )
+
+    async def _create_structured_message_once[T: BaseModel](
+        self, system: str, prompt: str, schema: type[T], extra: dict | None
+    ) -> T:
+        raise NotImplementedError
+
+    async def create_structured_message[T: BaseModel](
+        self, system: str, prompt: str, schema: type[T], extra: dict | None = None
+    ) -> T:
+        return await self._wrap(self._create_structured_message_once)(
+            system=system, prompt=prompt, schema=schema, extra=extra
+        )
+
+    async def _create_tool_use_once[T: BaseModel](
+        self, system: str, prompt: str, tool: Tool[T], extra: dict | None
+    ) -> T:
+        raise NotImplementedError
+
+    async def create_tool_use[T: BaseModel](
+        self, system: str, prompt: str, tool: Tool[T], extra: dict | None = None
+    ) -> T:
+        return await self._wrap(self._create_tool_use_once)(
+            system=system, prompt=prompt, tool=tool, extra=extra
+        )
+
+    async def _create_image_once(
+        self, prompt: str, ratio: ImageRatio, size: ImageSize, extra: dict | None
+    ) -> bytes:
+        raise NotImplementedError
+
+    async def create_image(
+        self, prompt: str, ratio: ImageRatio, size: ImageSize, extra: dict | None = None
+    ) -> bytes:
+        return await self._wrap(self._create_image_once)(
+            prompt=prompt, ratio=ratio, size=size, extra=extra
+        )

basekit/ai/utils.py

@@ -0,0 +1,63 @@
+from typing import cast
+
+from basekit.ai.schema import NotSupportedError
+from basekit.utils.misc import Any
+
+
+def transform_schema(schema: Any, /, *, gemini: bool = False) -> dict:
+    if not isinstance(schema, dict):
+        raise NotSupportedError(f"仅支持字典类型模式，SCHEMA：{type(schema)}")
+
+    strict_schema: dict = {}
+
+    if "$ref" in schema:
+        strict_schema["$ref"] = schema["$ref"]
+        return strict_schema
+
+    if "$defs" in schema:
+        strict_schema["$defs"] = {
+            key: transform_schema(value, gemini=gemini)
+            for key, value in schema["$defs"].items()
+        }
+
+    if "title" in schema:
+        strict_schema["title"] = schema["title"]
+
+    if "description" in schema:
+        strict_schema["description"] = schema["description"]
+
+    if "enum" in schema:
+        strict_schema["enum"] = schema["enum"]
+
+    if "const" in schema:
+        strict_schema["const"] = schema["const"]
+
+    if "default" in schema:
+        strict_schema["default"] = schema["default"]
+
+    if "anyOf" in schema:
+        strict_schema["anyOf"] = [
+            transform_schema(cast(dict, item), gemini=gemini)
+            for item in schema["anyOf"]
+        ]
+
+    if "type" not in schema:
+        return strict_schema
+
+    strict_schema["type"] = schema["type"]
+
+    match strict_schema["type"]:
+        case "object":
+            strict_schema["properties"] = {
+                key: transform_schema(value, gemini=gemini)
+                for key, value in schema["properties"].items()
+            }
+            strict_schema["required"] = list(strict_schema["properties"].keys())
+            if not gemini:
+                strict_schema["additionalProperties"] = False
+        case "array":
+            strict_schema["items"] = transform_schema(schema["items"], gemini=gemini)
+        case "string" | "number" | "integer" | "boolean" | "null":
+            pass
+
+    return strict_schema

basekit/cache/clients/sqlite.py

@@ -0,0 +1,90 @@
+import hashlib
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import override
+
+from pydantic import PrivateAttr
+from sqlalchemy import (
+    Column,
+    CursorResult,
+    Delete,
+    LargeBinary,
+    MetaData,
+    Row,
+    Select,
+    String,
+    Table,
+    delete,
+    select,
+)
+from sqlalchemy.dialects.sqlite import Insert, insert
+
+from basekit.cache.schema import Cache
+from basekit.database import DatabaseClient
+
+
+class SQLiteCache(Cache):
+    path: Path | None = None
+
+    _client: DatabaseClient = PrivateAttr()
+    _metadata: MetaData = PrivateAttr()
+    _table: Table = PrivateAttr()
+
+    @override
+    def model_post_init(self, context) -> None:
+        self.path = self.path.expanduser().resolve() if self.path else None
+        path: str = self.path.as_posix() if self.path else ":memory:"
+        url: str = f"sqlite+aiosqlite:///{path}"
+        self._client = DatabaseClient(url=url)
+        self._metadata = MetaData()
+        self._table = Table(
+            "cache",
+            self._metadata,
+            Column("sha256", String(length=64), nullable=False, primary_key=True),
+            Column("key", LargeBinary(), nullable=False),
+            Column("value", LargeBinary(), nullable=False),
+        )
+
+    @override
+    @asynccontextmanager
+    async def lifespan(self) -> AsyncIterator[None]:
+        async with self._client.lifespan():
+            await self._client.create_schema(self._metadata)
+            yield
+
+    def _sha256(self, key: bytes, /) -> str:
+        return hashlib.sha256(string=key).hexdigest()
+
+    @override
+    async def set(self, key: bytes, value: bytes, /) -> None:
+        async with self._client.connection() as connection:
+            statement: Insert = (
+                insert(table=self._table)
+                .values(sha256=self._sha256(key), key=key, value=value)
+                .on_conflict_do_update(
+                    index_elements=[self._table.c.sha256],
+                    set_={"key": key, "value": value},
+                )
+            )
+            await connection.execute(statement=statement)
+            await connection.commit()
+
+    @override
+    async def get(self, key: bytes, /) -> bytes | None:
+        async with self._client.connection() as connection:
+            statement: Select = select(self._table.c.value).where(
+                self._table.c.sha256 == self._sha256(key)
+            )
+            result: CursorResult = await connection.execute(statement=statement)
+        row: Row | None = result.one_or_none()
+        return row.value if row else None
+
+    @override
+    async def delete(self, key: bytes, /) -> None:
+        async with self._client.connection() as connection:
+            statement: Delete = delete(table=self._table).where(
+                self._table.c.sha256 == self._sha256(key)
+            )
+            await connection.execute(statement=statement)
+            await connection.commit()