Metadata-Version: 2.4
Name: my-llmkit
Version: 0.3.2
Summary: a kit for building LLM applications
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: anthropic>=0.75.0
Requires-Dist: google-genai>=2.14.0
Requires-Dist: httpx>=0.28.1
Requires-Dist: mcp>=1.12.4
Requires-Dist: openai>=2.14.0
Requires-Dist: openai-agents>=0.6.2
Requires-Dist: pytest>=9.0.2
Requires-Dist: pytest-asyncio>=1.3.0
Requires-Dist: python-dotenv>=1.2.1

# my-llmkit

一个统一的 LLM 聊天接口工具包，支持多种 AI 模型提供商。

## 架构概览

```
┌─────────────────────────────────────────────────────────────────────┐
│                          my-llmkit                                  │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  ┌──────────────────────────────────────────────────────────────┐  │
│  │                    chat 模块 (核心)                           │  │
│  ├──────────────────────────────────────────────────────────────┤  │
│  │                                                              │  │
│  │  LLMChatCompletion (抽象基类)                                │  │
│  │         ▲                 ▲                                  │  │
│  │         │                 │                                  │  │
│  │         │                 │                                  │  │
│  │  ┌──────┴──────┐   ┌──────┴──────┐                          │  │
│  │  │   OpenAI    │   │   Claude    │                          │  │
│  │  │ Compatible  │   │   Client    │                          │  │
│  │  └─────────────┘   └─────────────┘                          │  │
│  │                                                              │  │
│  │  ┌─────────────────────────────────────────────┐            │  │
│  │  │  ChatCompletionStreamRunner                 │            │  │
│  │  │  - 流式事件处理                              │            │  │
│  │  │  - 工具调用管理                              │            │  │
│  │  │  - 多轮对话控制                              │            │  │
│  │  └─────────────────────────────────────────────┘            │  │
│  │           │                                                  │  │
│  │           ▼                                                  │  │
│  │  ┌─────────────────────────────────────────────┐            │  │
│  │  │  ChatCompletionStreamProcessor              │            │  │
│  │  │  - 处理 UnifiedChunk 流                      │            │  │
│  │  │  - 生成统一事件                              │            │  │
│  │  │  - 结构化输出解析                            │            │  │
│  │  └─────────────────────────────────────────────┘            │  │
│  │                                                              │  │
│  │  核心组件:                                                   │  │
│  │  • UnifiedMessage    - 统一消息格式                          │  │
│  │  • UnifiedChunk      - 统一流式块                            │  │
│  │  • ToolFunctions     - 工具函数管理                          │  │
│  │  • ToolExecutor      - 工具执行器                            │  │
│  │  • ModelSettings     - 模型配置                              │  │
│  └──────────────────────────────────────────────────────────────┘  │
│                                                                     │
│  ┌──────────────────────────────────────────────────────────────┐  │
│  │                    mcp 模块 (扩展工具)                        │  │
│  ├──────────────────────────────────────────────────────────────┤  │
│  │                                                              │  │
│  │  MCPServerBase (抽象基类)                                    │  │
│  │         ▲                ▲               ▲                  │  │
│  │         │                │               │                  │  │
│  │    ┌────┴────┐     ┌────┴────┐    ┌────┴────┐             │  │
│  │    │  Stdio  │     │   SSE   │    │  HTTP   │             │  │
│  │    │ Server  │     │ Server  │    │ Server  │             │  │
│  │    └─────────┘     └─────────┘    └─────────┘             │  │
│  │                                                              │  │
│  │  MCPClientPool     - 按需连接、复用与自动重连                │  │
│  │  MCPServersContext - 短生命周期的上下文管理                  │  │
│  │  MCPManager        - 配置文件管理                            │  │
│  └──────────────────────────────────────────────────────────────┘  │
│                                                                     │
│  ┌──────────────────────────────────────────────────────────────┐  │
│  │                  models 模块 (能力查询)                       │  │
│  ├──────────────────────────────────────────────────────────────┤  │
│  │                                                              │  │
│  │  • supports_reasoning()        - 推理能力检测                │  │
│  │  • supports_vision()           - 视觉能力检测                │  │
│  │  • supports_function_calling() - 工具调用检测                │  │
│  │  • get_model_info()            - 获取模型信息                │  │
│  │                                                              │  │
│  │  数据源: litellm, models.dev (自动缓存)                      │  │
│  └──────────────────────────────────────────────────────────────┘  │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

数据流向:

  User Input (UnifiedMessage)
         │
         ▼
  LLMChatCompletion
         │
         ├──> OpenAI API / Claude API / ...
         │
         ▼
  Stream (UnifiedChunk)
         │
         ▼
  ChatCompletionStreamProcessor
         │
         ├──> Content Events
         ├──> Reasoning Events
         ├──> Tool Call Events
         ├──> Usage Events
         │
         ▼
  ChatCompletionStreamRunner
         │
         ├──> Tool Execution (本地函数 + MCP 服务器)
         ├──> Multi-turn Conversation
         │
         ▼
  Final Response (UnifiedResponse + 结构化输出)
```

