infinity_llm_client 0.1.0__py3-none-any.whl

llm/clients/open_ai/client.py

@@ -0,0 +1,261 @@
+"""OpenAI 兼容 API 的流式客户端实现。
+
+核心职责：
+- SSE 流解析与事件分发
+- tool_calls 跨 chunk 增量聚合
+- 流终止时的收尾处理
+"""
+
+import asyncio
+import json
+import logging
+import types
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+import aiohttp
+
+from ...models import (
+    DoneChunk,
+    FinishChunk,
+    Messages,
+    StreamChunk,
+    TextChunk,
+    ToolCall,
+    ToolCallCompleteChunk,
+    ToolDefinitions,
+    UsageChunk,
+    UsageStats,
+)
+from ..base import LLMClient
+from ..exceptions import (
+    LLMContentFilterError,
+    LLMNetworkError,
+    LLMStreamError,
+)
+from .aggregation import aggregate_tool_call_deltas
+from .config import OpenAIConfig
+from .connection import ConnectionManager
+from .open_ai_models import (
+    Choice,
+    StreamEvent,
+    ToolCallDelta,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _ToolCallAccumulator:
+    """跨 chunk 累积 tool_calls 增量并在完成时产出聚合结果。"""
+
+    __slots__ = ('_deltas',)
+
+    def __init__(self) -> None:
+        self._deltas: List[ToolCallDelta] = []
+
+    def extend(self, deltas: List[ToolCallDelta]) -> None:
+        """追加一批 tool_calls 增量片段。"""
+        self._deltas.extend(deltas)
+
+    def flush(self) -> List[ToolCall]:
+        """一次性聚合并清空已累积的 tool_calls，直接产出 ToolCall 列表。"""
+        if not self._deltas:
+            return []
+        result = aggregate_tool_call_deltas(self._deltas)
+        self._deltas.clear()
+        return result
+
+    def flush_as_chunk(self) -> Optional[ToolCallCompleteChunk]:
+        """将聚合结果包装为 ToolCallCompleteChunk，无待刷新时返回 None。"""
+        aggregated = self.flush()
+        if not aggregated:
+            return None
+        return ToolCallCompleteChunk(
+            tool_calls=aggregated,
+        )
+
+
+class OpenAIClient(LLMClient):
+    """OpenAI 兼容 API 的轻量异步客户端"""
+
+    def __init__(self, config: OpenAIConfig) -> None:
+        self._config = config
+        self._conn = ConnectionManager(
+            api_key=config.api_key,
+            base_url=config.base_url,
+            config=config.connection,
+        )
+
+    @property
+    def model(self) -> str:
+        """当前使用的模型名称。"""
+        return self._config.model
+
+    async def __aenter__(self) -> 'OpenAIClient':
+        await self._conn.ensure_session()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[types.TracebackType],
+    ) -> Optional[bool]:
+        await self._conn.close()
+
+    async def stream_chat(
+        self,
+        messages: Messages,
+        tools: Optional[ToolDefinitions] = None,
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """流式对话，自动聚合 tool_calls 分片并产出完整 AggregatedToolCall。
+
+        数据流::
+
+            HTTP SSE 行  →  :meth:`_parse_sse_stream`   →  StreamEvent
+            StreamEvent  →  :meth:`_handle_sse_event`   →  StreamChunk (yield)
+            流结束后     →  :meth:`_on_stream_end`      →  补刷 + DoneChunk
+
+        连接级重试（网络错误、可重试 HTTP 状态码）由
+        :class:`ConnectionManager` 内部透明处理，调用方无需关心。
+
+        异常沿三层异步生成器链向上传播，最终在调用方的
+        ``async for`` 循环中抛出。
+
+        :param messages: 对话消息列表
+        :param tools: 可选的工具定义列表
+        """
+        payload = self._build_stream_payload(messages, tools)
+        accumulator = _ToolCallAccumulator()
+
+        async with self._conn.request('chat/completions', payload) as resp:
+            async for event in self._parse_sse_stream(resp):
+                async for chunk in self._handle_sse_event(event, accumulator):
+                    yield chunk
+
+            # 流终止收尾
+            async for chunk in self._on_stream_end(accumulator):
+                yield chunk
+
+    # ------------------------------------------------------------------
+    # SSE 解析
+    # ------------------------------------------------------------------
+
+    async def _parse_sse_stream(
+        self,
+        resp: aiohttp.ClientResponse,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """
+        从 HTTP 响应中逐行读取 SSE 并解析为 StreamEvent。
+
+        - 自动跳过注释行与非 data 行
+        - 遇到 EOF 或 data: [DONE] 时正常结束
+        - 网络读取错误抛出 LLMNetworkError
+        - JSON 解析错误抛出 LLMStreamError
+        """
+        while True:
+            try:
+                line_bytes: bytes = await resp.content.readline()
+            except (asyncio.TimeoutError, aiohttp.ClientError) as e:
+                raise LLMNetworkError(f'Stream read error: {e}', original_error=e) from e
+
+            # EOF
+            if line_bytes == b'':
+                return
+
+            line: str = line_bytes.decode('utf-8').strip()
+
+            # 注释行
+            if line.startswith(':'):
+                continue
+
+            # 非 data 行
+            if not line.startswith('data: '):
+                continue
+
+            # [DONE] 终止标记
+            if line == 'data: [DONE]':
+                return
+
+            # 解析 JSON
+            data_str: str = line[5:].strip()
+            try:
+                yield StreamEvent.model_validate_json(data_str)
+            except json.JSONDecodeError as e:
+                raise LLMStreamError(
+                    f'JSON decode error in stream: {e}',
+                    original_error=e,
+                    response_body=data_str,
+                ) from e
+
+    def _build_stream_payload(
+        self,
+        messages: Messages,
+        tools: Optional[ToolDefinitions],
+    ) -> Dict[str, Any]:
+        """构建流式请求的 JSON payload。"""
+        payload: Dict[str, Any] = {
+            'model': self._config.model,
+            'messages': [m.model_dump(mode='json') for m in messages],
+            'stream': True,
+        }
+        if tools:
+            payload['tools'] = [t.model_dump_tool() for t in tools]
+        if self._config.include_usage:
+            payload['stream_options'] = {'include_usage': True}
+        if self._config.response_format:
+            payload['response_format'] = self._config.response_format.model_dump(
+                mode='json',
+                by_alias=True,
+                exclude_none=True,
+            )
+        return payload
+
+    async def _handle_sse_event(
+        self,
+        event: StreamEvent,
+        accumulator: _ToolCallAccumulator,
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """按固定管线处理单个 SSE 事件"""
+        if event.choices:
+            choice: Choice = event.choices[0]
+
+            # Step 1: 文本增量
+            if choice.delta.content:
+                yield TextChunk(text=choice.delta.content)
+
+            # Step 2: 工具调用累积
+            if choice.delta.tool_calls:
+                accumulator.extend(choice.delta.tool_calls)
+
+            # Step 3: 结束理由
+            if choice.finish_reason is not None:
+                if choice.finish_reason == 'content_filter':
+                    raise LLMContentFilterError(
+                        message='Content filtered by safety system',
+                        status_code=400,
+                    )
+                if choice.finish_reason == 'tool_calls':
+                    chunk = accumulator.flush_as_chunk()
+                    if chunk is not None:
+                        yield chunk
+                yield FinishChunk(finish_reason=choice.finish_reason)
+
+        # Step 4: Token 用量
+        if event.usage and self._config.include_usage:
+            yield UsageChunk(
+                usage=UsageStats(
+                    prompt_tokens=event.usage.prompt_tokens,
+                    completion_tokens=event.usage.completion_tokens,
+                    total_tokens=event.usage.total_tokens,
+                ),
+            )
+
+    async def _on_stream_end(
+        self,
+        accumulator: _ToolCallAccumulator,
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """流终止时的收尾：补刷未完成的 tool_calls，产出 DONE 标记。"""
+        chunk = accumulator.flush_as_chunk()
+        if chunk is not None:
+            yield chunk
+        yield DoneChunk()

llm/clients/open_ai/config.py

@@ -0,0 +1,34 @@
+"""OpenAI 兼容 API 配置模型"""
+
+from typing import Optional
+
+from pydantic import Field
+
+from ..base import LLMClient
+from ..config import LLMConfig
+from ..exceptions import LLMConfigError
+from ..provider import register_client
+from .response_format import ResponseFormat
+
+
+class OpenAIConfig(LLMConfig):
+    """OpenAI 兼容 API 配置"""
+
+    provider: str = 'openai'
+    api_key: str = Field(description='API 密钥')
+    base_url: str = Field(default='https://api.openai.com/v1', description='API 基础 URL')
+    response_format: Optional[ResponseFormat] = Field(
+        default=None,
+        description='响应格式约束（JSON 模式 / 结构化输出）',
+    )
+
+
+@register_client('openai')
+def create_openai_client(config: LLMConfig) -> LLMClient:
+    """根据 OpenAIConfig 创建 OpenAIClient 实例。"""
+    if not isinstance(config, OpenAIConfig):
+        raise LLMConfigError(f'openai factory received incompatible config type: {type(config).__name__}')
+
+    from .client import OpenAIClient
+
+    return OpenAIClient(config)

llm/clients/open_ai/connection.py

@@ -0,0 +1,167 @@
+"""OpenAI 兼容 API 的 HTTP 连接管理"""
+
+import asyncio
+import json
+import logging
+import random
+from contextlib import asynccontextmanager
+from typing import Any, AsyncGenerator, Dict, Optional
+
+import aiohttp
+
+from ..config import ConnectionConfig
+from ..exceptions import (
+    LLMContentFilterError,
+    LLMContextLengthError,
+    LLMHTTPError,
+    LLMNetworkError,
+    build_http_error,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionManager:
+    """管理 aiohttp ClientSession 的生命周期、HTTP 请求发送与连接级重试"""
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str,
+        config: Optional[ConnectionConfig] = None,
+    ) -> None:
+        self._api_key: str = api_key
+        self._base_url: str = base_url.rstrip('/')
+        self._config: ConnectionConfig = config or ConnectionConfig()
+
+        self._session: Optional[aiohttp.ClientSession] = None
+        self._session_lock: asyncio.Lock = asyncio.Lock()
+
+    async def ensure_session(self) -> aiohttp.ClientSession:
+        """懒初始化"""
+        async with self._session_lock:
+            if self._session is None or self._session.closed:
+                timeout_config = aiohttp.ClientTimeout(
+                    total=None,
+                    sock_connect=self._config.connect_timeout,
+                    sock_read=self._config.timeout,
+                )
+                self._session = aiohttp.ClientSession(
+                    headers={
+                        'Authorization': f'Bearer {self._api_key}',
+                        'Content-Type': 'application/json',
+                    },
+                    timeout=timeout_config,
+                )
+                logger.debug('Created new aiohttp ClientSession')
+            return self._session
+
+    async def close(self) -> None:
+        """关闭底层 HTTP 会话。"""
+        async with self._session_lock:
+            if self._session and not self._session.closed:
+                await self._session.close()
+                logger.debug('Closed aiohttp ClientSession')
+            self._session = None
+
+    def _is_retryable(self, error: BaseException) -> bool:
+        if isinstance(error, LLMNetworkError):
+            return True
+        if isinstance(error, LLMHTTPError) and error.status_code is not None:
+            return error.status_code in self._config.retryable_status
+        return False
+
+    def _calc_delay(self, attempt: int) -> float:
+        """计算第 attempt 次重试的等待时间（指数退避 + 可选抖动）。"""
+        cfg = self._config
+        delay = min(cfg.base_delay * (2 ** (attempt - 1)), cfg.max_delay)
+        if cfg.jitter:
+            delay = delay * (0.75 + random.random() * 0.5)  # ±25% 抖动
+        return delay
+
+    async def _do_request(self, endpoint: str, payload: Dict[str, Any]) -> aiohttp.ClientResponse:
+        """执行单次 POST 请求并校验 HTTP 状态码。
+
+        Returns:
+            ClientResponse
+
+        Raises:
+            LLMNetworkError: 连接/超时失败
+            LLMHTTPError: 非 200 状态码
+        """
+        try:
+            session: aiohttp.ClientSession = await self.ensure_session()
+            url: str = f'{self._base_url}/{endpoint.lstrip("/")}'
+            resp: aiohttp.ClientResponse = await session.post(url, json=payload)
+        except asyncio.TimeoutError as e:
+            raise LLMNetworkError(f'Request timeout: {e}', original_error=e) from e
+        except aiohttp.ClientError as e:
+            raise LLMNetworkError(f'Connection error: {e}', original_error=e) from e
+
+        if resp.status == 200:
+            return resp
+
+        error_text: str = await resp.text()
+        resp.close()
+
+        try:
+            error_type: str = json.loads(error_text).get('error', {}).get('type', '')
+        except (json.JSONDecodeError, AttributeError):
+            error_type = ''
+
+        match error_type:
+            case 'content_filter':
+                raise LLMContentFilterError(
+                    message=f'Content filtered: {error_text}',
+                    status_code=resp.status,
+                    response_body=error_text,
+                )
+            case 'context_length_exceeded':
+                raise LLMContextLengthError(
+                    message=f'Context length exceeded: {error_text}',
+                    status_code=resp.status,
+                    response_body=error_text,
+                )
+            case _:
+                raise build_http_error(
+                    status_code=resp.status,
+                    message=f'API error {resp.status}: {error_text}',
+                    response_body=error_text,
+                )
+
+    async def _backoff(self, attempt: int, last_error: Optional[Exception]) -> None:
+        delay: float = self._calc_delay(attempt)
+        logger.warning(
+            'Request failed (attempt %d/%d), retrying in %.1fs: %s',
+            attempt,
+            self._config.max_retries,
+            delay,
+            last_error,
+        )
+        if isinstance(last_error, LLMNetworkError):
+            await self.close()
+        await asyncio.sleep(delay)
+
+    @asynccontextmanager
+    async def request(self, endpoint: str, payload: Dict[str, Any]) -> AsyncGenerator[aiohttp.ClientResponse, None]:
+        last_error: Optional[Exception] = None
+        total_attempts = self._config.max_retries + 1
+
+        for attempt in range(1, total_attempts + 1):
+            try:
+                resp = await self._do_request(endpoint, payload)
+            except (LLMNetworkError, LLMHTTPError) as e:
+                if not self._is_retryable(e):
+                    raise
+                last_error = e
+            else:
+                try:
+                    yield resp
+                    return
+                finally:
+                    resp.close()
+
+            await self._backoff(attempt, last_error)
+
+        assert last_error is not None
+        raise last_error

llm/clients/open_ai/open_ai_models.py

@@ -0,0 +1,121 @@
+"""OpenAI 兼容 API 的请求/响应模型定义"""
+
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class ChatCompletionTokenLogprob(BaseModel):
+    """单个 token 的对数概率信息"""
+
+    token: str = Field(description='Token 字符串')
+    bytes: Optional[List[int]] = Field(default=None, description='Token 的字节表示')
+    logprob: float = Field(description='Token 的对数概率')
+    top_logprobs: Optional[List['ChatCompletionTokenLogprob']] = Field(
+        default=None, description='最可能的替代 token 列表'
+    )
+
+
+class LogprobsContent(BaseModel):
+    """消息内容 token 的对数概率列表"""
+
+    content: Optional[List[ChatCompletionTokenLogprob]] = Field(
+        default=None, description='消息内容 token 的对数概率信息'
+    )
+    refusal: Optional[List[ChatCompletionTokenLogprob]] = Field(
+        default=None, description='消息拒绝 token 的对数概率信息'
+    )
+
+
+class ToolCallFunction(BaseModel):
+    """工具调用中的函数信息（流式增量）"""
+
+    name: Optional[str] = Field(default=None, description='函数名称')
+    arguments: str = Field(default='', description='JSON 参数片段（流式拼接）')
+
+
+class ToolCallDelta(BaseModel):
+    """单条工具调用的流式增量片段"""
+
+    index: int = Field(default=0, description='工具调用在列表中的索引位置')
+    id: Optional[str] = Field(default=None, description='工具调用唯一标识符')
+    type: Literal['function'] = Field(default='function', description='调用类型')
+    function: Optional[ToolCallFunction] = Field(default=None, description='函数名称与参数增量')
+
+
+class Delta(BaseModel):
+    """流式返回的 completion 增量"""
+
+    content: Optional[str] = Field(default=None, description='增量的文本内容')
+    role: Optional[Literal['assistant']] = Field(default=None, description='消息角色')
+    function_call: Optional[Dict[Any, Any]] = Field(default=None, description='函数调用信息（已弃用）')
+    refusal: Optional[str] = Field(default=None, description='拒绝消息内容')
+    tool_calls: Optional[List[ToolCallDelta]] = Field(default=None, description='工具调用增量列表')
+
+
+class CompletionTokensDetails(BaseModel):
+    """completion 中 token 的细分统计"""
+
+    accepted_prediction_tokens: Optional[int] = Field(
+        default=None,
+        description='使用预测输出时，预测中出现在 completion 中的 token 数量',
+    )
+    audio_tokens: Optional[int] = Field(default=None, description='模型生成的音频 token 数量')
+    reasoning_tokens: Optional[int] = Field(default=None, description='模型用于推理的 token 数量')
+    rejected_prediction_tokens: Optional[int] = Field(
+        default=None,
+        description='使用预测输出时，预测中未出现在 completion 中的 token 数量',
+    )
+
+
+class PromptTokensDetails(BaseModel):
+    """prompt 中 token 的细分统计"""
+
+    audio_tokens: Optional[int] = Field(default=None, description='prompt 中的音频 token 数量')
+    cached_tokens: Optional[int] = Field(default=None, description='prompt 中的缓存 token 数量')
+
+
+class CompletionUsage(BaseModel):
+    """Token 使用统计"""
+
+    completion_tokens: int = Field(description='生成的 completion 的 token 数量')
+    prompt_tokens: int = Field(description='prompt 的 token 数量')
+    total_tokens: int = Field(description='总计 token 数量')
+    completion_tokens_details: Optional[CompletionTokensDetails] = Field(
+        default=None, description='completion token 的细分信息'
+    )
+    prompt_tokens_details: Optional[PromptTokensDetails] = Field(default=None, description='prompt token 的细分信息')
+
+
+class Choice(BaseModel):
+    """completion 选择项"""
+
+    delta: Delta = Field(description='流式增量内容')
+    finish_reason: Optional[Literal['stop', 'length', 'content_filter', 'tool_calls', 'function_call']] = Field(
+        default=None, description='模型停止生成 token 的原因'
+    )
+    index: int = Field(description='选择项的索引')
+    logprobs: Optional[LogprobsContent] = Field(default=None, description='该 choice 的对数概率信息')
+
+
+class StreamEvent(BaseModel):
+    """OpenAI 兼容流式响应的单个数据块"""
+
+    id: str = Field(description='对话的唯一标识符，每个 chunk 相同')
+    choices: List[Choice] = Field(
+        description='completion 选择列表，若 n>1 可包含多个元素；若设置了 include_usage，最后 chunk 可能为空'
+    )
+    created: int = Field(description='创建时间戳（秒）')
+    model: str = Field(description='模型名称')
+    object: Literal['chat.completion.chunk'] = Field(
+        default='chat.completion.chunk',
+        description='对象类型，固定为 chat.completion.chunk',
+    )
+    service_tier: Optional[Literal['auto', 'default', 'flex', 'scale', 'priority']] = Field(
+        default=None, description='实际处理请求的服务层级'
+    )
+    system_fingerprint: Optional[str] = Field(default=None, description='后端配置指纹，可用于确定性调试')
+    usage: Optional[CompletionUsage] = Field(
+        default=None,
+        description='仅在 stream_options.include_usage 为 true 时存在，通常最后 chunk 包含完整统计',
+    )

llm/clients/open_ai/response_format.py

@@ -0,0 +1,42 @@
+"""LLM 响应格式约束模型"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, field_serializer
+
+
+class JsonSchema(BaseModel):
+    """JSON Schema 模式 — 结构化输出的 schema 定义
+
+    用法::
+
+        # 直接传入 Pydantic 模型类，自动推导 JSON Schema
+        schema = JsonSchema(name='tick_output', schema=TickOutput)
+    """
+
+    name: str = Field(description='schema 名称，模型可见')
+    schema_: type[BaseModel] = Field(alias='schema', description='Pydantic 模型类，自动推导 JSON Schema')
+    strict: bool = Field(default=True, description='是否启用严格模式')
+
+    @field_serializer('schema_')
+    def _serialize_schema(self, model: type[BaseModel], _info: Any) -> dict[str, Any]:
+        return model.model_json_schema()
+
+
+class ResponseFormat(BaseModel):
+    """响应格式约束
+
+    用法::
+
+        # JSON 模式（模型自由输出 JSON）
+        ResponseFormat(type='json_object')
+
+        # 结构化输出（传入 Pydantic 模型，自动推导 schema）
+        ResponseFormat(
+            type='json_schema',
+            json_schema=JsonSchema(name='tick_output', schema=TickOutput),
+        )
+    """
+
+    type: Literal['json_object', 'json_schema'] = Field(default='json_object')
+    json_schema: JsonSchema | None = Field(default=None)

llm/clients/provider.py

@@ -0,0 +1,38 @@
+"""LLM 客户端工厂注册表"""
+
+from typing import Callable, Dict
+
+from .base import LLMClient
+from .config import LLMConfig
+
+_ClientFactory = Callable[[LLMConfig], LLMClient]
+_registry: Dict[str, _ClientFactory] = {}
+
+
+def register_client(provider: str) -> Callable[[_ClientFactory], _ClientFactory]:
+    """装饰器：将工厂函数注册到对应的 provider。
+
+    用法::
+
+        @register_client("openai")
+        def _create_openai(config: OpenAIConfig) -> LLMClient:
+            ...
+    """
+
+    def decorator(factory: _ClientFactory) -> _ClientFactory:
+        _registry[provider] = factory
+        return factory
+
+    return decorator
+
+
+def create_client(config: LLMConfig) -> LLMClient:
+    """根据配置创建对应的 LLM 客户端实例。
+
+    :param config: LLMConfig 子类实例
+    :raises ValueError: 未注册的 provider
+    """
+    factory = _registry.get(config.provider)
+    if factory is None:
+        raise ValueError(f'Unknown LLM provider: {config.provider!r}. Available: {list(_registry.keys())}')
+    return factory(config)