mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/lmproto: add openai-compatible and deepseek clients
OpenAI-compatible models, config, and async client via niquests; DeepSeek OpenAI-compat extension with reasoning fields.
This commit is contained in:
@@ -0,0 +1,8 @@
|
||||
from .openai_compat.client import DeepseekOpenAIClient as DeepseekOpenAIClient
|
||||
from .openai_compat.model import AnyChatMessage as AnyChatMessage
|
||||
from .openai_compat.model import AssistantChatMessage as AssistantChatMessage
|
||||
from .openai_compat.model import ChatCompletionsParam as ChatCompletionsParam
|
||||
from .openai_compat.model import DeepSeekReasoningEffort as DeepSeekReasoningEffort
|
||||
from .openai_compat.model import NamedChatMessage as NamedChatMessage
|
||||
from .openai_compat.model import ThinkingOptions as ThinkingOptions
|
||||
from .openai_compat.model import ToolChatMessage as ToolChatMessage
|
||||
@@ -0,0 +1,49 @@
|
||||
from typing import TYPE_CHECKING, Literal, overload
|
||||
|
||||
import msgspec
|
||||
|
||||
from ...openai_compatible.client import BaseOpenAIClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from ...openai_compatible.config import OpenAIConfig
|
||||
from ...openai_compatible.model import ChatCompletionChunk, ChatCompletionResponse
|
||||
from .model import ChatCompletionsParam
|
||||
|
||||
|
||||
class DeepseekOpenAIClient(BaseOpenAIClient):
|
||||
def __init__(self, config: OpenAIConfig) -> None:
|
||||
super().__init__(config)
|
||||
|
||||
@overload
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
||||
) -> ChatCompletionResponse: ...
|
||||
|
||||
@overload
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
||||
) -> AsyncIterator[ChatCompletionChunk]: ...
|
||||
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: bool = False
|
||||
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
||||
param = msgspec.structs.replace(param, stream=stream)
|
||||
data = self.encoder.encode(param)
|
||||
if stream:
|
||||
resp = await self.session.post(
|
||||
"/chat/completions",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
stream=True,
|
||||
)
|
||||
return self._parse_sse(resp)
|
||||
resp = await self.session.post(
|
||||
"/chat/completions",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
stream=False,
|
||||
)
|
||||
assert resp.content is not None
|
||||
return self.decoder.decode(resp.content)
|
||||
@@ -0,0 +1,60 @@
|
||||
from typing import Literal
|
||||
|
||||
from msgspec import UNSET, Struct
|
||||
|
||||
from plyngent.typedef import Unset # noqa: TC001
|
||||
|
||||
from ...openai_compatible.model import AssistantChatMessage as BaseAssistantChatMessage
|
||||
from ...openai_compatible.model import (
|
||||
ChatMessage,
|
||||
ReasoningEffort,
|
||||
ResponseFormat,
|
||||
StreamOptions,
|
||||
ToolChoiceMode,
|
||||
ToolFunctionItem,
|
||||
)
|
||||
from ...openai_compatible.model import ToolChatMessage as BaseToolChatMessage
|
||||
|
||||
type DeepSeekReasoningEffort = ReasoningEffort | Literal["max"]
|
||||
|
||||
|
||||
class NamedChatMessage(ChatMessage):
|
||||
role: Literal["system", "user"]
|
||||
name: str | Unset = UNSET
|
||||
|
||||
|
||||
class AssistantChatMessage(BaseAssistantChatMessage):
|
||||
prefix: bool | Unset = UNSET
|
||||
reasoning_content: str | Unset = UNSET
|
||||
|
||||
|
||||
class ToolChatMessage(BaseToolChatMessage):
|
||||
pass
|
||||
|
||||
|
||||
type AnyChatMessage = NamedChatMessage | AssistantChatMessage | ToolChatMessage
|
||||
|
||||
|
||||
class ThinkingOptions(Struct):
|
||||
type: Literal["enabled", "disabled"]
|
||||
|
||||
|
||||
class ChatCompletionsParam(Struct):
|
||||
messages: list[AnyChatMessage]
|
||||
model: str
|
||||
thinking: ThinkingOptions | Unset = UNSET
|
||||
reasoning_effort: DeepSeekReasoningEffort | Unset = UNSET
|
||||
max_tokens: int | Unset = UNSET
|
||||
response_format: ResponseFormat | Unset = UNSET
|
||||
stop: str | list[str] | Unset = UNSET
|
||||
stream: bool | Unset = UNSET
|
||||
stream_options: StreamOptions | Unset = UNSET
|
||||
temperature: float | Unset = UNSET
|
||||
top_p: int | Unset = UNSET
|
||||
tool_choice: ToolChoiceMode | ToolFunctionItem | Unset = UNSET
|
||||
tools: list[ToolFunctionItem] | Unset = UNSET
|
||||
logprobs: bool | Unset = UNSET
|
||||
top_logprobs: int | Unset = UNSET
|
||||
user_id: str | Unset = UNSET
|
||||
frequency_penalty: float | Unset = UNSET
|
||||
presence_penalty: float | Unset = UNSET
|
||||
@@ -0,0 +1,35 @@
|
||||
from .client import BaseOpenAIClient as BaseOpenAIClient
|
||||
from .client import OpenAIClient as OpenAIClient
|
||||
from .config import OpenAIConfig as OpenAIConfig
|
||||
from .model import AnyAssistantToolCall as AnyAssistantToolCall
|
||||
from .model import AnyChatMessage as AnyChatMessage
|
||||
from .model import AnyResponseFormat as AnyResponseFormat
|
||||
from .model import AnyToolItem as AnyToolItem
|
||||
from .model import AssistantChatMessage as AssistantChatMessage
|
||||
from .model import AudioFormatStr as AudioFormatStr
|
||||
from .model import AudioOptions as AudioOptions
|
||||
from .model import CacheRetention as CacheRetention
|
||||
from .model import ChatCompletionChoice as ChatCompletionChoice
|
||||
from .model import ChatCompletionChunk as ChatCompletionChunk
|
||||
from .model import ChatCompletionResponse as ChatCompletionResponse
|
||||
from .model import ChatCompletionsParam as ChatCompletionsParam
|
||||
from .model import ChatMessage as ChatMessage
|
||||
from .model import ChunkChoice as ChunkChoice
|
||||
from .model import DeltaMessage as DeltaMessage
|
||||
from .model import FinishReason as FinishReason
|
||||
from .model import GrammarSyntax as GrammarSyntax
|
||||
from .model import Modality as Modality
|
||||
from .model import NamedChatMessage as NamedChatMessage
|
||||
from .model import NamedRole as NamedRole
|
||||
from .model import ReasoningEffort as ReasoningEffort
|
||||
from .model import ResponseFormat as ResponseFormat
|
||||
from .model import RoleAssistant as RoleAssistant
|
||||
from .model import RoleTool as RoleTool
|
||||
from .model import SchemaResponseFormat as SchemaResponseFormat
|
||||
from .model import ServiceTier as ServiceTier
|
||||
from .model import StreamOptions as StreamOptions
|
||||
from .model import ToolChatMessage as ToolChatMessage
|
||||
from .model import ToolChoiceMode as ToolChoiceMode
|
||||
from .model import ToolFunctionItem as ToolFunctionItem
|
||||
from .model import Verbosity as Verbosity
|
||||
from .model import VoiceName as VoiceName
|
||||
@@ -0,0 +1,75 @@
|
||||
from typing import TYPE_CHECKING, Literal, overload
|
||||
|
||||
import msgspec
|
||||
import niquests
|
||||
from niquests.auth import BearerTokenAuth
|
||||
|
||||
from .config import OpenAIConfig # noqa: TC001
|
||||
from .model import ChatCompletionChunk, ChatCompletionResponse, ChatCompletionsParam
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from niquests.async_session import AsyncSession
|
||||
from niquests.models import AsyncResponse
|
||||
|
||||
|
||||
class BaseOpenAIClient:
|
||||
session: AsyncSession
|
||||
encoder: msgspec.json.Encoder
|
||||
decoder: msgspec.json.Decoder[ChatCompletionResponse]
|
||||
chunk_decoder: msgspec.json.Decoder[ChatCompletionChunk]
|
||||
|
||||
def __init__(self, config: OpenAIConfig) -> None:
|
||||
self.session = niquests.AsyncSession(
|
||||
base_url=config.base_url,
|
||||
auth=BearerTokenAuth(config.access_key_or_token),
|
||||
)
|
||||
self.encoder = msgspec.json.Encoder()
|
||||
self.decoder = msgspec.json.Decoder(ChatCompletionResponse)
|
||||
self.chunk_decoder = msgspec.json.Decoder(ChatCompletionChunk)
|
||||
|
||||
async def _parse_sse(self, resp: AsyncResponse) -> AsyncIterator[ChatCompletionChunk]:
|
||||
lines = resp.iter_lines()
|
||||
async for line in lines:
|
||||
if not line or line == b"data: [DONE]":
|
||||
continue
|
||||
if line.startswith(b"data: "):
|
||||
yield self.chunk_decoder.decode(line[6:])
|
||||
|
||||
|
||||
class OpenAIClient(BaseOpenAIClient):
|
||||
def __init__(self, config: OpenAIConfig) -> None:
|
||||
super().__init__(config)
|
||||
|
||||
@overload
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
||||
) -> ChatCompletionResponse: ...
|
||||
|
||||
@overload
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
||||
) -> AsyncIterator[ChatCompletionChunk]: ...
|
||||
|
||||
async def chat_completions(
|
||||
self, param: ChatCompletionsParam, *, stream: bool = False
|
||||
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
||||
param = msgspec.structs.replace(param, stream=stream)
|
||||
data = self.encoder.encode(param)
|
||||
if stream:
|
||||
resp = await self.session.post(
|
||||
"/chat/completions",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
stream=True,
|
||||
)
|
||||
return self._parse_sse(resp)
|
||||
resp = await self.session.post(
|
||||
"/chat/completions",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
stream=False,
|
||||
)
|
||||
assert resp.content is not None
|
||||
return self.decoder.decode(resp.content)
|
||||
@@ -0,0 +1,7 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpenAIConfig:
|
||||
access_key_or_token: str
|
||||
base_url: str = "https://api.openai.com/v1"
|
||||
@@ -0,0 +1,232 @@
|
||||
from typing import Any, Literal
|
||||
|
||||
from msgspec import UNSET, Struct
|
||||
|
||||
from plyngent.typedef import JSONSchema, Unset # noqa: TC001
|
||||
|
||||
type NamedRole = Literal["developer", "system", "user"]
|
||||
type RoleAssistant = Literal["assistant"]
|
||||
type RoleTool = Literal["tool"]
|
||||
type ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
|
||||
type ServiceTier = Literal["auto", "default", "flex", "scale", "priority"]
|
||||
type ToolChoiceMode = Literal["none", "auto", "required"]
|
||||
type AudioFormatStr = Literal["wav", "aac", "mp3", "flac", "opus", "pcm16"]
|
||||
type VoiceName = Literal["alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar"]
|
||||
type FinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
|
||||
type Modality = Literal["text", "audio"]
|
||||
type CacheRetention = Literal["in_memory", "24h"]
|
||||
type Verbosity = Literal["low", "medium", "high"]
|
||||
type GrammarSyntax = Literal["lark", "regex"]
|
||||
|
||||
|
||||
class ChatMessage(Struct):
|
||||
content: str
|
||||
|
||||
|
||||
class NamedChatMessage(ChatMessage):
|
||||
role: NamedRole
|
||||
name: str | Unset = UNSET
|
||||
|
||||
|
||||
class IDObject(Struct):
|
||||
id: str
|
||||
|
||||
|
||||
class AssistantFunctionTool(Struct):
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class AssistantFunctionToolCall(Struct):
|
||||
id: str
|
||||
type: Literal["function"]
|
||||
function: AssistantFunctionTool
|
||||
|
||||
|
||||
class AssistantCustomTool(Struct):
|
||||
name: str
|
||||
input: str
|
||||
|
||||
|
||||
class AssistantCustomToolCall(Struct):
|
||||
id: str
|
||||
type: Literal["custom"]
|
||||
custom: AssistantCustomTool
|
||||
|
||||
|
||||
type AnyAssistantToolCall = AssistantFunctionToolCall | AssistantCustomToolCall
|
||||
|
||||
|
||||
class AssistantChatMessage(ChatMessage):
|
||||
role: RoleAssistant
|
||||
name: str | Unset = UNSET
|
||||
audio: IDObject | Unset = UNSET
|
||||
refusal: str | Unset = UNSET
|
||||
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
|
||||
|
||||
|
||||
class ToolChatMessage(ChatMessage):
|
||||
role: RoleTool
|
||||
tool_call_id: str
|
||||
|
||||
|
||||
type AnyChatMessage = NamedChatMessage | AssistantChatMessage | ToolChatMessage
|
||||
|
||||
|
||||
class ResponseFormat(Struct):
|
||||
type: Literal["text", "json_object"]
|
||||
|
||||
|
||||
class SchemaResponseFormat(Struct):
|
||||
type: Literal["json_schema"]
|
||||
json_schema: JSONSchema
|
||||
|
||||
|
||||
type AnyResponseFormat = ResponseFormat | SchemaResponseFormat
|
||||
|
||||
|
||||
class ToolFunction(Struct):
|
||||
name: str
|
||||
description: str | Unset = UNSET
|
||||
parameters: JSONSchema | Unset = UNSET
|
||||
strict: bool | Unset = UNSET
|
||||
|
||||
|
||||
class ToolFunctionItem(Struct):
|
||||
type: Literal["function"]
|
||||
function: ToolFunction
|
||||
|
||||
|
||||
class TextFormat(Struct):
|
||||
type: Literal["text"]
|
||||
|
||||
|
||||
class GrammarDefinition(Struct):
|
||||
syntax: GrammarSyntax
|
||||
definition: str
|
||||
|
||||
|
||||
class GrammarFormat(Struct):
|
||||
type: Literal["grammar"]
|
||||
grammar: GrammarDefinition
|
||||
|
||||
|
||||
class ToolCustom(Struct):
|
||||
name: str
|
||||
description: str | Unset = UNSET
|
||||
format: TextFormat | GrammarFormat | Unset = UNSET
|
||||
|
||||
|
||||
class ToolCustomItem(Struct):
|
||||
type: Literal["custom"]
|
||||
custom: ToolCustom
|
||||
|
||||
|
||||
type AnyToolItem = ToolFunctionItem | ToolCustomItem
|
||||
|
||||
|
||||
class AudioOptions(Struct):
|
||||
format: AudioFormatStr
|
||||
voice: VoiceName | IDObject
|
||||
|
||||
|
||||
class ModerationOptions(Struct):
|
||||
model: str
|
||||
|
||||
|
||||
class PredictionOptions(Struct):
|
||||
type: Literal["content"]
|
||||
content: str
|
||||
|
||||
|
||||
class StreamOptions(Struct):
|
||||
include_obfuscation: bool | Unset = UNSET
|
||||
include_usage: bool | Unset = UNSET
|
||||
|
||||
|
||||
class AllowedTools(Struct):
|
||||
mode: Literal["auto", "required"]
|
||||
tools: list[AnyToolItem]
|
||||
|
||||
|
||||
class AllowedToolChoice(Struct):
|
||||
type: Literal["allowed_tools"]
|
||||
allowed_tools: AllowedTools
|
||||
|
||||
|
||||
class ChatCompletionsParam(Struct):
|
||||
messages: list[AnyChatMessage]
|
||||
model: str
|
||||
audio: AudioOptions | Unset = UNSET
|
||||
frequency_penalty: float | Unset = UNSET
|
||||
logit_bias: dict[int, int] | Unset = UNSET
|
||||
logprobs: bool | Unset = UNSET
|
||||
max_completion_tokens: int | Unset = UNSET
|
||||
max_tokens: int | Unset = UNSET
|
||||
metadata: dict[str, str] | Unset = UNSET
|
||||
modalities: set[Modality] | Unset = UNSET
|
||||
moderation: ModerationOptions | Unset = UNSET
|
||||
n: int | Unset = UNSET
|
||||
parallel_tool_calls: bool | Unset = UNSET
|
||||
prediction: PredictionOptions | Unset = UNSET
|
||||
presence_penalty: float | Unset = UNSET
|
||||
prompt_cache_key: str | Unset = UNSET
|
||||
prompt_cache_retention: CacheRetention | Unset = UNSET
|
||||
reasoning_effort: ReasoningEffort | Unset = UNSET
|
||||
response_format: AnyResponseFormat | Unset = UNSET
|
||||
safety_identifier: str | Unset = UNSET
|
||||
seed: int | Unset = UNSET
|
||||
service_tier: ServiceTier | Unset = UNSET
|
||||
stop: str | list[str] | Unset = UNSET
|
||||
store: bool | Unset = UNSET
|
||||
stream: bool | Unset = UNSET
|
||||
stream_options: StreamOptions | Unset = UNSET
|
||||
temperature: float | Unset = UNSET
|
||||
tool_choice: ToolChoiceMode | AllowedToolChoice | AnyToolItem | Unset = UNSET
|
||||
tools: list[AnyToolItem] | Unset = UNSET
|
||||
top_logprobs: int | Unset = UNSET
|
||||
top_p: int | Unset = UNSET
|
||||
user: str | Unset = UNSET
|
||||
verbosity: Verbosity | Unset = UNSET
|
||||
web_search_options: dict[str, Any] | Unset = UNSET
|
||||
|
||||
|
||||
class ChatCompletionChoice(Struct):
|
||||
index: int
|
||||
message: AssistantChatMessage
|
||||
logprobs: dict[str, Any]
|
||||
finish_reason: FinishReason
|
||||
|
||||
|
||||
class ChatCompletionResponse(Struct):
|
||||
id: str
|
||||
object: Literal["chat.completion"]
|
||||
created: int
|
||||
model: str
|
||||
choices: list[ChatCompletionChoice]
|
||||
system_fingerprint: str
|
||||
usage: dict[str, Any]
|
||||
moderation: dict[str, Any] | Unset = UNSET
|
||||
service_tier: ServiceTier | Unset = UNSET
|
||||
|
||||
|
||||
class DeltaMessage(Struct):
|
||||
role: RoleAssistant | Unset = UNSET
|
||||
content: str | Unset = UNSET
|
||||
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
|
||||
|
||||
|
||||
class ChunkChoice(Struct):
|
||||
index: int
|
||||
delta: DeltaMessage
|
||||
logprobs: dict[str, Any] | Unset = UNSET
|
||||
finish_reason: FinishReason | None | Unset = UNSET
|
||||
|
||||
|
||||
class ChatCompletionChunk(Struct):
|
||||
id: str
|
||||
object: Literal["chat.completion.chunk"]
|
||||
created: int
|
||||
model: str
|
||||
choices: list[ChunkChoice]
|
||||
usage: dict[str, Any] | Unset = UNSET
|
||||
Reference in New Issue
Block a user