lmkit 0.1.0

Multi-provider AI API client (OpenAI, Anthropic, Google Gemini, Aliyun, Ollama, Zhipu; chat, embed incl. Gemini, rerank, image, audio stubs)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! 文本向量嵌入：非流式 JSON，默认 HTTP 超时约 30 秒（可用 [`ProviderConfig::timeout`] 覆盖）。
//!
//! # HTTP 与厂商分支
//!
//! 多数厂商请求地址为 `POST {base_url}/embeddings`，`base_url` 会先 `trim_end_matches('/')` 再拼接路径。**Google Gemini** 见下文专用路径。
//!
//! **`OpenAI` / `Aliyun` / `Ollama`**（启用对应厂商 feature 与 `embed`）：OpenAI 兼容请求体，含 `model`、`input`（字符串数组）、**`dimensions`**（等于配置中的 [`ProviderConfig::dimension`]，序列化进 JSON）。成功时解析 `data[].embedding`。
//!
//! **`Zhipu`**：路径仍为 `…/embeddings`，请求体仅 `model` 与 `input`，**不发送 `dimensions` 字段**；配置中的 `dimension` 仍必填，用于 [`EmbedProvider::dimension`] 返回值，且须与模型实际输出维数一致。未启用 `zhipu` feature 时选择智谱会得到 [`Error::ProviderDisabled`]。
//!
//! **`Google`**（`google` + `embed`）：**Gemini** `embedContent` / `batchEmbedContents`，query 参数 **`key`**；请求体字段与官方 REST 一致，见源码 `google_gemini.rs`。未启用 `google` feature 时选 Google 为 [`Error::ProviderDisabled`]。
//!
//! **`Anthropic`**：工厂返回 [`Error::Unsupported`]（`capability` 为 `"embed"`）；未启用 `anthropic` feature 时选该厂商为 [`Error::ProviderDisabled`]。
//!
//! # 文本预处理
//!
//! [`EmbedProvider::encode`] 与 [`EmbedProvider::encode_batch`] 在组请求前会对每条文本做首尾空白裁剪与连续空白折叠。
//!
//! # 鉴权
//!
//! 已实现嵌入的厂商：`OpenAI` / `Aliyun` / `Ollama` / `Zhipu` 使用 `Authorization: Bearer {api_key}`；**`Google`** 使用 query `key`（见上）。**`Anthropic`** 无嵌入实现，工厂阶段即返回错误，不发起 HTTP。

#[cfg(feature = "google")]
mod google_gemini;
mod openai_compat;
#[cfg(feature = "zhipu")]
mod zhipu;

use async_trait::async_trait;
use std::time::Duration;

use crate::client::HttpClient;
use crate::config::Provider;
use crate::config::ProviderConfig;
use crate::error::{Error, Result};

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);

#[async_trait]
pub trait EmbedProvider: Send + Sync {
    /// 单条文本嵌入；内部先规范化空白再调用 [`encode_batch`](Self::encode_batch)。
    async fn encode(&self, text: &str) -> Result<Vec<f32>>;
    /// 批量嵌入；顺序与 `texts` 一致，长度与 [`dimension`](Self::dimension) 由配置与上游决定。
    async fn encode_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
    /// 配置维度（智谱请求体虽不含 `dimensions` 字段，仍须与模型输出一致）。
    fn dimension(&self) -> usize;
}

fn http_client(config: &ProviderConfig) -> Result<HttpClient> {
    HttpClient::new(config.timeout.unwrap_or(DEFAULT_TIMEOUT))
}

pub(crate) fn create(config: &ProviderConfig) -> Result<Box<dyn EmbedProvider>> {
    let dimension = config.dimension.ok_or(Error::MissingConfig("dimension"))?;

    match config.provider {
        #[cfg(feature = "openai")]
        Provider::OpenAI => Ok(Box::new(openai_compat::OpenaiCompatEmbed::new(
            config,
            dimension,
            http_client(config)?,
        ))),
        #[cfg(not(feature = "openai"))]
        Provider::OpenAI => Err(Error::ProviderDisabled("openai".to_string())),

        #[cfg(feature = "aliyun")]
        Provider::Aliyun => Ok(Box::new(openai_compat::OpenaiCompatEmbed::new(
            config,
            dimension,
            http_client(config)?,
        ))),
        #[cfg(not(feature = "aliyun"))]
        Provider::Aliyun => Err(Error::ProviderDisabled("aliyun".to_string())),

        #[cfg(feature = "ollama")]
        Provider::Ollama => Ok(Box::new(openai_compat::OpenaiCompatEmbed::new(
            config,
            dimension,
            http_client(config)?,
        ))),
        #[cfg(not(feature = "ollama"))]
        Provider::Ollama => Err(Error::ProviderDisabled("ollama".to_string())),

        #[cfg(feature = "zhipu")]
        Provider::Zhipu => Ok(Box::new(zhipu::ZhipuEmbed::new(
            config,
            dimension,
            http_client(config)?,
        ))),
        #[cfg(not(feature = "zhipu"))]
        Provider::Zhipu => Err(Error::ProviderDisabled("zhipu".to_string())),

        #[cfg(feature = "anthropic")]
        Provider::Anthropic => Err(Error::Unsupported {
            provider: config.provider.to_string(),
            capability: "embed",
        }),
        #[cfg(not(feature = "anthropic"))]
        Provider::Anthropic => Err(Error::ProviderDisabled("anthropic".to_string())),

        #[cfg(feature = "google")]
        Provider::Google => Ok(Box::new(google_gemini::GoogleGeminiEmbed::new(
            config,
            dimension,
            http_client(config)?,
        ))),
        #[cfg(not(feature = "google"))]
        Provider::Google => Err(Error::ProviderDisabled("google".to_string())),
    }
}