citrus-code 0.1.0__py3-none-any.whl

citrus_code/ai/providers/deepseek.py

@@ -0,0 +1,104 @@
+from collections.abc import AsyncIterator
+
+import httpx
+from pydantic import SecretStr
+
+from ..events import AssistantMessageEvent
+from ..types import (
+    ImageContent, Message, ModelContext, ModelOptions, ModelSpec,
+    TextContent, ToolResultMessage, UserMessage,
+)
+from .chat_completions import (
+    build_chat_request,
+    stream_chat_data,
+    stream_assistant_events,
+)
+
+
+DEEPSEEK_VISION_MODEL_IDS = frozenset({
+    "deepseek-flash",
+    "deepseek-v4-flash",
+    "deepseek-v4-flash-vision-exp",
+})
+DEEPSEEK_IMAGE_MIME_TYPES = frozenset({
+    "image/png", "image/jpeg", "image/gif", "image/webp",
+})
+
+
+class DeepSeekProvider:
+    id: str = "deepseek"
+    name: str = "DeepSeek"
+
+    def __init__(
+        self,
+        client: httpx.AsyncClient,
+        model: ModelSpec,
+        api_key: SecretStr,
+    ) -> None:
+        if model.provider != self.id:
+            raise ValueError(
+                "The model does not belong to this provider"
+            )
+
+        if model.api != "chat-completions":
+            raise ValueError(
+                "This provider requires the chat-completions API"
+            )
+
+        self._client = client
+        self._model = model
+        self._api_key = api_key
+
+    def get_models(self) -> list[ModelSpec]:
+        return [self._model]
+
+    def stream(
+        self,
+        model: ModelSpec,
+        context: ModelContext,
+        options: ModelOptions | None = None,
+    ) -> AsyncIterator[AssistantMessageEvent]:
+        if model != self._model:
+            raise ValueError(
+                "The requested model is not registered "
+                "with this provider"
+            )
+
+        if model.supports_images:
+            messages: list[Message] = []
+            for message in context.messages:
+                if isinstance(message, (UserMessage, ToolResultMessage)) and not isinstance(message.content, str):
+                    content = [
+                        TextContent(text=(
+                            f"[图片未发送：DeepSeek 不支持 {block.mime_type}；"
+                            "请先转换为 PNG、JPEG、GIF 或 WebP。]"
+                        ))
+                        if isinstance(block, ImageContent) and block.mime_type not in DEEPSEEK_IMAGE_MIME_TYPES
+                        else block
+                        for block in message.content
+                    ]
+                    message = message.model_copy(update={"content": content})
+                messages.append(message)
+            # Transform only the request. An unsupported old attachment must
+            # not make every later request in the session fail with HTTP 400.
+            context = ModelContext(messages=messages, tools=context.tools)
+
+        body = build_chat_request(
+            model=model,
+            context=context,
+            options=options,
+        )
+
+        body["thinking"] = {"type": "disabled"}
+
+        data_stream = stream_chat_data(
+            client=self._client,
+            model=model,
+            api_key=self._api_key,
+            body=body,
+        )
+
+        return stream_assistant_events(
+            model=model,
+            data_stream=data_stream,
+        )

citrus_code/ai/providers/fake.py