## 模块

- `my_llmkit.chat` - 统一的聊天接口
- `my_llmkit.image` - 多供应商图片生成与编辑接口
- `my_llmkit.mcp` - MCP (Model Context Protocol) 支持
- `my_llmkit.models` - 模型能力查询

详细文档：

- [MCP 配置、连接池与生命周期](docs/mcp.md)
- [统一工具结果与多模态 MCP 返回值](docs/tool-results.md)
- [工具执行审批](docs/tool-approval.md)
- [每轮请求消息准备](docs/message-preparer.md)

## 安装

```bash
pip install -e .
```

## Chat 模块使用说明

### 基本用法

#### 1. OpenAI 兼容接口

支持所有兼容 OpenAI API 的模型，包括 GPT、Gemini、DeepSeek、Kimi、Grok 等。

```python
from my_llmkit.chat import OpenAICompatibleChatCompletion, UnifiedMessage

# 创建客户端
client = OpenAICompatibleChatCompletion(
    api_key="your-api-key",
    api_base="https://api.openai.com/v1",
    model="gpt-4"
)

# 发送消息
messages = [UnifiedMessage(role="user", content="你好")]
result = client.run_stream(messages=messages)

# 处理流式响应
async for event in result.stream_event():
    if event.type == "content":
        print(event.content, end="", flush=True)
    elif event.type == "usage":
        print(f"\n使用量: {event.usage}")
```

#### 2. Claude 接口

```python
from my_llmkit.chat import ClaudeChatCompletion, UnifiedMessage
from my_llmkit.chat.model_settings import ModelSettings

# 创建客户端（带思考模式）
model_settings = ModelSettings(
    include_usage=True,
    max_tokens=30000,
    extra_body={
        "thinking": {"type": "enabled", "budget_tokens": 10000}
    }
)

client = ClaudeChatCompletion(
    api_key="your-api-key",
    api_base="https://api.anthropic.com",
    model="claude-sonnet-4.5",
    model_settings=model_settings
)

messages = [UnifiedMessage(role="user", content="解释一下量子计算")]
result = client.run_stream(messages=messages)

async for event in result.stream_event():
    if event.type == "reasoning_content":
        print(f"[思考] {event.content}", end="", flush=True)
    elif event.type == "content":
        print(event.content, end="", flush=True)
```

### 高级功能

#### 1. 工具调用

```python
from my_llmkit.chat import ToolFunctions
from datetime import datetime

def get_weather_tool(day: str):
    """
    获取指定日期的天气信息

    Args:
        day: 日期，格式为 YYYY-MM-DD
    """
    return f"{day} 的天气是晴天，温度 6°C"

def now_tool() -> str:
    """
    获取当前日期和时间
    """
    now = datetime.now()
    return f"当前时间: {now.strftime('%Y-%m-%d %H:%M:%S')}"

# 使用工具
messages = [UnifiedMessage(role="user", content="明天天气怎么样")]
result = client.run_stream(
    messages=messages,
    tools=ToolFunctions(now_tool, get_weather_tool)
)

async for event in result.stream_event():
    if event.type == "tool_call_start":
        print(f"\n[调用工具] {event.function_name}")
    elif event.type == "tool_call_result":
        print(f"\n[工具结果] {event.function_name}: {event.function_result}")
    elif event.type == "content":
        print(event.content, end="", flush=True)
```

#### 2. MCP 工具集成

短生命周期任务可以使用 `MCPServersContext`：

```python
from my_llmkit.mcp import MCPServersContext

async with MCPServersContext("~/mcp.json") as servers:
    messages = [UnifiedMessage(role="user", content="~/Downloads 里有哪些文件")]
    result = client.run_stream(
        messages=messages,
        tools=ToolFunctions(now_tool),
        mcp_servers=servers
    )

    async for event in result.stream_event():
        if event.type == "content":
            print(event.content, end="", flush=True)
```

