feat(evaluators): db migration for evaluator tables (#9960)

Co-authored-by: Mikyo King <[email protected]>
Co-authored-by: Xander Song <[email protected]>

Author: 3 people

app/schema.graphql

@@ -387,6 +387,22 @@ input ClusterInput {
   id: ID
 }
 
+type CodeEvaluator implements Evaluator & Node {
+  """The Globally Unique ID of this object"""
+  id: ID!
+  name: Identifier!
+  description: String
+  kind: EvaluatorKind!
+  createdAt: DateTime!
+  updatedAt: DateTime!
+  user: User
+}
+
+type CodeEvaluatorMutationPayload {
+  evaluator: CodeEvaluator!
+  query: Query!
+}
+
 union ContentPart = TextContentPart | ToolCallContentPart | ToolResultContentPart
 
 input ContentPartInput @oneOf {
@@ -449,12 +465,25 @@ input CreateChatPromptVersionInput {
   tags: [SetPromptVersionTagInput!] = null
 }
 
+input CreateCodeEvaluatorInput {
+  datasetId: ID!
+  name: Identifier!
+  description: String
+}
+
 input CreateDatasetInput {
   name: String!
   description: String
   metadata: JSON
 }
 
+input CreateDatasetLLMEvaluatorInput {
+  datasetId: ID!
+  name: Identifier!
+  description: String
+  promptVersion: ChatPromptVersionInput!
+}
+
 input CreateDatasetLabelInput {
   name: String!
   description: String
@@ -633,6 +662,7 @@ type Dataset implements Node {
   experiments(first: Int = 50, last: Int, after: String, before: String, filterCondition: String, filterIds: [ID!]): ExperimentConnection!
   experimentAnnotationSummaries: [DatasetExperimentAnnotationSummary!]!
   labels: [DatasetLabel!]!
+  evaluators(first: Int = 50, last: Int, after: String, before: String): EvaluatorConnection!
   lastUpdatedAt: DateTime
 }
 
@@ -1236,6 +1266,39 @@ input EvalResultKey {
   attr: EvalAttr!
 }
 
+interface Evaluator implements Node {
+  """The Globally Unique ID of this object"""
+  id: ID!
+  name: Identifier!
+  description: String
+  kind: EvaluatorKind!
+  createdAt: DateTime!
+  updatedAt: DateTime!
+}
+
+"""A connection to a list of items."""
+type EvaluatorConnection {
+  """Pagination data for this connection"""
+  pageInfo: PageInfo!
+
+  """Contains the nodes in this connection"""
+  edges: [EvaluatorEdge!]!
+}
+
+"""An edge in a connection."""
+type EvaluatorEdge {
+  """A cursor for use in pagination"""
+  cursor: String!
+
+  """The item at the end of the edge"""
+  node: Evaluator!
+}
+
+enum EvaluatorKind {
+  LLM
+  CODE
+}
+
 type Event {
   id: ID!
   eventMetadata: EventMetadata!
@@ -1768,6 +1831,25 @@ type JSONInvocationParameter implements InvocationParameterBase {
   defaultValue: JSON
 }
 
+type LLMEvaluator implements Evaluator & Node {
+  """The Globally Unique ID of this object"""
+  id: ID!
+  name: Identifier!
+  description: String
+  kind: EvaluatorKind!
+  createdAt: DateTime!
+  updatedAt: DateTime!
+  prompt: Prompt!
+  promptVersionTag: PromptVersionTag
+  user: User
+  promptVersion: PromptVersion!
+}
+
+type LLMEvaluatorMutationPayload {
+  evaluator: LLMEvaluator!
+  query: Query!
+}
+
 type LabelFraction {
   label: String!
   fraction: Float!
@@ -1831,6 +1913,8 @@ type Mutation {
   addDatasetExamplesToDatasetSplits(input: AddDatasetExamplesToDatasetSplitsInput!): AddDatasetExamplesToDatasetSplitsMutationPayload!
   removeDatasetExamplesFromDatasetSplits(input: RemoveDatasetExamplesFromDatasetSplitsInput!): RemoveDatasetExamplesFromDatasetSplitsMutationPayload!
   createDatasetSplitWithExamples(input: CreateDatasetSplitWithExamplesInput!): DatasetSplitMutationPayloadWithExamples!
+  createDatasetCodeEvaluator(input: CreateCodeEvaluatorInput!): CodeEvaluatorMutationPayload!
+  createDatasetLlmEvaluator(input: CreateDatasetLLMEvaluatorInput!): LLMEvaluatorMutationPayload!
   deleteExperiments(input: DeleteExperimentsInput!): ExperimentMutationPayload!
 
   """
@@ -2501,6 +2585,7 @@ type Query {
   promptLabels(first: Int = 50, last: Int, after: String, before: String): PromptLabelConnection!
   datasetLabels(first: Int = 50, last: Int, after: String, before: String): DatasetLabelConnection!
   datasetSplits(first: Int = 50, last: Int, after: String, before: String): DatasetSplitConnection!
+  evaluators(first: Int = 50, last: Int, after: String, before: String): EvaluatorConnection!
   annotationConfigs(first: Int = 50, last: Int = null, after: String = null, before: String = null): AnnotationConfigConnection!
   clusters(clusters: [ClusterInput!]!): [Cluster!]!
   hdbscanClustering(

scripts/ddl/postgresql_schema.sql

@@ -650,6 +650,70 @@ CREATE INDEX ix_document_annotations_span_rowid ON public.document_annotations
     USING btree (span_rowid);
 
 
+-- Table: evaluators
+-- -----------------
+CREATE TABLE public.evaluators (
+    id bigserial NOT NULL,
+    name VARCHAR NOT NULL,
+    description VARCHAR,
+    kind VARCHAR NOT NULL,
+    user_id BIGINT,
+    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
+    CONSTRAINT pk_evaluators PRIMARY KEY (id),
+    CONSTRAINT uq_evaluators_kind_id
+        UNIQUE (kind, id),
+    CONSTRAINT uq_evaluators_name
+        UNIQUE (name),
+    CHECK (((kind)::text = ANY ((ARRAY[
+            'LLM'::character varying,
+            'CODE'::character varying
+        ])::text[]))),
+    CONSTRAINT fk_evaluators_user_id_users FOREIGN KEY
+        (user_id)
+        REFERENCES public.users (id)
+        ON DELETE SET NULL
+);
+
+CREATE INDEX ix_evaluators_user_id ON public.evaluators
+    USING btree (user_id);
+
+
+-- Table: code_evaluators
+-- ----------------------
+CREATE TABLE public.code_evaluators (
+    id BIGINT NOT NULL,
+    kind VARCHAR NOT NULL DEFAULT 'CODE'::character varying,
+    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
+    CONSTRAINT pk_code_evaluators PRIMARY KEY (id),
+    CHECK (((kind)::text = 'CODE'::text)),
+    CONSTRAINT fk_code_evaluators_kind_evaluators FOREIGN KEY
+        (kind, id)
+        REFERENCES public.evaluators (kind, id)
+        ON DELETE CASCADE
+);
+
+
+-- Table: datasets_evaluators
+-- --------------------------
+CREATE TABLE public.datasets_evaluators (
+    dataset_id BIGINT NOT NULL,
+    evaluator_id BIGINT NOT NULL,
+    input_config JSONB NOT NULL,
+    CONSTRAINT pk_datasets_evaluators PRIMARY KEY (dataset_id, evaluator_id),
+    CONSTRAINT fk_datasets_evaluators_dataset_id_datasets FOREIGN KEY
+        (dataset_id)
+        REFERENCES public.datasets (id)
+        ON DELETE CASCADE,
+    CONSTRAINT fk_datasets_evaluators_evaluator_id_evaluators FOREIGN KEY
+        (evaluator_id)
+        REFERENCES public.evaluators (id)
+        ON DELETE CASCADE
+);
+
+CREATE INDEX ix_datasets_evaluators_evaluator_id ON public.datasets_evaluators
+    USING btree (evaluator_id);
+
+
 -- Table: experiments
 -- ------------------
 CREATE TABLE public.experiments (
@@ -976,6 +1040,38 @@ CREATE INDEX ix_prompt_version_tags_user_id ON public.prompt_version_tags
     USING btree (user_id);
 
 
+-- Table: llm_evaluators
+-- ---------------------
+CREATE TABLE public.llm_evaluators (
+    id BIGINT NOT NULL,
+    kind VARCHAR NOT NULL DEFAULT 'LLM'::character varying,
+    prompt_id BIGINT NOT NULL,
+    prompt_version_tag_id BIGINT,
+    output_config JSONB NOT NULL,
+    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
+    CONSTRAINT pk_llm_evaluators PRIMARY KEY (id),
+    CHECK (((kind)::text = 'LLM'::text)),
+    CONSTRAINT fk_llm_evaluators_kind_evaluators FOREIGN KEY
+        (kind, id)
+        REFERENCES public.evaluators (kind, id)
+        ON DELETE CASCADE,
+    CONSTRAINT fk_llm_evaluators_prompt_id_prompts FOREIGN KEY
+        (prompt_id)
+        REFERENCES public.prompts (id)
+        ON DELETE RESTRICT,
+    CONSTRAINT fk_llm_evaluators_prompt_version_tag_id_prompt_version_tags
+        FOREIGN KEY
+        (prompt_version_tag_id)
+        REFERENCES public.prompt_version_tags (id)
+        ON DELETE SET NULL
+);
+
+CREATE INDEX ix_llm_evaluators_prompt_id ON public.llm_evaluators
+    USING btree (prompt_id);
+CREATE INDEX ix_llm_evaluators_prompt_version_tag_id ON public.llm_evaluators
+    USING btree (prompt_version_tag_id);
+
+
 -- Table: refresh_tokens
 -- ---------------------
 CREATE TABLE public.refresh_tokens (