@@ -0,0 +1,101 @@
+import asyncio
+import time
+from collections.abc import AsyncIterator
+
+from ..events import (
+    AssistantMessageEvent,
+    DoneEvent,
+    StartEvent,
+    TextDeltaEvent,
+    TextEndEvent,
+    TextStartEvent,
+)
+from ..types import (
+    AssistantMessage,
+    ModelContext,
+    ModelOptions,
+    ModelSpec,
+    TextContent,
+)
+class FakeProvider:
+    id: str = "fake"
+    name: str = "Fake Provider"
+
+    def __init__(
+            self,
+            chunks: list[str],
+            delay_seconds: float=0.2
+    ) -> None:
+        if delay_seconds < 0:
+            raise ValueError("delay_seconds must not be negative")
+
+        self._chunks = chunks.copy()
+        self._delay_seconds = delay_seconds
+
+        self._model = ModelSpec(
+            id="fake-text",
+            name="Fake Text Model",
+            provider=self.id,
+            api="fake",
+            base_url="fake://local",
+        )
+
+    def get_models(self) -> list[ModelSpec]:
+        return [self._model]
+
+    async def stream(
+        self,
+        model: ModelSpec,
+        context: ModelContext,
+        options: ModelOptions | None = None,
+    ) -> AsyncIterator[AssistantMessageEvent]:
+        if(
+            model.provider != self.id
+            or model.id !=self._model.id
+            or model.api !=self._model.api
+        ):
+            raise ValueError("The request model is not supported")
+
+        message = AssistantMessage(
+            content=[],
+            api=model.api,
+            provider=model.provider,
+            model=model.id,
+            usage=None,
+            stop_reason="pending",
+            timestamp=int(time.time()*1000),
+        )
+        yield StartEvent(partial=message)
+
+        text_block = TextContent(text="")
+        message.content.append(text_block)
+
+        yield TextStartEvent(
+            content_index=0,
+            partial=message,
+        )
+
+        for chunk in self._chunks:
+            await asyncio.sleep(self._delay_seconds)
+
+            text_block.text += chunk
+
+            yield TextDeltaEvent(
+                content_index=0,
+                delta=chunk,
+                partial=message,
+            )
+
+        yield TextEndEvent(
+            content_index=0,
+            content=text_block.text,
+            partial=message,
+        )
+
+        message.stop_reason = "stop"
+
+        yield DoneEvent(
+            reason="stop",
+            message=message,
+        ) 
+        

citrus_code/ai/providers/tool_call_buffer.py

@@ -0,0 +1,73 @@
+from pydantic import JsonValue, TypeAdapter
+
+from ..types import ToolCall
+from .chat_stream_types import ChatToolCallDelta
+
+
+ARGUMENTS_ADAPTER = TypeAdapter(dict[str, JsonValue])
+
+
+class ToolCallBuffer:
+    def __init__(
+            self,
+            tool_index: int,
+            content_index: int
+    ) -> None:
+        self.tool_index = tool_index
+        self.content_index = content_index
+        self.arguments_text = ""
+
+        self.tool_call = ToolCall(
+            id="",
+            name="",
+            arguments={},
+        )
+
+
+    def append(self, delta: ChatToolCallDelta) -> str:
+        if delta.index != self.tool_index:
+            raise ValueError(
+                "The delta belongs to a different tool call"
+            )
+
+        if delta.id is not None:
+            if self.tool_call.id and self.tool_call.id !=delta.id:
+                raise ValueError("The tool call ID changed")
+
+            self.tool_call.id = delta.id
+
+        if delta.function is None:
+            return ""
+
+        if delta.function.name is not None:
+            name = delta.function.name
+
+            if self.tool_call.name and self.tool_call.name != name:
+                raise ValueError("The tool name changed")
+
+            self.tool_call.name = name
+
+        fragment = delta.function.arguments
+
+        if fragment is None:
+            return ""
+
+        self.arguments_text += fragment
+
+        return fragment
+
+
+    def finish(self) -> ToolCall:
+        if not self.tool_call.id:
+            raise ValueError("The tool call ID is missing")
+
+        if not self.tool_call.name:
+            raise ValueError("The tool name is missing")
+
+        arguments = ARGUMENTS_ADAPTER.validate_json(
+            self.arguments_text
+        )
+
+        self.tool_call.arguments = arguments
+
+        return self.tool_call

citrus_code/ai/registry.py

