use std::sync::Arc;
use jammi_ai::session::InferenceSession;
use jammi_ai::Session;
use tonic::{Request, Response, Status};
use crate::grpc::proto::eval as pb;
use crate::grpc::proto::eval::eval_service_server::EvalService;
use crate::grpc::wire::{map_engine_error, scoped, session_tenant_traced};
pub struct EvalServer {
session: Arc<InferenceSession>,
}
impl EvalServer {
pub fn new(session: Arc<InferenceSession>) -> Self {
Self { session }
}
fn local(&self) -> Session {
Session::new(Arc::clone(&self.session))
}
}
#[tonic::async_trait]
impl EvalService for EvalServer {
#[tracing::instrument(skip(self, request), fields(tenant_id = tracing::field::Empty))]
async fn eval_embeddings(
&self,
request: Request<pb::EvalEmbeddingsRequest>,
) -> Result<Response<pb::EmbeddingEvalReport>, Status> {
let tenant = session_tenant_traced(&request);
let args = jammi_ai::wire::eval_embeddings_from_proto(request.into_inner())?;
let session = self.local();
let report = scoped(&self.session, tenant, || {
session.eval_embeddings(
&args.source_id,
args.embedding_table.as_deref(),
&args.golden_source,
args.k,
&args.cohorts,
)
})
.await
.map_err(map_engine_error)?;
Ok(Response::new(report.into()))
}
#[tracing::instrument(skip(self, request), fields(tenant_id = tracing::field::Empty))]
async fn eval_per_query(
&self,
request: Request<pb::EvalPerQueryRequest>,
) -> Result<Response<pb::EvalPerQueryResponse>, Status> {
let tenant = session_tenant_traced(&request);
let eval_run_id = jammi_ai::wire::eval_per_query_from_proto(request.into_inner())?;
let session = self.local();
let records = scoped(&self.session, tenant, || {
session.eval_per_query(&eval_run_id)
})
.await
.map_err(map_engine_error)?;
Ok(Response::new(pb::EvalPerQueryResponse {
records: records.into_iter().map(Into::into).collect(),
}))
}
#[tracing::instrument(skip(self, request), fields(tenant_id = tracing::field::Empty))]
async fn eval_inference(
&self,
request: Request<pb::EvalInferenceRequest>,
) -> Result<Response<pb::InferenceEvalReport>, Status> {
let tenant = session_tenant_traced(&request);
let args = jammi_ai::wire::eval_inference_from_proto(request.into_inner())?;
let session = self.local();
let report = scoped(&self.session, tenant, || {
session.eval_inference(
&args.model_id,
&args.source_id,
&args.columns,
args.task,
&args.golden_source,
&args.label_column,
)
})
.await
.map_err(map_engine_error)?;
Ok(Response::new(report.into()))
}
#[tracing::instrument(skip(self, request), fields(tenant_id = tracing::field::Empty))]
async fn eval_compare(
&self,
request: Request<pb::EvalCompareRequest>,
) -> Result<Response<pb::CompareEvalReport>, Status> {
let tenant = session_tenant_traced(&request);
let args = jammi_ai::wire::eval_compare_from_proto(request.into_inner())?;
let session = self.local();
let report = scoped(&self.session, tenant, || {
session.eval_compare(
&args.embedding_tables,
&args.source_id,
&args.golden_source,
args.k,
)
})
.await
.map_err(map_engine_error)?;
Ok(Response::new(report.into()))
}
#[tracing::instrument(skip(self, request), fields(tenant_id = tracing::field::Empty))]
async fn eval_calibration(
&self,
request: Request<pb::EvalCalibrationRequest>,
) -> Result<Response<pb::CalibrationEvalReport>, Status> {
let tenant = session_tenant_traced(&request);
let args = jammi_ai::wire::eval_calibration_from_proto(request.into_inner())?;
let report = scoped(&self.session, tenant, || async {
self.session
.eval_calibration(
&args.source_id,
&args.golden_source,
args.shape,
&args.cohorts,
)
.await
})
.await
.map_err(map_engine_error)?;
Ok(Response::new(report.into()))
}
}