feat: make datasetId optional for experiments

Chainlit · Jul 24, 2024 · 9888db1 · 9888db1
1 parent 28fc8c4
commit 9888db1
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 19 deletions.
diff --git a/src/api.ts b/src/api.ts
@@ -32,6 +32,7 @@ import {
   PaginatedResponse,
   Prompt,
   Score,
+  ScoreConstructor,
   Step,
   StepType,
   Thread,
@@ -234,7 +235,7 @@ function ingestStepsQueryBuilder(steps: Step[]) {
     `;
 }
 
-function createScoresFieldsBuilder(scores: Score[]) {
+function createScoresFieldsBuilder(scores: ScoreConstructor[]) {
   let generated = '';
   for (let id = 0; id < scores.length; id++) {
     generated += `$name_${id}: String!
@@ -251,7 +252,7 @@ function createScoresFieldsBuilder(scores: Score[]) {
   return generated;
 }
 
-function createScoresArgsBuilder(scores: Score[]) {
+function createScoresArgsBuilder(scores: ScoreConstructor[]) {
   let generated = '';
   for (let id = 0; id < scores.length; id++) {
     generated += `
@@ -278,7 +279,7 @@ function createScoresArgsBuilder(scores: Score[]) {
   return generated;
 }
 
-function createScoresQueryBuilder(scores: Score[]) {
+function createScoresQueryBuilder(scores: ScoreConstructor[]) {
   return `
     mutation CreateScores(${createScoresFieldsBuilder(scores)}) {
       ${createScoresArgsBuilder(scores)}
@@ -749,6 +750,7 @@ export class API {
               error
               type
               tags
+              metadata
             }
           }
         }
@@ -1741,12 +1743,12 @@ export class API {
 
   public async createExperiment(datasetExperiment: {
     name: string;
-    datasetId: string;
+    datasetId?: string;
     promptId?: string;
     params?: Record<string, any> | Array<Record<string, any>>;
   }) {
     const query = `
-      mutation CreateDatasetExperiment($name: String!, $datasetId: String! $promptId: String, $params: Json) {
+      mutation CreateDatasetExperiment($name: String!, $datasetId: String $promptId: String, $params: Json) {
         createDatasetExperiment(name: $name, datasetId: $datasetId, promptId: $promptId, params: $params) {
           id
         }
@@ -1771,7 +1773,7 @@ export class API {
     scores
   }: DatasetExperimentItem) {
     const query = `
-      mutation CreateDatasetExperimentItem($datasetExperimentId: String!, $datasetItemId: String!, $input: Json, $output: Json) {
+      mutation CreateDatasetExperimentItem($datasetExperimentId: String!, $datasetItemId: String, $input: Json, $output: Json) {
         createDatasetExperimentItem(datasetExperimentId: $datasetExperimentId, datasetItemId: $datasetItemId, input: $input, output: $output) {
           id
           input
@@ -1791,7 +1793,7 @@ export class API {
       scores.map((score) => {
         score.datasetExperimentItemId =
           result.data.createDatasetExperimentItem.id;
-        return score;
+        return new Score(score);
       })
     );
 

diff --git a/src/generation.ts b/src/generation.ts
@@ -52,6 +52,7 @@ export class BaseGeneration extends Utils {
   model?: Maybe<string>;
   id?: Maybe<string>;
   tags?: Maybe<string[]>;
+  metadata?: Maybe<Record<string, any>>;
   error?: Maybe<string>;
   variables?: Maybe<Record<string, any>>;
   settings?: Maybe<ILLMSettings>;

diff --git a/src/instrumentation/openai.ts b/src/instrumentation/openai.ts
@@ -281,6 +281,7 @@ export type OpenAIOutput =
 
 export interface InstrumentOpenAIOptions {
   tags?: Maybe<string[]>;
+  metadata?: Maybe<Record<string, any>>;
 }
 
 export interface ProcessOpenAIOutput extends InstrumentOpenAIOptions {
@@ -301,13 +302,14 @@ function isStream(obj: any): boolean {
 const processOpenAIOutput = async (
   client: LiteralClient,
   output: OpenAIOutput,
-  { start, tags, inputs }: ProcessOpenAIOutput
+  { start, tags, metadata, inputs }: ProcessOpenAIOutput
 ) => {
   const baseGeneration = {
     provider: 'openai',
     model: inputs.model,
     settings: getSettings(inputs),
-    tags: tags
+    tags: tags,
+    metadata: metadata
   };
 
   let threadFromStore: Thread | null = null;
@@ -336,7 +338,8 @@ const processOpenAIOutput = async (
       output: output,
       startTime: new Date(start).toISOString(),
       endTime: new Date().toISOString(),
-      tags: tags
+      tags: tags,
+      metadata: metadata
     };
 
     const step = parent

diff --git a/src/types.ts b/src/types.ts
@@ -78,11 +78,7 @@ export class Utils {
 
 export type ScoreType = 'HUMAN' | 'AI';
 
-/**
- * Represents a score entity with properties to track various aspects of scoring.
- * It extends the `Utils` class for serialization capabilities.
- */
-export class Score extends Utils {
+class ScoreFields extends Utils {
   id?: Maybe<string>;
   stepId?: Maybe<string>;
   generationId?: Maybe<string>;
@@ -93,13 +89,23 @@ export class Score extends Utils {
   scorer?: Maybe<string>;
   comment?: Maybe<string>;
   tags?: Maybe<string[]>;
+}
+
+export type ScoreConstructor = OmitUtils<ScoreFields>;
 
-  constructor(data: OmitUtils<Score>) {
+/**
+ * Represents a score entity with properties to track various aspects of scoring.
+ * It extends the `Utils` class for serialization capabilities.
+ */
+export class Score extends ScoreFields {
+  constructor(data: ScoreConstructor) {
     super();
     Object.assign(this, data);
   }
 }
 
+export type ScoreCons = OmitUtils<StepFields>;
+
 /**
  * Represents an attachment with optional metadata, MIME type, and other properties.
  * It extends the `Utils` class for serialization capabilities.
@@ -631,8 +637,8 @@ export class DatasetItem extends Utils {
 class DatasetExperimentItemFields extends Utils {
   id?: string;
   datasetExperimentId!: string;
-  datasetItemId!: string;
-  scores!: Score[];
+  datasetItemId?: string;
+  scores!: ScoreConstructor[];
   input?: Record<string, any>;
   output?: Record<string, any>;
 }
@@ -641,7 +647,7 @@ export class DatasetExperiment extends Utils {
   id!: string;
   createdAt!: string;
   name!: string;
-  datasetId!: string;
+  datasetId?: string;
   promptId?: string;
   api: API;
   params!: Record<string, any> | Array<Record<string, any>>;