@@ -0,0 +1,40 @@
+from .provider import Provider
+from .types import ModelSpec
+
+
+class ProviderRegistry:
+    def __init__(self) -> None:
+        self._providers: dict[str,Provider] = {}
+
+    def register(self,provider: Provider) -> None:
+        if provider.id in self._providers:
+            raise ValueError(
+                f"Provider already registered: {provider.id}"
+            )
+        self._providers[provider.id]=provider
+
+    def get_provider(self,provider_id: str) -> Provider:
+        provider = self._providers.get(provider_id)
+
+        if provider is None:
+            raise KeyError(
+                f"Provider not found: {provider_id}"
+            ) 
+
+        return provider
+
+    def get_model(
+            self,
+            provider_id: str,
+            model_id: str,
+    ) -> ModelSpec:
+        provider = self.get_provider(provider_id)
+
+        for model in provider.get_models():
+            if model.id == model_id:
+                return model
+
+        raise KeyError(
+            f"Model not found: {provider_id}/{model_id}"
+        )
+        

citrus_code/ai/types.py

@@ -0,0 +1,142 @@
+from typing import Annotated,Literal
+
+from pydantic import BaseModel,ConfigDict,Field,JsonValue
+
+
+class TextContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ImageContent(BaseModel):
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+    type: Literal["image"] = "image"
+    data: str
+    mime_type: str = Field(alias="mimeType")
+
+
+ToolContent = Annotated[
+    TextContent | ImageContent,
+    Field(discriminator="type"),
+]
+
+
+class ToolCall(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    type: Literal["toolCall"]="toolCall"
+    id: str
+    name: str
+    arguments: dict[str,JsonValue]
+
+class SystemMessage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    role: Literal["system"] = "system"
+    content: str | list[TextContent]
+    timestamp: int 
+
+
+class UserMessage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    role: Literal["user"] = "user"
+    content: str | list[ToolContent]
+    timestamp: int
+
+
+class ToolResultMessage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    role: Literal["toolResult"] = "toolResult"
+    tool_call_id: str
+    tool_name: str
+    content: list[ToolContent]
+    is_error: bool
+    timestamp: int
+
+StopReason = Literal[
+    "pending",
+    "stop",
+    "length",
+    "toolUse",
+    "error",
+    "aborted",
+]
+
+
+class Usage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    input_tokens: int = Field(ge=0)
+    output_tokens: int = Field(ge=0)
+    total_tokens: int = Field(ge=0)
+
+    cache_read_tokens: int | None = Field(default=None,ge=0)
+    cache_write_tokens: int | None = Field(default=None,ge=0)
+
+
+AssistantContent = Annotated[
+    TextContent | ToolCall,
+    Field(discriminator="type"),
+]
+
+
+class AssistantMessage(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    role: Literal["assistant"] = "assistant"
+    content: list[AssistantContent]
+
+    api: str
+    provider: str
+    model: str
+
+    usage: Usage | None
+    stop_reason: StopReason
+    error_message: str | None = None
+    timestamp: int
+
+
+Message = Annotated[
+    SystemMessage | UserMessage | AssistantMessage | ToolResultMessage,
+    Field(discriminator="role"),
+]
+
+
+class ModelSpec(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    id: str
+    name: str
+    provider: str
+    api: str
+    base_url: str
+
+    context_window: int | None = Field(default=None,gt=0)
+    max_output_tokens: int | None = Field(default=None,gt=0)
+    supports_images: bool = False
+
+
+class ToolDefinition(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    name: str
+    description: str
+    parameters: dict[str,JsonValue]
+
+
+class ModelContext(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    messages: list[Message]
+    tools: list[ToolDefinition] = Field(default_factory=list)
+
+
+class ModelOptions(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    max_output_tokens: int | None = Field(default=None,gt=0)
+    temperature: float | None = None

citrus_code/cli/__init__.py

@@ -0,0 +1,5 @@
+"""Command-line application for Citrus Code."""
+
+from .main import main
+
+__all__ = ["main"]