feat: add embedding and rerank support

This commit is contained in:
Kayphoon
2026-05-03 17:32:41 +08:00
parent 3e2eca4fd0
commit 5abe664d65
87 changed files with 5520 additions and 184 deletions

View File

@@ -1,7 +1,8 @@
use crate::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -49,6 +50,20 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec>
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: false,
}),
OPENAI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "openai:embedding",
decision_kind: OPENAI_EMBEDDING_SYNC_PLAN_KIND,
report_kind: "openai_embedding_sync_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
OPENAI_RERANK_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "openai:rerank",
decision_kind: OPENAI_RERANK_SYNC_PLAN_KIND,
report_kind: "openai_rerank_sync_success",
family: LocalSameFormatProviderFamily::Standard,
require_streaming: false,
}),
_ => None,
}
}
@@ -106,4 +121,20 @@ mod tests {
assert_eq!(spec.report_kind, "gemini_cli_stream_success");
assert!(spec.require_streaming);
}
#[test]
fn resolves_openai_embedding_sync_same_format_spec() {
let spec = resolve_sync_spec("openai_embedding_sync").expect("spec");
assert_eq!(spec.api_format, "openai:embedding");
assert_eq!(spec.report_kind, "openai_embedding_sync_success");
assert!(!spec.require_streaming);
}
#[test]
fn resolves_openai_rerank_sync_same_format_spec() {
let spec = resolve_sync_spec("openai_rerank_sync").expect("spec");
assert_eq!(spec.api_format, "openai:rerank");
assert_eq!(spec.report_kind, "openai_rerank_sync_success");
assert!(!spec.require_streaming);
}
}

View File

@@ -7,12 +7,12 @@ use crate::contracts::{
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
use crate::request::specialized::image::is_openai_image_stream_request;
@@ -173,6 +173,22 @@ pub fn resolve_execution_runtime_sync_plan_kind(
return Some(OPENAI_CHAT_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("embedding")
&& *method == Method::POST
&& path == "/v1/embeddings"
{
return Some(OPENAI_EMBEDDING_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("rerank")
&& *method == Method::POST
&& path == "/v1/rerank"
{
return Some(OPENAI_RERANK_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("image")
&& *method == Method::POST
@@ -327,6 +343,8 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool {
matches!(
plan_kind,
OPENAI_CHAT_SYNC_PLAN_KIND
| OPENAI_EMBEDDING_SYNC_PLAN_KIND
| OPENAI_RERANK_SYNC_PLAN_KIND
| OPENAI_IMAGE_SYNC_PLAN_KIND
| OPENAI_RESPONSES_SYNC_PLAN_KIND
| OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND
@@ -377,7 +395,8 @@ mod tests {
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND,
OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
};
@@ -650,6 +669,40 @@ mod tests {
));
}
#[test]
fn resolves_openai_embedding_sync_plan_kind() {
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("openai"),
Some("embedding"),
&Method::POST,
"/v1/embeddings",
),
Some(OPENAI_EMBEDDING_SYNC_PLAN_KIND)
);
assert!(supports_sync_scheduler_decision_kind(
OPENAI_EMBEDDING_SYNC_PLAN_KIND
));
}
#[test]
fn resolves_openai_rerank_sync_plan_kind() {
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("openai"),
Some("rerank"),
&Method::POST,
"/v1/rerank",
),
Some(OPENAI_RERANK_SYNC_PLAN_KIND)
);
assert!(supports_sync_scheduler_decision_kind(
OPENAI_RERANK_SYNC_PLAN_KIND
));
}
#[test]
fn resolves_openai_image_stream_plan_kind() {
assert_eq!(

View File

@@ -8,6 +8,7 @@ pub enum LocalStandardSourceFamily {
pub enum LocalStandardSourceMode {
Chat,
Cli,
Embedding,
}
#[derive(Debug, Clone, Copy)]