evals on query engine; more datasets apis (#1021)

* evals on query engine; more datasets apis

* fix build, address dataset id comment

* minor fix in types

Author: dinmukhamedm

app-server/src/api/v1/datasets.rs

@@ -19,9 +19,33 @@ use crate::{
 
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
-pub struct GetDatapointsRequestParams {
-    #[serde(alias = "dataset_name", alias = "name")]
-    dataset_name: String,
+struct GetDatasetsRequest {
+    #[serde(default)]
+    id: Option<Uuid>,
+    #[serde(default)]
+    name: Option<String>,
+}
+
+#[get("/datasets")]
+async fn get_datasets(
+    db: web::Data<DB>,
+    project_api_key: ProjectApiKey,
+    req: web::Query<GetDatasetsRequest>,
+) -> ResponseResult {
+    let project_id = project_api_key.project_id;
+    let db = db.into_inner();
+    let request = req.into_inner();
+    let datasets =
+        db::datasets::get_datasets(&db.pool, project_id, request.id, request.name).await?;
+
+    Ok(HttpResponse::Ok().json(datasets))
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct GetDatapointsRequestParams {
+    #[serde(flatten)]
+    dataset: DatasetIdentifier,
     limit: i64,
     offset: i64,
 }
@@ -46,13 +70,19 @@ async fn get_datapoints(
     let query_engine = query_engine.into_inner().as_ref().clone();
     let query = params.into_inner();
 
-    let dataset_id =
-        db::datasets::get_dataset_id_by_name(&db.pool, &query.dataset_name, project_id).await?;
-
-    let Some(dataset_id) = dataset_id else {
-        return Ok(HttpResponse::NotFound().json(serde_json::json!({
-            "error": "Dataset not found"
-        })));
+    let dataset_id = match query.dataset {
+        DatasetIdentifier::Name(name) => {
+            let Some(dataset_id) =
+                db::datasets::get_dataset_id_by_name(&db.pool, &name.dataset_name, project_id)
+                    .await?
+            else {
+                return Ok(HttpResponse::NotFound().json(serde_json::json!({
+                    "error": "Dataset not found"
+                })));
+            };
+            dataset_id
+        }
+        DatasetIdentifier::Id(id) => id.dataset_id,
     };
 
     let select_query = "
@@ -129,28 +159,48 @@ async fn get_datapoints(
 
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
-pub struct CreateDatapointsRequest {
-    // The alias is added to support the old endpoint (dataset_name)
-    #[serde(alias = "dataset_name")]
+pub struct DatasetName {
+    #[serde(alias = "dataset_name", alias = "name")]
     pub dataset_name: String,
-    pub datapoints: Vec<CreateDatapointRequest>,
 }
 
 #[derive(Deserialize)]
-pub struct CreateDatapointRequest {
+#[serde(rename_all = "camelCase")]
+pub struct DatasetId {
+    #[serde(alias = "dataset_id")]
+    pub dataset_id: Uuid,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+#[serde(untagged)]
+enum DatasetIdentifier {
+    Name(DatasetName),
+    Id(DatasetId),
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct CreateDatapointsRequest {
+    #[serde(flatten)]
+    dataset: DatasetIdentifier,
+    datapoints: Vec<RequestDatapoint>,
+    #[serde(default)]
+    create_dataset: bool,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct RequestDatapoint {
     #[serde(default)]
-    pub id: Option<Uuid>,
-    pub data: serde_json::Value,
-    pub target: Option<serde_json::Value>,
+    id: Option<Uuid>,
+    data: serde_json::Value,
+    target: Option<serde_json::Value>,
     #[serde(default)]
-    pub metadata: std::collections::HashMap<String, serde_json::Value>,
+    metadata: std::collections::HashMap<String, serde_json::Value>,
 }
 
 /// Create datapoints in a dataset
-///
-/// Request body should contain:
-/// - dataset_name: The name of the dataset to add datapoints to
-/// - datapoints: Array of datapoint objects with data, optional target, and optional metadata
 #[post("/datasets/datapoints")]
 async fn create_datapoints(
     req: web::Json<CreateDatapointsRequest>,
@@ -162,6 +212,7 @@ async fn create_datapoints(
     let db = db.into_inner();
     let clickhouse = clickhouse.into_inner().as_ref().clone();
     let request = req.into_inner();
+    let mut created = false;
 
     // Validate that we have datapoints to insert
     if request.datapoints.is_empty() {
@@ -170,21 +221,55 @@ async fn create_datapoints(
         })));
     }
 
-    let dataset_id =
-        db::datasets::get_dataset_id_by_name(&db.pool, &request.dataset_name, project_id).await?;
-
-    let Some(dataset_id) = dataset_id else {
-        return Ok(HttpResponse::NotFound().json(serde_json::json!({
-            "error": "Dataset not found"
-        })));
+    let dataset_id = match request.dataset {
+        DatasetIdentifier::Name(name) => {
+            match db::datasets::get_dataset_id_by_name(&db.pool, &name.dataset_name, project_id)
+                .await?
+            {
+                Some(dataset_id) => {
+                    if request.create_dataset {
+                        return Ok(HttpResponse::Conflict().json(serde_json::json!({
+                            "error": "Dataset with this name already exists"
+                        })));
+                    }
+                    dataset_id
+                }
+                None => {
+                    if request.create_dataset {
+                        let dataset =
+                            db::datasets::create_dataset(&db.pool, &name.dataset_name, project_id)
+                                .await?;
+                        created = true;
+                        dataset.id
+                    } else {
+                        return Ok(HttpResponse::NotFound().json(serde_json::json!({
+                            "error": "Dataset not found"
+                        })));
+                    }
+                }
+            }
+        }
+        DatasetIdentifier::Id(id) => {
+            if request.create_dataset {
+                return Ok(HttpResponse::BadRequest().json(serde_json::json!({
+                    "error": "When creating a new dataset, the name must be provided"
+                })));
+            }
+            if !db::datasets::dataset_exists(&db.pool, id.dataset_id, project_id).await? {
+                return Ok(HttpResponse::NotFound().json(serde_json::json!({
+                    "error": "Dataset not found"
+                })));
+            }
+            id.dataset_id
+        }
     };
 
     // Convert request datapoints to Datapoint structs
     let datapoints: Vec<Datapoint> = request
         .datapoints
         .into_iter()
         .map(|dp_req| Datapoint {
-            // now_v7 is guaranteed to be sorted by creation time
+            // `now_v7` is guaranteed to be sorted by creation time
             id: dp_req.id.unwrap_or(Uuid::now_v7()),
             created_at: Utc::now(),
             dataset_id,
@@ -201,9 +286,26 @@ async fn create_datapoints(
 
     ch_datapoints::insert_datapoints(clickhouse, ch_datapoints).await?;
 
-    Ok(HttpResponse::Created().json(serde_json::json!({
+    let datapoint_info = datapoints
+        .iter()
+        .map(|dp| {
+            serde_json::json!({
+                "id": dp.id,
+                "createdAt": dp.created_at,
+            })
+        })
+        .collect::<Vec<_>>();
+
+    let mut response = if created {
+        HttpResponse::Created()
+    } else {
+        HttpResponse::Ok()
+    };
+    Ok(response.json(serde_json::json!({
         "message": "Datapoints created successfully",
-        "count": datapoints.len()
+        "datasetId": dataset_id,
+        "count": datapoints.len(),
+        "datapointInfo": datapoint_info,
     })))
 }
 

app-server/src/ch/evaluation_datapoints.rs

@@ -31,6 +31,11 @@ pub struct CHEvaluationDatapoint {
     pub data: String,
     pub target: String,
     pub metadata: String,
+    #[serde(with = "clickhouse::serde::uuid")]
+    pub dataset_id: Uuid,
+    #[serde(with = "clickhouse::serde::uuid")]
+    pub dataset_datapoint_id: Uuid,
+    pub dataset_datapoint_created_at: i64,
 }
 
 #[derive(Row, Deserialize)]
@@ -57,6 +62,22 @@ impl CHEvaluationDatapoint {
             metadata: json_value_to_string(
                 &serde_json::to_value(result.metadata.unwrap_or_default()).unwrap_or_default(),
             ),
+            dataset_id: result
+                .dataset_link
+                .as_ref()
+                .map(|link| link.dataset_id)
+                .unwrap_or_default(),
+            dataset_datapoint_id: result
+                .dataset_link
+                .as_ref()
+                .map(|link| link.datapoint_id)
+                .unwrap_or_default(),
+            dataset_datapoint_created_at: chrono_to_nanoseconds(
+                result
+                    .dataset_link
+                    .map(|link| link.created_at)
+                    .unwrap_or_default(),
+            ),
         }
     }
 }

app-server/src/db/datasets.rs

@@ -1,7 +1,60 @@
 use anyhow::Result;
-use sqlx::PgPool;
+use chrono::{DateTime, Utc};
+use serde::Serialize;
+use sqlx::{FromRow, PgPool, QueryBuilder};
 use uuid::Uuid;
 
+#[derive(Serialize, FromRow)]
+#[serde(rename_all = "camelCase")]
+pub struct Dataset {
+    pub id: Uuid,
+    pub name: String,
+    pub project_id: Uuid,
+    pub created_at: DateTime<Utc>,
+}
+
+/// Get datasets by project id. If ID or name is provided, filter the results accordingly.
+pub async fn get_datasets(
+    pool: &PgPool,
+    project_id: Uuid,
+    id: Option<Uuid>,
+    name: Option<String>,
+) -> Result<Vec<Dataset>> {
+    let mut query_builder = QueryBuilder::new(
+        r#"SELECT id, name, project_id, created_at FROM datasets WHERE project_id = "#,
+    );
+
+    query_builder.push_bind(project_id);
+    if let Some(id) = id {
+        query_builder.push(" AND id = ");
+        query_builder.push_bind(id);
+    }
+    if let Some(name) = name {
+        query_builder.push(" AND name = ");
+        query_builder.push_bind(name);
+    }
+    query_builder.push(" ORDER BY created_at DESC");
+    let datasets = query_builder
+        .build_query_as::<Dataset>()
+        .fetch_all(pool)
+        .await?;
+
+    Ok(datasets)
+}
+
+pub async fn create_dataset(pool: &PgPool, name: &str, project_id: Uuid) -> Result<Dataset> {
+    let dataset = sqlx::query_as::<_, Dataset>(
+        "INSERT INTO datasets (name, project_id) VALUES ($1, $2)
+        RETURNING id, name, project_id, created_at",
+    )
+    .bind(name)
+    .bind(project_id)
+    .fetch_one(pool)
+    .await?;
+
+    Ok(dataset)
+}
+
 pub async fn get_dataset_id_by_name(
     pool: &PgPool,
     name: &str,
@@ -18,6 +71,18 @@ pub async fn get_dataset_id_by_name(
     Ok(dataset_id)
 }
 
+pub async fn dataset_exists(pool: &PgPool, dataset_id: Uuid, project_id: Uuid) -> Result<bool> {
+    let dataset_exists = sqlx::query_scalar::<_, bool>(
+        "SELECT EXISTS(SELECT 1 FROM datasets WHERE id = $1 AND project_id = $2)",
+    )
+    .bind(dataset_id)
+    .bind(project_id)
+    .fetch_one(pool)
+    .await?;
+
+    Ok(dataset_exists)
+}
+
 pub async fn get_parquet_path(
     pool: &PgPool,
     project_id: Uuid,

app-server/src/evaluations/utils.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use chrono::{DateTime, Utc};
 use serde::Deserialize;
 use serde_json::Value;
 use uuid::Uuid;
@@ -8,6 +9,14 @@ use crate::utils::json_value_to_string;
 
 const MAX_JSON_VALUE_LENGTH: usize = 1000;
 
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct EvaluationDatapointDatasetLink {
+    pub dataset_id: Uuid,
+    pub datapoint_id: Uuid,
+    pub created_at: DateTime<Utc>,
+}
+
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct EvaluationDatapointResult {
@@ -26,6 +35,8 @@ pub struct EvaluationDatapointResult {
     pub trace_id: Uuid,
     #[serde(default)]
     pub scores: HashMap<String, Option<f64>>,
+    #[serde(default)]
+    pub dataset_link: Option<EvaluationDatapointDatasetLink>,
 }
 
 pub struct DatapointColumns {

app-server/src/main.rs

@@ -797,6 +797,7 @@ fn main() -> anyhow::Result<()> {
                             .service(
                                 web::scope("/v1")
                                     .wrap(project_auth.clone())
+                                    .service(api::v1::datasets::get_datasets)
                                     .service(api::v1::datasets::get_datapoints)
                                     .service(api::v1::datasets::create_datapoints)
                                     .service(api::v1::datasets::get_parquet)