服务端、TUI 等长生命周期进程建议使用 `MCPClientPool`。连接池会按需连接，
在多个请求间复用连接，并在工具调用失败后让稳定 handle 在下次调用时自动重连：

```python
from my_llmkit.mcp import MCPClientPool

pool = MCPClientPool.from_file("~/mcp.json")
try:
    servers = await pool.resolve(["filesystem"])
    result = client.run_stream(messages=messages, mcp_servers=servers)
    async for event in result.stream_event():
        if event.type == "content":
            print(event.content, end="", flush=True)
finally:
    await pool.close()
```

配置中的 `tool_timeout` 控制单次 MCP 工具调用超时，默认 120 秒；
设为 `0`、负数或 `null` 可禁用。完整配置和生命周期说明见
[MCP 文档](docs/mcp.md)。

#### 3. 结构化输出

##### JSON Schema 模式（推荐）

```python
from pydantic import BaseModel

class TimeResult(BaseModel):
    local_time: str
    utc_time: str
    tz: str
    weekday: str

messages = [UnifiedMessage(role="user", content="现在几点？")]
result = client.run_stream(
    messages=messages,
    tools=ToolFunctions(now_tool),
    response_format=TimeResult
)

# 处理流式输出
async for event in result.stream_event():
    if event.type == "content":
        print(event.content, end="", flush=True)

# 获取结构化输出结果
output: TimeResult = result.output_result
print(f"\n结构化输出: {output.local_time}, {output.weekday}")
```

##### JSON Object 模式

```python
messages = [UnifiedMessage(role="user", content="""现在几点？
请使用以下 JSON 格式来回答：
{
  "local_time": "本地时间",
  "utc_time": "UTC时间",
  "tz": "时区",
  "weekday": "星期几"
}
""")]

result = client.run_stream(
    messages=messages,
    tools=ToolFunctions(now_tool),
    response_format={"type": "json_object"}
)

async for event in result.stream_event():
    if event.type == "content":
        print(event.content, end="", flush=True)

# 获取 JSON 输出（返回 dict）
output: dict = result.output_result
print(f"\nJSON 输出: {output}")
```

#### 4. 图片输入

```python
from my_llmkit.chat import TextContent, ImageContent

messages = [
    UnifiedMessage(
        role="user",
        content=[
            TextContent(text="图片里有什么"),
            ImageContent(image_url="https://example.com/image.png"),
            # 或从本地文件加载
            # ImageContent.from_file("/path/to/image.png"),
        ]
    )
]

result = client.run_stream(messages=messages)
```

#### 5. 文档输入（PDF）

支持 PDF 文档输入，可以通过 URL 或本地文件方式。

##### Claude（支持 URL 和 Base64）

```python
from my_llmkit.chat import TextContent, DocumentContent

# 通过 URL
messages = [
    UnifiedMessage(
        role="user",
        content=[
            TextContent(text="总结这个文档的主要内容"),
            DocumentContent.from_url("https://example.com/document.pdf"),
        ]
    )
]

# 通过本地文件
messages = [
    UnifiedMessage(
        role="user",
        content=[
            TextContent(text="分析这个研究报告"),
            DocumentContent.from_file("/path/to/report.pdf"),
        ]
    )
]

result = client.run_stream(messages=messages)
async for event in result.stream_event():
    if event.type == "content":
        print(event.content, end="", flush=True)
```

##### OpenAI（仅支持 Base64）

```python
from my_llmkit.chat import TextContent, DocumentContent

# OpenAI 只支持 base64 方式，建议使用 from_file
messages = [
    UnifiedMessage(
        role="user",
        content=[
            TextContent(text="这个文档讲了什么？"),
            DocumentContent.from_file("/path/to/document.pdf"),
        ]
    )
]

# 或手动指定 base64 数据
messages = [
    UnifiedMessage(
        role="user",
        content=[
            TextContent(text="分析这个文档"),
            DocumentContent.from_base64(
                base64_data="your-base64-encoded-pdf-data",
                filename="document.pdf"
            ),
        ]
    )
]

result = client.run_stream(messages=messages)
async for event in result.stream_event():
    if event.type == "content":
        print(event.content, end="", flush=True)
```

**注意事项：**

