use super::*;
#[cfg(feature = "embed-api")]
pub(crate) async fn query_with_mode(
upstream: &GenericRestAdapter,
request: &QueryRequest,
smart_embedder: &Option<Arc<SmartEmbedder>>,
) -> Result<QueryResponse, super::super::upstream::UpstreamError> {
let mode = upstream.query_mode();
match mode {
QueryMode::TextNative => {
upstream.query(request).await
}
QueryMode::VectorOnly => {
let embedder = smart_embedder.as_ref().ok_or_else(|| {
super::super::upstream::UpstreamError::EmbeddingRequired(
"VectorOnly upstream requires proxy-embed feature with embedder configured"
.to_string(),
)
})?;
let vector = embedder.embed(&request.query).await.map_err(
|e: crate::error::ConproxyError| {
super::super::upstream::UpstreamError::EmbeddingFailed(e.to_string())
},
)?;
upstream.query_vector(request, &vector).await
}
QueryMode::Unknown => {
match upstream.query(request).await {
Ok(response) => {
upstream.set_query_mode(QueryMode::TextNative);
Ok(response)
}
Err(super::super::upstream::UpstreamError::UnsupportedQueryType(_)) => {
if let Ok(discovered) = upstream.discover_query_mode().await {
upstream.set_query_mode(discovered);
if discovered == QueryMode::VectorOnly {
let embedder = smart_embedder.as_ref().ok_or_else(|| {
super::super::upstream::UpstreamError::EmbeddingRequired(
"VectorOnly upstream requires embedding support".to_string(),
)
})?;
let vector = embedder.embed(&request.query).await.map_err(
|e: crate::error::ConproxyError| {
super::super::upstream::UpstreamError::EmbeddingFailed(
e.to_string(),
)
},
)?;
return upstream.query_vector(request, &vector).await;
}
}
Err(super::super::upstream::UpstreamError::UnsupportedQueryType(
"Could not determine upstream query mode".to_string(),
))
}
Err(e) => Err(e),
}
}
}
}
#[instrument(skip(state, headers, agent_identity), fields(query_len = request.query.len(), top_k = request.top_k))]
pub(super) async fn handle_query(
State(state): State<AppState>,
headers: axum::http::HeaderMap,
agent_identity: Option<axum::Extension<AgentIdentity>>,
Json(request): Json<QueryRequest>,
) -> impl IntoResponse {
let request_id = extract_request_id(&headers);
let agent_id = agent_identity.as_ref().map(|a| a.id.clone());
debug!(request_id = %request_id, agent = ?agent_id, "Processing query request");
let ctx_id = resolve_context(&headers, agent_identity.as_deref());
let source = agent_id.unwrap_or_else(|| {
headers
.get("x-forwarded-for")
.and_then(|v| v.to_str().ok())
.unwrap_or("local")
.to_string()
});
let result = super::query_core::execute_query(
&state,
request,
ctx_id,
request_id,
agent_identity.as_deref(),
source,
false,
)
.await;
let status_code = match result.status {
200 => StatusCode::OK,
400 => StatusCode::BAD_REQUEST,
403 => StatusCode::FORBIDDEN,
502 => StatusCode::BAD_GATEWAY,
503 => StatusCode::SERVICE_UNAVAILABLE,
_ => StatusCode::INTERNAL_SERVER_ERROR,
};
(status_code, Json(result.response))
}