diff --git a/clients/ts-sdk/openapi.json b/clients/ts-sdk/openapi.json index 3e87dfe0ae..1bfcd66a1b 100644 --- a/clients/ts-sdk/openapi.json +++ b/clients/ts-sdk/openapi.json @@ -4935,6 +4935,295 @@ ] } }, + "/api/experiment": { + "get": { + "tags": [ + "Experiment" + ], + "summary": "Get Experiments", + "description": "Get all experiments for a dataset. Auth'ed user must be an owner of the organization to get experiments.", + "operationId": "get_experiments", + "parameters": [ + { + "name": "TR-Dataset", + "in": "header", + "description": "The dataset id to use for the request", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + } + ], + "responses": { + "200": { + "description": "Experiments retrieved successfully", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Experiment" + } + } + } + } + }, + "400": { + "description": "Service error relating to getting the experiments", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponseBody" + } + } + } + } + }, + "security": [ + { + "ApiKey": [ + "owner" + ] + } + ] + }, + "post": { + "tags": [ + "Experiment" + ], + "summary": "Create Experiment", + "description": "Experiment will be created in the dataset specified via the TR-Dataset header. Auth'ed user must be an owner of the organization to create an experiment.", + "operationId": "create_experiment", + "parameters": [ + { + "name": "TR-Dataset", + "in": "header", + "description": "The dataset id to use for the request", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + } + ], + "requestBody": { + "description": "JSON request payload to create a new experiment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CreateExperimentReqBody" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Experiment created successfully", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Experiment" + } + } + } + }, + "400": { + "description": "Service error relating to creating the experiment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponseBody" + } + } + } + } + }, + "security": [ + { + "ApiKey": [ + "owner" + ] + } + ] + }, + "put": { + "tags": [ + "Experiment" + ], + "summary": "Update Experiment", + "description": "Update an experiment. Auth'ed user must be an owner of the organization to update an experiment.", + "operationId": "update_experiment", + "parameters": [ + { + "name": "TR-Dataset", + "in": "header", + "description": "The dataset id to use for the request", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + } + ], + "requestBody": { + "description": "JSON request payload to update an experiment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UpdateExperimentReqBody" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Experiment updated successfully", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/Experiment" + } + } + } + }, + "400": { + "description": "Service error relating to updating the experiment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponseBody" + } + } + } + } + }, + "security": [ + { + "ApiKey": [ + "owner" + ] + } + ] + } + }, + "/api/experiment/ab-test": { + "post": { + "tags": [ + "Experiment" + ], + "summary": "Ab Test", + "description": "Get a user's treatment for an experiment. Auth'ed user must be an owner of the organization to get a user's treatment.", + "operationId": "ab_test", + "parameters": [ + { + "name": "TR-Dataset", + "in": "header", + "description": "The dataset id to use for the request", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + } + ], + "requestBody": { + "description": "JSON request payload to get a user's treatment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AbTestReqBody" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "User treatment response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/UserTreatmentResponse" + } + } + } + }, + "400": { + "description": "Service error relating to getting the user's treatment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponseBody" + } + } + } + } + }, + "security": [ + { + "ApiKey": [ + "owner" + ] + } + ] + } + }, + "/api/experiment/{experiment_id}": { + "delete": { + "tags": [ + "Experiment" + ], + "summary": "Delete Experiment", + "description": "Delete an experiment. Auth'ed user must be an owner of the organization to delete an experiment.", + "operationId": "delete_experiment", + "parameters": [ + { + "name": "TR-Dataset", + "in": "header", + "description": "The dataset id to use for the request", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + }, + { + "name": "experiment_id", + "in": "path", + "required": true, + "schema": { + "type": "string", + "format": "uuid" + } + } + ], + "responses": { + "204": { + "description": "Experiment deleted successfully" + }, + "400": { + "description": "Service error relating to deleting the experiment", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ErrorResponseBody" + } + } + } + } + }, + "security": [ + { + "ApiKey": [ + "owner" + ] + } + ] + } + }, "/api/file": { "post": { "tags": [ @@ -7586,6 +7875,22 @@ "V2" ] }, + "AbTestReqBody": { + "type": "object", + "required": [ + "experiment_id", + "user_id" + ], + "properties": { + "experiment_id": { + "type": "string", + "format": "uuid" + }, + "user_id": { + "type": "string" + } + } + }, "AddChunkToGroupReqPayload": { "type": "object", "properties": { @@ -10761,6 +11066,21 @@ } } }, + "CreateExperimentReqBody": { + "type": "object", + "required": [ + "name", + "experiment_config" + ], + "properties": { + "experiment_config": { + "$ref": "#/components/schemas/ExperimentConfig" + }, + "name": { + "type": "string" + } + } + }, "CreateFormWithoutFile": { "type": "object", "description": "Will use [chunkr.ai](https://chunkr.ai) to process the file when this object is defined. See [docs.chunkr.ai/api-references/task/create-task](https://docs.chunkr.ai/api-references/task/create-task) for detailed information about what each field on this request payload does.", @@ -12833,6 +13153,80 @@ } } }, + "Experiment": { + "type": "object", + "required": [ + "id", + "name", + "t1_name", + "t1_split", + "control_name", + "control_split", + "dataset_id", + "created_at", + "updated_at" + ], + "properties": { + "control_name": { + "type": "string" + }, + "control_split": { + "type": "number", + "format": "float" + }, + "created_at": { + "type": "string", + "format": "date-time" + }, + "dataset_id": { + "type": "string", + "format": "uuid" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "name": { + "type": "string" + }, + "t1_name": { + "type": "string" + }, + "t1_split": { + "type": "number", + "format": "float" + }, + "updated_at": { + "type": "string", + "format": "date-time" + } + } + }, + "ExperimentConfig": { + "type": "object", + "required": [ + "t1_name", + "t1_split", + "control_name", + "control_split" + ], + "properties": { + "control_name": { + "type": "string" + }, + "control_split": { + "type": "number", + "format": "float" + }, + "t1_name": { + "type": "string" + }, + "t1_split": { + "type": "number", + "format": "float" + } + } + }, "ExtendedOrganizationUsageCount": { "type": "object", "required": [ @@ -22082,6 +22476,30 @@ } } }, + "UpdateExperimentReqBody": { + "type": "object", + "required": [ + "id" + ], + "properties": { + "experiment_config": { + "allOf": [ + { + "$ref": "#/components/schemas/ExperimentConfig" + } + ], + "nullable": true + }, + "id": { + "type": "string", + "format": "uuid" + }, + "name": { + "type": "string", + "nullable": true + } + } + }, "UpdateGroupByTrackingIDReqPayload": { "type": "object", "required": [ @@ -22458,6 +22876,26 @@ "user_id": "e3e3e3e3-e3e3-e3e3-e3e3-e3e3e3e3e3e3" } }, + "UserTreatmentResponse": { + "type": "object", + "required": [ + "treatment_name", + "experiment_id", + "user_id" + ], + "properties": { + "experiment_id": { + "type": "string", + "format": "uuid" + }, + "treatment_name": { + "type": "string" + }, + "user_id": { + "type": "string" + } + } + }, "V1RecommendChunksResponseBody": { "type": "array", "items": { @@ -22587,6 +23025,10 @@ { "name": "Analytics", "description": "Analytics endpoint. Used to get information for search and RAG analytics" + }, + { + "name": "Experiment", + "description": "Experiment endpoint. Used to create and manage experiments" } ] } diff --git a/clients/ts-sdk/src/types.gen.ts b/clients/ts-sdk/src/types.gen.ts index c0e0476ca2..07deafdfdb 100644 --- a/clients/ts-sdk/src/types.gen.ts +++ b/clients/ts-sdk/src/types.gen.ts @@ -2,6 +2,11 @@ export type APIVersion = 'V1' | 'V2'; +export type AbTestReqBody = { + experiment_id: string; + user_id: string; +}; + export type AddChunkToGroupReqPayload = { /** * Id of the chunk to make a member of the group. @@ -970,6 +975,11 @@ export type CreateDatasetReqPayload = { tracking_id?: (string) | null; }; +export type CreateExperimentReqBody = { + experiment_config: ExperimentConfig; + name: string; +}; + /** * Will use [chunkr.ai](https://chunkr.ai) to process the file when this object is defined. See [docs.chunkr.ai/api-references/task/create-task](https://docs.chunkr.ai/api-references/task/create-task) for detailed information about what each field on this request payload does. */ @@ -1908,6 +1918,25 @@ export type EventsForTopicResponse = { events: Array; }; +export type Experiment = { + control_name: string; + control_split: number; + created_at: string; + dataset_id: string; + id: string; + name: string; + t1_name: string; + t1_split: number; + updated_at: string; +}; + +export type ExperimentConfig = { + control_name: string; + control_split: number; + t1_name: string; + t1_split: number; +}; + export type ExtendedOrganizationUsageCount = { bytes_ingested: number; chunk_count: number; @@ -4643,6 +4672,12 @@ export type UpdateDatasetReqPayload = { tracking_id?: (string) | null; }; +export type UpdateExperimentReqBody = { + experiment_config?: ((ExperimentConfig) | null); + id: string; + name?: (string) | null; +}; + export type UpdateGroupByTrackingIDReqPayload = { /** * Description to assign to the chunk_group. Convenience field for you to avoid having to remember what the group is for. If not provided, the description will not be updated. @@ -4796,6 +4831,12 @@ export type UserOrganization = { user_id: string; }; +export type UserTreatmentResponse = { + experiment_id: string; + treatment_name: string; + user_id: string; +}; + export type V1RecommendChunksResponseBody = Array; export type WorkerEvent = { @@ -5857,6 +5898,64 @@ export type CreateEtlJobData = { export type CreateEtlJobResponse = (void); +export type GetExperimentsData = { + /** + * The dataset id to use for the request + */ + trDataset: string; +}; + +export type GetExperimentsResponse = (Array); + +export type CreateExperimentData = { + /** + * JSON request payload to create a new experiment + */ + requestBody: CreateExperimentReqBody; + /** + * The dataset id to use for the request + */ + trDataset: string; +}; + +export type CreateExperimentResponse = (Experiment); + +export type UpdateExperimentData = { + /** + * JSON request payload to update an experiment + */ + requestBody: UpdateExperimentReqBody; + /** + * The dataset id to use for the request + */ + trDataset: string; +}; + +export type UpdateExperimentResponse = (Experiment); + +export type AbTestData = { + /** + * JSON request payload to get a user's treatment + */ + requestBody: AbTestReqBody; + /** + * The dataset id to use for the request + */ + trDataset: string; +}; + +export type AbTestResponse = (UserTreatmentResponse); + +export type DeleteExperimentData = { + experimentId: string; + /** + * The dataset id to use for the request + */ + trDataset: string; +}; + +export type DeleteExperimentResponse = (void); + export type UploadFileHandlerData = { /** * JSON request payload to upload a file @@ -7572,6 +7671,77 @@ export type $OpenApiTs = { }; }; }; + '/api/experiment': { + get: { + req: GetExperimentsData; + res: { + /** + * Experiments retrieved successfully + */ + 200: Array; + /** + * Service error relating to getting the experiments + */ + 400: ErrorResponseBody; + }; + }; + post: { + req: CreateExperimentData; + res: { + /** + * Experiment created successfully + */ + 200: Experiment; + /** + * Service error relating to creating the experiment + */ + 400: ErrorResponseBody; + }; + }; + put: { + req: UpdateExperimentData; + res: { + /** + * Experiment updated successfully + */ + 200: Experiment; + /** + * Service error relating to updating the experiment + */ + 400: ErrorResponseBody; + }; + }; + }; + '/api/experiment/ab-test': { + post: { + req: AbTestData; + res: { + /** + * User treatment response + */ + 200: UserTreatmentResponse; + /** + * Service error relating to getting the user's treatment + */ + 400: ErrorResponseBody; + }; + }; + }; + '/api/experiment/{experiment_id}': { + delete: { + req: DeleteExperimentData; + res: { + /** + * Experiment deleted successfully + */ + 204: void; + /** + * Service error relating to deleting the experiment + */ + 400: ErrorResponseBody; + }; + }; + }; '/api/file': { post: { req: UploadFileHandlerData; diff --git a/frontends/dashboard/src/pages/dataset/PublicPageSettings.tsx b/frontends/dashboard/src/pages/dataset/PublicPageSettings.tsx index 5e50e5c262..6c1fdee2a7 100644 --- a/frontends/dashboard/src/pages/dataset/PublicPageSettings.tsx +++ b/frontends/dashboard/src/pages/dataset/PublicPageSettings.tsx @@ -303,7 +303,7 @@ const PublicPageControls = () => { {loadingDefaultConfig() ? "Loading..." : "Auto Configure"} -
+
Published Url:{" "} {publicUrl()} @@ -359,11 +359,11 @@ const PublicPageControls = () => { options={["light", "dark"]} />
-
+
{(color) => ( diff --git a/server/ch_migrations/1747440134_add_experiments_table/down.sql b/server/ch_migrations/1747440134_add_experiments_table/down.sql new file mode 100644 index 0000000000..b4db80709a --- /dev/null +++ b/server/ch_migrations/1747440134_add_experiments_table/down.sql @@ -0,0 +1,2 @@ +DROP TABLE IF EXISTS experiments; +DROP TABLE IF EXISTS experiment_user_assignments; diff --git a/server/ch_migrations/1747440134_add_experiments_table/up.sql b/server/ch_migrations/1747440134_add_experiments_table/up.sql new file mode 100644 index 0000000000..a9588f9acb --- /dev/null +++ b/server/ch_migrations/1747440134_add_experiments_table/up.sql @@ -0,0 +1,28 @@ +CREATE TABLE IF NOT EXISTS experiments ( + id UUID, + name String, + t1_name String, + t1_split Float32, + control_name String, + control_split Float32, + dataset_id UUID, + created_at DateTime DEFAULT now(), + updated_at DateTime DEFAULT now(), +) +ORDER BY (created_at, id) +PARTITION BY + (dataset_id); + +CREATE TABLE IF NOT EXISTS experiment_user_assignments ( + id UUID, + experiment_id UUID, + user_id String, + dataset_id UUID, + treatment_name String, + created_at DateTime DEFAULT now(), + updated_at DateTime DEFAULT now(), +) +ORDER BY (created_at, id) +PARTITION BY + (experiment_id); + \ No newline at end of file diff --git a/server/src/data/models.rs b/server/src/data/models.rs index 86463acbe8..2f4ec73d0c 100644 --- a/server/src/data/models.rs +++ b/server/src/data/models.rs @@ -8901,6 +8901,92 @@ impl Default for ContextOptions { } } +#[derive(Debug, Serialize, Deserialize, ToSchema, Clone, Row)] +pub struct ExperimentClickhouse { + #[serde(with = "clickhouse::serde::uuid")] + pub id: uuid::Uuid, + pub name: String, + pub t1_name: String, + pub t1_split: f32, + pub control_name: String, + pub control_split: f32, + #[serde(with = "clickhouse::serde::uuid")] + pub dataset_id: uuid::Uuid, + #[serde(with = "clickhouse::serde::time::datetime")] + pub created_at: OffsetDateTime, + #[serde(with = "clickhouse::serde::time::datetime")] + pub updated_at: OffsetDateTime, +} + +#[derive(Debug, Serialize, Deserialize, ToSchema, Clone)] +pub struct Experiment { + pub id: uuid::Uuid, + pub name: String, + pub t1_name: String, + pub t1_split: f32, + pub control_name: String, + pub control_split: f32, + pub dataset_id: uuid::Uuid, + pub created_at: chrono::NaiveDateTime, + pub updated_at: chrono::NaiveDateTime, +} + +impl From for ExperimentClickhouse { + fn from(experiment: Experiment) -> Self { + ExperimentClickhouse { + id: experiment.id, + name: experiment.name, + t1_name: experiment.t1_name, + t1_split: experiment.t1_split, + control_name: experiment.control_name, + control_split: experiment.control_split, + dataset_id: experiment.dataset_id, + created_at: OffsetDateTime::from_unix_timestamp(experiment.created_at.timestamp()) + .unwrap(), + updated_at: OffsetDateTime::from_unix_timestamp(experiment.updated_at.timestamp()) + .unwrap(), + } + } +} + +impl From for Experiment { + fn from(experiment: ExperimentClickhouse) -> Self { + Experiment { + id: experiment.id, + name: experiment.name, + t1_name: experiment.t1_name, + t1_split: experiment.t1_split, + control_name: experiment.control_name, + control_split: experiment.control_split, + dataset_id: experiment.dataset_id, + created_at: chrono::NaiveDateTime::from_timestamp( + experiment.created_at.unix_timestamp(), + 0, + ), + updated_at: chrono::NaiveDateTime::from_timestamp( + experiment.updated_at.unix_timestamp(), + 0, + ), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Row, ToSchema)] +pub struct ExperimentUserAssignment { + #[serde(with = "clickhouse::serde::uuid")] + pub id: uuid::Uuid, + #[serde(with = "clickhouse::serde::uuid")] + pub experiment_id: uuid::Uuid, + pub user_id: String, + #[serde(with = "clickhouse::serde::uuid")] + pub dataset_id: uuid::Uuid, + pub treatment_name: String, + #[serde(with = "clickhouse::serde::time::datetime")] + pub created_at: OffsetDateTime, + #[serde(with = "clickhouse::serde::time::datetime")] + pub updated_at: OffsetDateTime, +} + #[derive(Debug, Serialize, Deserialize, ToSchema, Clone, Default)] /// LLM options to use for the completion. If not specified, this defaults to the dataset's LLM options. pub struct LLMOptions { diff --git a/server/src/handlers/experiment_handler.rs b/server/src/handlers/experiment_handler.rs new file mode 100644 index 0000000000..a4e4b8376e --- /dev/null +++ b/server/src/handlers/experiment_handler.rs @@ -0,0 +1,234 @@ +use actix_web::{web, HttpResponse}; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +use crate::{ + data::models::{DatasetAndOrgWithSubAndPlan, Experiment}, + errors::ServiceError, + operators::experiment_operator::{ + ab_test_query, create_experiment_query, delete_experiment_query, get_experiments_query, + update_experiment_query, + }, +}; + +use super::auth_handler::AdminOnly; + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct ExperimentConfig { + pub t1_name: String, + pub t1_split: f32, + pub control_name: String, + pub control_split: f32, +} + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct CreateExperimentReqBody { + pub name: String, + pub experiment_config: ExperimentConfig, +} + +impl CreateExperimentReqBody { + pub fn to_experiment(&self, dataset_id: uuid::Uuid) -> Experiment { + Experiment { + id: uuid::Uuid::new_v4(), + name: self.name.clone(), + t1_name: self.experiment_config.t1_name.clone(), + t1_split: self.experiment_config.t1_split, + control_name: self.experiment_config.control_name.clone(), + control_split: self.experiment_config.control_split, + dataset_id, + created_at: chrono::Utc::now().naive_utc(), + updated_at: chrono::Utc::now().naive_utc(), + } + } +} + +/// Create Experiment +/// +/// Experiment will be created in the dataset specified via the TR-Dataset header. Auth'ed user must be an owner of the organization to create an experiment. +#[utoipa::path( + post, + path = "/experiment", + context_path = "/api", + tag = "Experiment", + request_body(content = CreateExperimentReqBody, description = "JSON request payload to create a new experiment", content_type = "application/json"), + responses( + (status = 200, description = "Experiment created successfully", body = Experiment), + (status = 400, description = "Service error relating to creating the experiment", body = ErrorResponseBody), + ), + params( + ("TR-Dataset" = uuid::Uuid, Header, description = "The dataset id to use for the request"), + ), + security( + ("ApiKey" = ["owner"]), + ) +)] +pub async fn create_experiment( + data: web::Json, + _user: AdminOnly, + clickhouse_client: web::Data, + dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan, +) -> Result { + let experiment = create_experiment_query( + data.into_inner(), + dataset_org_plan_sub.dataset.id, + clickhouse_client.get_ref(), + ) + .await?; + Ok(HttpResponse::Ok().json(experiment)) +} + +/// Get Experiments +/// +/// Get all experiments for a dataset. Auth'ed user must be an owner of the organization to get experiments. +#[utoipa::path( + get, + path = "/experiment", + context_path = "/api", + tag = "Experiment", + responses( + (status = 200, description = "Experiments retrieved successfully", body = Vec), + (status = 400, description = "Service error relating to getting the experiments", body = ErrorResponseBody), + ), + params( + ("TR-Dataset" = uuid::Uuid, Header, description = "The dataset id to use for the request"), + ), + security( + ("ApiKey" = ["owner"]), + ) +)] +pub async fn get_experiments( + _user: AdminOnly, + clickhouse_client: web::Data, + dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan, +) -> Result { + let experiments = + get_experiments_query(dataset_org_plan_sub.dataset.id, clickhouse_client.get_ref()).await?; + Ok(HttpResponse::Ok().json(experiments)) +} + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct UpdateExperimentReqBody { + pub id: uuid::Uuid, + pub name: Option, + pub experiment_config: Option, +} + +/// Update Experiment +/// +/// Update an experiment. Auth'ed user must be an owner of the organization to update an experiment. +#[utoipa::path( + put, + path = "/experiment", + context_path = "/api", + tag = "Experiment", + request_body(content = UpdateExperimentReqBody, description = "JSON request payload to update an experiment", content_type = "application/json"), + responses( + (status = 200, description = "Experiment updated successfully", body = Experiment), + (status = 400, description = "Service error relating to updating the experiment", body = ErrorResponseBody), + ), + params( + ("TR-Dataset" = uuid::Uuid, Header, description = "The dataset id to use for the request"), + ), + security( + ("ApiKey" = ["owner"]), + ) +)] +pub async fn update_experiment( + data: web::Json, + _user: AdminOnly, + clickhouse_client: web::Data, + dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan, +) -> Result { + let experiment = update_experiment_query( + data.into_inner(), + dataset_org_plan_sub.dataset.id, + clickhouse_client.get_ref(), + ) + .await?; + Ok(HttpResponse::Ok().json(experiment)) +} + +/// Delete Experiment +/// +/// Delete an experiment. Auth'ed user must be an owner of the organization to delete an experiment. +#[utoipa::path( + delete, + path = "/experiment/{experiment_id}", + context_path = "/api", + tag = "Experiment", + responses( + (status = 204, description = "Experiment deleted successfully"), + (status = 400, description = "Service error relating to deleting the experiment", body = ErrorResponseBody), + ), + params( + ("TR-Dataset" = uuid::Uuid, Header, description = "The dataset id to use for the request"), + ), + security( + ("ApiKey" = ["owner"]), + ) +)] +pub async fn delete_experiment( + data: web::Path, + _user: AdminOnly, + clickhouse_client: web::Data, + dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan, +) -> Result { + delete_experiment_query( + data.into_inner(), + dataset_org_plan_sub.dataset.id, + clickhouse_client.get_ref(), + ) + .await?; + Ok(HttpResponse::NoContent().finish()) +} + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct AbTestReqBody { + pub experiment_id: uuid::Uuid, + pub user_id: String, +} + +#[derive(Debug, Serialize, Deserialize, ToSchema)] +pub struct UserTreatmentResponse { + pub treatment_name: String, + #[schema(value_type = String, format = Uuid)] + pub experiment_id: uuid::Uuid, + pub user_id: String, +} + +/// Ab Test +/// +/// Get a user's treatment for an experiment. Auth'ed user must be an owner of the organization to get a user's treatment. +#[utoipa::path( + post, + path = "/experiment/ab-test", + context_path = "/api", + tag = "Experiment", + request_body(content = AbTestReqBody, description = "JSON request payload to get a user's treatment", content_type = "application/json"), + responses( + (status = 200, description = "User treatment response", body = UserTreatmentResponse), + (status = 400, description = "Service error relating to getting the user's treatment", body = ErrorResponseBody), + ), + params( + ("TR-Dataset" = uuid::Uuid, Header, description = "The dataset id to use for the request"), + ), + security( + ("ApiKey" = ["owner"]), + ) +)] +pub async fn ab_test( + data: web::Json, + _user: AdminOnly, + clickhouse_client: web::Data, + dataset_org_plan_sub: DatasetAndOrgWithSubAndPlan, +) -> Result { + let treatment_response = ab_test_query( + data.experiment_id, + dataset_org_plan_sub.dataset.id, + data.user_id.clone(), + clickhouse_client.get_ref(), + ) + .await?; + Ok(HttpResponse::Ok().json(treatment_response)) +} diff --git a/server/src/handlers/mod.rs b/server/src/handlers/mod.rs index 9a0c8693e3..a4daba1064 100644 --- a/server/src/handlers/mod.rs +++ b/server/src/handlers/mod.rs @@ -5,6 +5,7 @@ pub mod crawl_handler; pub mod dataset_handler; pub mod etl_handler; pub mod event_handler; +pub mod experiment_handler; pub mod file_handler; pub mod group_handler; pub mod invitation_handler; diff --git a/server/src/lib.rs b/server/src/lib.rs index f93a7c73b7..bb78bfa3b2 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -285,6 +285,11 @@ impl Modify for SecurityAddon { handlers::page_handler::public_page, handlers::etl_handler::create_etl_job, handlers::payment_handler::handle_shopify_plan_change, + handlers::experiment_handler::create_experiment, + handlers::experiment_handler::get_experiments, + handlers::experiment_handler::update_experiment, + handlers::experiment_handler::delete_experiment, + handlers::experiment_handler::ab_test, ), components( schemas( @@ -478,6 +483,12 @@ impl Modify for SecurityAddon { handlers::payment_handler::ShopifyPlanChangePayload, handlers::payment_handler::ShopifyPlan, handlers::dataset_handler::CloneDatasetRequest, + handlers::experiment_handler::UpdateExperimentReqBody, + handlers::experiment_handler::AbTestReqBody, + handlers::experiment_handler::UserTreatmentResponse, + handlers::experiment_handler::CreateExperimentReqBody, + handlers::experiment_handler::ExperimentConfig, + data::models::Experiment, data::models::UserApiKey, data::models::CrawlStatus, data::models::CrawlType, @@ -695,6 +706,7 @@ impl Modify for SecurityAddon { (name = "Health", description = "Health check endpoint. Used to check if the server is up and running."), (name = "Metrics", description = "Metrics endpoint. Used to get information for monitoring"), (name = "Analytics", description = "Analytics endpoint. Used to get information for search and RAG analytics"), + (name = "Experiment", description = "Experiment endpoint. Used to create and manage experiments"), ), )] pub struct ApiDoc; @@ -1526,7 +1538,24 @@ pub fn main() -> std::io::Result<()> { web::resource("/ctr") .route(web::put().to(handlers::analytics_handler::send_ctr_data)) ) - ), + ) + .service( + web::scope("/experiment") + .service( + web::resource("/ab-test") + .route(web::post().to(handlers::experiment_handler::ab_test)), + ) + .service( + web::resource("") + .route(web::post().to(handlers::experiment_handler::create_experiment)) + .route(web::get().to(handlers::experiment_handler::get_experiments)) + .route(web::put().to(handlers::experiment_handler::update_experiment)) + ) + .service( + web::resource("/{experiment_id}") + .route(web::delete().to(handlers::experiment_handler::delete_experiment)) + ) + ) ) }) .workers(num_workers) diff --git a/server/src/operators/experiment_operator.rs b/server/src/operators/experiment_operator.rs new file mode 100644 index 0000000000..b43116bfe2 --- /dev/null +++ b/server/src/operators/experiment_operator.rs @@ -0,0 +1,229 @@ +use serde_json::Value; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use time::OffsetDateTime; + +use crate::data::models::{Experiment, ExperimentClickhouse, ExperimentUserAssignment}; +use crate::handlers::experiment_handler::UserTreatmentResponse; +use crate::{ + errors::ServiceError, + handlers::experiment_handler::{CreateExperimentReqBody, UpdateExperimentReqBody}, +}; + +pub async fn create_experiment_query( + data: CreateExperimentReqBody, + dataset_id: uuid::Uuid, + clickhouse_client: &clickhouse::Client, +) -> Result { + let experiment: Experiment = data.to_experiment(dataset_id); + let experiment_clickhouse: ExperimentClickhouse = experiment.clone().into(); + let mut insert = clickhouse_client + .insert("experiments") + .map_err(|e| ServiceError::InternalServerError(e.to_string()))?; + + insert.write(&experiment_clickhouse).await.map_err(|e| { + ServiceError::InternalServerError(format!("Failed to insert experiment: {}", e)) + })?; + + insert.end().await.map_err(|e| { + ServiceError::InternalServerError(format!("Failed to end experiment insert: {}", e)) + })?; + + Ok(experiment) +} + +pub async fn get_experiments_query( + dataset_id: uuid::Uuid, + clickhouse_client: &clickhouse::Client, +) -> Result, ServiceError> { + let experiments = clickhouse_client + .query("SELECT ?fields FROM experiments WHERE dataset_id = ?") + .bind(dataset_id) + .fetch_all() + .await + .map_err(|e| ServiceError::InternalServerError(e.to_string()))?; + Ok(experiments + .into_iter() + .map(|e: ExperimentClickhouse| e.into()) + .collect()) +} + +pub async fn get_experiment_by_id_query( + experiment_id: uuid::Uuid, + dataset_id: uuid::Uuid, + clickhouse_client: &clickhouse::Client, +) -> Result { + let experiment: ExperimentClickhouse = clickhouse_client + .query("SELECT ?fields FROM experiments WHERE id = ? AND dataset_id = ?") + .bind(experiment_id) + .bind(dataset_id) + .fetch_one() + .await + .map_err(|e| { + ServiceError::NotFound(format!( + "Experiment with id {} not found in dataset {}: {}", + experiment_id, dataset_id, e + )) + })?; + Ok(experiment.into()) +} + +pub async fn update_experiment_query( + data: UpdateExperimentReqBody, + dataset_id: uuid::Uuid, + clickhouse_client: &clickhouse::Client, +) -> Result { + let mut experiment: ExperimentClickhouse = clickhouse_client + .query("SELECT ?fields FROM experiments WHERE id = ? AND dataset_id = ?") + .bind(data.id) + .bind(dataset_id) + .fetch_one() + .await + .map_err(|e| { + ServiceError::NotFound(format!( + "Experiment with id {} not found in dataset {}: {}", + data.id, dataset_id, e + )) + })?; + + let mut set_clauses: Vec = Vec::new(); + let mut params: Vec = Vec::new(); + + if let Some(new_name) = &data.name { + if *new_name != experiment.name { + set_clauses.push("name = ?".to_string()); + params.push(new_name.clone().into()); + experiment.name = new_name.clone(); + } + } + + if let Some(config) = &data.experiment_config { + if config.t1_name != experiment.t1_name { + set_clauses.push("t1_name = ?".to_string()); + params.push(config.t1_name.clone().into()); + experiment.t1_name = config.t1_name.clone(); + } + if config.t1_split != experiment.t1_split { + set_clauses.push("t1_split = ?".to_string()); + params.push(config.t1_split.into()); + experiment.t1_split = config.t1_split; + } + if config.control_name != experiment.control_name { + set_clauses.push("control_name = ?".to_string()); + params.push(config.control_name.clone().into()); + experiment.control_name = config.control_name.clone(); + } + if config.control_split != experiment.control_split { + set_clauses.push("control_split = ?".to_string()); + params.push(config.control_split.into()); + experiment.control_split = config.control_split; + } + } + + if set_clauses.is_empty() { + return Ok(experiment.into()); + } + + let new_updated_at = OffsetDateTime::now_utc(); + set_clauses.push("updated_at = ?".to_string()); + + experiment.updated_at = new_updated_at; + + let query_str = format!( + "ALTER TABLE experiments UPDATE {} WHERE id = ? AND dataset_id = ?", + set_clauses.join(", ") + ); + + let mut exec_query = clickhouse_client.query(&query_str); + for param_value in params { + exec_query = exec_query.bind(param_value); + } + exec_query = exec_query.bind(new_updated_at.unix_timestamp()); + exec_query = exec_query.bind(data.id); + exec_query = exec_query.bind(dataset_id); + + exec_query.execute().await.map_err(|e| { + ServiceError::InternalServerError(format!("Failed to update experiment: {}", e)) + })?; + + Ok(experiment.into()) +} + +pub async fn delete_experiment_query( + experiment_id: uuid::Uuid, + dataset_id: uuid::Uuid, + clickhouse_client: &clickhouse::Client, +) -> Result<(), ServiceError> { + clickhouse_client + .query("DELETE FROM experiments WHERE id = ? AND dataset_id = ?") + .bind(experiment_id) + .bind(dataset_id) + .execute() + .await + .map_err(|e| ServiceError::InternalServerError(e.to_string()))?; + Ok(()) +} + +pub async fn ab_test_query( + experiment_id: uuid::Uuid, + dataset_id: uuid::Uuid, + user_id: String, + clickhouse_client: &clickhouse::Client, +) -> Result { + let existing_assignment: Option = clickhouse_client + .query("SELECT ?fields FROM experiment_user_assignments WHERE experiment_id = ? AND user_id = ? AND dataset_id = ?") + .bind(experiment_id) + .bind(user_id.clone()) + .bind(dataset_id) + .fetch_optional() + .await + .map_err(|e| ServiceError::InternalServerError(format!("Error fetching user assignment: {}", e)))?; + + if let Some(assignment) = existing_assignment { + return Ok(UserTreatmentResponse { + treatment_name: assignment.treatment_name, + experiment_id: assignment.experiment_id, + user_id: assignment.user_id, + }); + } + + let experiment = + get_experiment_by_id_query(experiment_id, dataset_id, clickhouse_client).await?; + + let mut hasher = DefaultHasher::new(); + user_id.hash(&mut hasher); + experiment_id.hash(&mut hasher); + let hash_value = hasher.finish(); + + let assigned_treatment_name = + if experiment.t1_split > 0.0 && (hash_value % 100) < (experiment.t1_split * 100.0) as u64 { + experiment.t1_name.clone() + } else { + experiment.control_name.clone() + }; + + let new_assignment = ExperimentUserAssignment { + id: uuid::Uuid::new_v4(), + experiment_id, + user_id: user_id.clone(), + dataset_id, + treatment_name: assigned_treatment_name.clone(), + created_at: OffsetDateTime::now_utc(), + updated_at: OffsetDateTime::now_utc(), + }; + + clickhouse_client + .insert("experiment_user_assignments") // Ensure this table name matches your schema + .map_err(|e| ServiceError::InternalServerError(format!("Insert setup failed: {}", e)))? + .write(&new_assignment) + .await + .map_err(|e| { + ServiceError::InternalServerError(format!("Error inserting user assignment: {}", e)) + })?; + + Ok(UserTreatmentResponse { + treatment_name: assigned_treatment_name, + experiment_id, + user_id, + }) +} diff --git a/server/src/operators/mod.rs b/server/src/operators/mod.rs index d572332e6d..f2da6487c3 100644 --- a/server/src/operators/mod.rs +++ b/server/src/operators/mod.rs @@ -7,6 +7,7 @@ pub mod dittofeed_operator; pub mod email_operator; pub mod etl_operator; pub mod event_operator; +pub mod experiment_operator; pub mod file_operator; pub mod group_operator; pub mod invitation_operator;