- Claude 支持 URL 和 Base64 两种方式
- OpenAI 仅支持 Base64 方式，使用 URL 会抛出异常
- 当前仅支持 PDF 格式（`application/pdf`）
- 使用 `from_file()` 方法会自动读取文件并转换为 base64

#### 6. 推理模式（Reasoning）

```python
from openai.types import Reasoning

# 创建支持推理的客户端
from my_llmkit.chat.model_settings import ModelSettings

model_settings = ModelSettings(
    include_usage=True,
    reasoning=Reasoning(effort="medium")  # 可选: low, medium, high
)

client = OpenAICompatibleChatCompletion(
    api_key="your-api-key",
    api_base="https://api.provider.com/v1",
    model="gpt-5.2",  # 或其他支持推理的模型
    model_settings=model_settings
)

messages = [UnifiedMessage(role="user", content="解决这个数学问题：...")]
result = client.run_stream(messages=messages)

async for event in result.stream_event():
    if event.type == "reasoning_content":
        # 推理过程
        print(f"[推理] {event.content}", end="", flush=True)
    elif event.type == "content":
        # 最终回答
        print(event.content, end="", flush=True)
```

### 事件类型

流式响应中可能返回的事件类型：

- `content` - 普通文本内容
- `reasoning_content` - 推理过程内容（仅支持推理的模型）
- `tool_call_start` - 工具调用开始
- `tool_call_result` - 工具调用结果
- `usage` - Token 使用量统计

### 支持的模型提供商

- **OpenAI**: GPT-4, GPT-5.2, etc.
- **Anthropic**: Claude Sonnet 4.5, Claude Sonnet 4.6, Claude Haiku 4.5
- **Google**: Gemini 3 Pro, Gemini 3 Flash
- **DeepSeek**: DeepSeek Reasoner, DeepSeek Chat
- **Moonshot**: Kimi K2 Thinking, Kimi K2.5, Kimi K2 Turbo
- **字节跳动**: Doubao Seed 1.8, Doubao Seed 2.0 Pro
- **百度**: ERNIE X 1.1, ERNIE 5 Thinking
- **阿里**: Qwen Plus
- **xAI**: Grok 4.1 Fast, Grok 4 Fast, Grok 4
- **MiniMax**: MiniMax M2.5

以及通过 OpenRouter、ZenMux、302AI 等聚合平台访问的其他模型。

## Image 模块使用说明

### Python API

`my_llmkit.image` 提供统一的异步图片生成与编辑接口，支持 OpenAI
兼容接口、OpenRouter、Google Gemini，以及 ZenMux、302AI 等聚合平台。

```python
import asyncio
from pathlib import Path

from my_llmkit.image import draw, suffix_for_mime_type


async def main() -> None:
    images = await draw(
        model_path="openrouter/gpt-image-2",
        prompt="一只放在白色桌面上的陶瓷杯，产品摄影风格",
        api_key="your-openrouter-api-key",
        size="1024x1024",
        number=1,
        input_images=[
            "./reference.png",
            "https://example.com/reference.webp",
        ],
    )

    image = images[0]
    output_path = Path(f"result{suffix_for_mime_type(image.mime_type)}")
    output_path.write_bytes(image.data)


asyncio.run(main())
```

`draw()` 返回 `GeneratedImage` 列表，每项包含图片二进制数据 `data` 和
MIME 类型 `mime_type`。`input_images` 可以传入本地文件路径或 HTTP/HTTPS
URL；URL 图片会缓存在 `/tmp/my_llmkit/input_images`。

当前注册的模型路径：

- `openrouter/gpt-image-2`
- `openrouter/gemini-3-pro-image`
- `openrouter/gemini-3.1-flash-image`
- `openrouter/grok-image`
- `zenmux/gpt-image-2`
- `zenmux/gemini-3-pro-image`
- `zenmux/gemini-3.1-flash-image`
- `302ai/gpt-image-2`
- `google/gemini-3-pro-image`
- `google/gemini-3.1-flash-image`

根据模型配置对应的 API Key：

```shell
OPENROUTER_API_KEY=
ZENMUX_API_KEY=
AI302_API_KEY=
GEMINI_API_KEY=
```

Python API 优先使用 `draw(api_key=...)` 显式传入的 Key；未传入时，
才会根据模型配置读取已存在的进程环境变量。Python API 不会主动
加载 `.env` 文件；需要使用 `.env` 时，由调用方先行加载。

