mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
feat: add embedding and rerank support
This commit is contained in:
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -8,6 +8,7 @@ pub enum LocalStandardSourceFamily {
|
||||
pub enum LocalStandardSourceMode {
|
||||
Chat,
|
||||
Cli,
|
||||
Embedding,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
|
||||
Reference in New Issue
Block a user