OpenAI 兼容接口和 OpenRouter 会直接使用 `size`。Gemini 支持以下格式：

- `16:9`：设置宽高比
- `2K` 或 `4K`：设置图片尺寸
- `16:9@2K`：同时设置宽高比和图片尺寸
- `auto`：不显式指定 Gemini 图片配置

`web_search` 参数当前尚未实现，传入非 `False` 值会抛出
`NotImplementedError`。

### Image CLI

安装项目后可以使用 `draw` 命令生成图片：

```bash
draw \
  --model-path openrouter/gpt-image-2 \
  --prompt "一只放在白色桌面上的陶瓷杯，产品摄影风格" \
  --output ./result.png \
  --size 1024x1024
```

CLI 启动时会依次加载项目根目录 `.env` 和 `~/.gede/config/.env`，
已存在的进程环境变量不会被 dotenv 文件覆盖。

使用 `-` 从标准输入读取提示词：

```bash
echo "白色背景上的红色立方体" | draw \
  --model-path openrouter/gpt-image-2 \
  --prompt - \
  --output ./result.png
```

通过重复 `--input-image` 提供本地或 URL 参考图：

```bash
draw \
  --model-path openrouter/gpt-image-2 \
  --prompt "根据参考对象生成产品照片" \
  --input-image ./reference.png \
  --input-image https://example.com/reference.webp \
  --output ./result.png
```

`--number` 可以一次生成多张图片。第一张使用指定输出路径，后续文件依次
增加 `-2`、`-3` 后缀。`--log-level` 支持 `DEBUG`、`INFO`、`WARNING`、
`ERROR` 和 `CRITICAL`。

## 测试

测试代码位于 `tests/`，当前入口是 `tests/model_tests.py`。这些用例会真实调用模型接口，运行前需要在 `~/.gede/config/.env` 配好对应 provider 的 API Key 和 Base URL。

`.env` 中会读取的变量包括：

```shell
OPENROUTER_API_KEY=
OPENROUTER_BASE_URL=
ZENMUX_API_KEY=
ZENMUX_BASE_URL_OPENAI=
ZENMUX_BASE_URL_ANTHROPIC=
AI302_API_KEY=
AI302_BASE_URL=
DEEPSEEK_API_KEY=
DEEPSEEK_BASE_URL=
MOONSHOT_API_KEY=
MOONSHOT_BASE_URL=
ARK_API_KEY=
ARK_BASE_URL=
QIANFAN_API_KEY=
QIANFAN_BASE_URL=
DASHSCOPE_API_KEY=
DASHSCOPE_BASE_URL=
GOOGLE_API_KEY=
GOOGLE_BASE_URL=
MINIMAX_API_KEY=
MINIMAX_BASE_URL_ANTHROPIC=
```

部分图片和 PDF 测试依赖外部 URL；文件输入测试默认读取本地文件：

- `/Users/reynoldqin/Downloads/1.png`
- `/Users/reynoldqin/Downloads/planning-with-files.pdf`

```shell
# 收集当前模型集成测试
pytest --collect-only -q tests/model_tests.py

# 运行全部模型集成测试
pytest -s --log-cli-level=INFO tests/model_tests.py

# 单独运行某个模型
pytest -s --log-cli-level=INFO tests/model_tests.py::test_gpt_5_2
pytest -s --log-cli-level=INFO tests/model_tests.py::test_claude_4_6_sonnet_zenmux
pytest -s --log-cli-level=INFO tests/model_tests.py::test_qwen_plus
```

当前覆盖的测试场景包括：

- 工具调用：流式和非流式
- 推理模式：OpenAI 兼容接口、Claude、Qwen
- 结构化输出：Pydantic JSON Schema 和 JSON Object 模式
- 图片输入：URL 和本地文件
- PDF 文档输入：URL 和本地文件

当前模型用例：

- `test_gpt_5_2`
- `test_gemini_3_pro`
- `test_kimi_k2_thinking`
- `test_kimi_k2_5`
- `test_deepseek_reasoner`
- `test_doubao_seed_2_pro`
- `test_doubao_seed_2_lite`
- `test_doubao_seed_2_1_pro`
- `test_doubao_seed_1_8`
- `test_ernie_x_1_1`
- `test_grok_4_1_fast`
- `test_qwen_plus`
- `test_claude_4_5_sonnet_zenmux`
- `test_claude_4_6_sonnet_zenmux`
- `test_minimax_m2_5`
