From a094b2f0a2727f57457d5d831bf7eccc3f69184b Mon Sep 17 00:00:00 2001 From: Willy Douhard Date: Wed, 24 Jul 2024 12:05:32 +0200 Subject: [PATCH 1/4] feat: make datasetId optional for experiments --- src/api.ts | 15 ++++++++------- src/types.ts | 24 +++++++++++++++--------- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/src/api.ts b/src/api.ts index dd4426c..6ae6425 100644 --- 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)} @@ -1741,12 +1742,12 @@ export class API { public async createExperiment(datasetExperiment: { name: string; - datasetId: string; + datasetId?: string; promptId?: string; params?: Record | Array>; }) { 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 +1772,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 +1792,7 @@ export class API { scores.map((score) => { score.datasetExperimentItemId = result.data.createDatasetExperimentItem.id; - return score; + return new Score(score); }) ); diff --git a/src/types.ts b/src/types.ts index 536f0d3..97e902b 100644 --- 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; stepId?: Maybe; generationId?: Maybe; @@ -93,13 +89,23 @@ export class Score extends Utils { scorer?: Maybe; comment?: Maybe; tags?: Maybe; +} + +export type ScoreConstructor = OmitUtils; - constructor(data: OmitUtils) { +/** + * 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; + /** * 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; output?: Record; } @@ -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 | Array>; From 1d79bae9f91db5deeadb4c0eabf6db607d3fbdec Mon Sep 17 00:00:00 2001 From: Willy Douhard Date: Fri, 26 Jul 2024 10:58:57 +0200 Subject: [PATCH 2/4] feat: wip experiment run --- src/api.ts | 39 +++++++++----- src/index.ts | 32 ++++++++++-- src/types.ts | 78 ++++++++++++++++++++++++++-- tests/api.test.ts | 2 +- tests/integration/llamaindex.test.ts | 6 +-- tests/integration/openai.test.ts | 14 ++--- tests/integration/vercel-sdk.test.ts | 6 +-- tests/wrappers.test.ts | 36 ++++++++++++- 8 files changed, 175 insertions(+), 38 deletions(-) diff --git a/src/api.ts b/src/api.ts index 6ae6425..31aca35 100644 --- a/src/api.ts +++ b/src/api.ts @@ -27,6 +27,7 @@ import { DatasetExperimentItem, DatasetItem, DatasetType, + Environment, Maybe, OmitUtils, PaginatedResponse, @@ -56,6 +57,7 @@ const stepFields = ` input output metadata + environment scores { id type @@ -330,7 +332,7 @@ function addGenerationsToDatasetQueryBuilder(generationIds: string[]) { export class API { /** @ignore */ - private client: LiteralClient; + public client: LiteralClient; /** @ignore */ private apiKey: string; /** @ignore */ @@ -340,28 +342,33 @@ export class API { /** @ignore */ private restEndpoint: string; /** @ignore */ + public environment: Environment | undefined; + /** @ignore */ public disabled: boolean; /** @ignore */ constructor( client: LiteralClient, - apiKey: string, - url: string, + apiKey?: string, + url?: string, + environment?: Environment, disabled?: boolean ) { this.client = client; - this.apiKey = apiKey; - this.url = url; - this.graphqlEndpoint = `${url}/api/graphql`; - this.restEndpoint = `${url}/api`; - this.disabled = !!disabled; - if (!this.apiKey) { + if (!apiKey) { throw new Error('LITERAL_API_KEY not set'); } - if (!this.url) { + if (!url) { throw new Error('LITERAL_API_URL not set'); } + + this.apiKey = apiKey; + this.url = url; + this.environment = environment; + this.graphqlEndpoint = `${url}/api/graphql`; + this.restEndpoint = `${url}/api`; + this.disabled = !!disabled; } /** @ignore */ @@ -369,8 +376,9 @@ export class API { return { 'Content-Type': 'application/json', 'x-api-key': this.apiKey, - 'x-client-name': 'js-literal-client', - 'x-client-version': version + 'x-client-name': 'ts-literal-client', + 'x-client-version': version, + 'x-env': this.environment }; } @@ -1767,16 +1775,18 @@ export class API { public async createExperimentItem({ datasetExperimentId, datasetItemId, + experimentRunId, input, output, scores }: DatasetExperimentItem) { const query = ` - mutation CreateDatasetExperimentItem($datasetExperimentId: String!, $datasetItemId: String, $input: Json, $output: Json) { - createDatasetExperimentItem(datasetExperimentId: $datasetExperimentId, datasetItemId: $datasetItemId, input: $input, output: $output) { + mutation CreateDatasetExperimentItem($datasetExperimentId: String!, $datasetItemId: String, $experimentRunId: String, $input: Json, $output: Json) { + createDatasetExperimentItem(datasetExperimentId: $datasetExperimentId, datasetItemId: $datasetItemId, experimentRunId: $experimentRunId, input: $input, output: $output) { id input output + experimentRunId } } `; @@ -1784,6 +1794,7 @@ export class API { const result = await this.makeGqlCall(query, { datasetExperimentId, datasetItemId, + experimentRunId, input, output }); diff --git a/src/index.ts b/src/index.ts index 9750dfa..6215387 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,7 +3,14 @@ import { AsyncLocalStorage } from 'node:async_hooks'; import { API } from './api'; import instrumentation from './instrumentation'; import openai from './openai'; -import { Step, StepConstructor, Thread, ThreadConstructor } from './types'; +import { + Environment, + ExperimentRun, + Step, + StepConstructor, + Thread, + ThreadConstructor +} from './types'; export * from './types'; export * from './generation'; @@ -13,6 +20,7 @@ export type * from './instrumentation'; type StoredContext = { currentThread: Thread | null; currentStep: Step | null; + currentExperimentRunId: string | null; }; const storage = new AsyncLocalStorage(); @@ -23,7 +31,17 @@ export class LiteralClient { instrumentation: ReturnType; store: AsyncLocalStorage = storage; - constructor(apiKey?: string, apiUrl?: string, disabled?: boolean) { + constructor({ + apiKey, + apiUrl, + environment, + disabled + }: { + apiKey?: string; + apiUrl?: string; + environment?: Environment; + disabled?: boolean; + } = {}) { if (!apiKey) { apiKey = process.env.LITERAL_API_KEY; } @@ -32,7 +50,7 @@ export class LiteralClient { apiUrl = process.env.LITERAL_API_URL || 'https://cloud.getliteral.ai'; } - this.api = new API(this, apiKey!, apiUrl!, disabled); + this.api = new API(this, apiKey, apiUrl, environment, disabled); this.openai = openai(this); this.instrumentation = instrumentation(this); } @@ -49,6 +67,14 @@ export class LiteralClient { return this.step({ ...data, type: 'run' }); } + experimentRun(data?: Omit) { + return new ExperimentRun(this, { + ...(data || {}), + name: 'Experiment Run', + type: 'run' + }); + } + /** * Gets the current thread from the context. * WARNING : this will throw if run outside of a thread context. diff --git a/src/types.ts b/src/types.ts index 97e902b..fd66d3e 100644 --- a/src/types.ts +++ b/src/types.ts @@ -14,6 +14,8 @@ export type Maybe = T | null | undefined; export type OmitUtils = Omit; +export type Environment = 'dev' | 'staging' | 'prod' | 'experiment'; + export type PageInfo = { hasNextPage: boolean; startCursor: string; @@ -104,8 +106,6 @@ export class Score extends ScoreFields { } } -export type ScoreCons = OmitUtils; - /** * Represents an attachment with optional metadata, MIME type, and other properties. * It extends the `Utils` class for serialization capabilities. @@ -226,7 +226,7 @@ export class Thread extends ThreadFields { | ((output: Output) => Promise) ) { const output = await this.client.store.run( - { currentThread: this, currentStep: null }, + { currentThread: this, currentStep: null, currentExperimentRunId: null }, () => cb(this) ); @@ -268,6 +268,7 @@ class StepFields extends Utils { createdAt?: Maybe; startTime?: Maybe; id?: Maybe; + environment?: Maybe; error?: Maybe>; input?: Maybe>; output?: Maybe>; @@ -411,7 +412,11 @@ export class Step extends StepFields { const currentStore = this.client.store.getStore(); const output = await this.client.store.run( - { currentThread: currentStore?.currentThread ?? null, currentStep: this }, + { + currentThread: currentStore?.currentThread ?? null, + currentStep: this, + currentExperimentRunId: null + }, () => cb(this) ); @@ -447,6 +452,64 @@ export class Step extends StepFields { } } +/** + * Represents a step in a process or workflow, extending the fields and methods from StepFields. + */ +export class ExperimentRun extends Step { + api: API; + client: LiteralClient; + + /** + * Constructs a new ExperimentRun instance. + * @param api The API instance to be used for sending and managing steps. + * @param data The initial data for the step, excluding utility properties. + */ + constructor( + client: LiteralClient, + data: StepConstructor, + ignoreContext?: true + ) { + super(client, data, ignoreContext); + this.client = client; + this.api = client.api; + } + + async wrap( + cb: (step: Step) => Output | Promise, + updateStep?: + | Partial + | ((output: Output) => Partial) + | ((output: Output) => Promise>) + ) { + const originalEnvironment = this.api.environment; + this.api.environment = 'experiment'; + + const currentStore = this.client.store.getStore(); + const output: Output = await this.client.store.run( + { + currentThread: currentStore?.currentThread ?? null, + currentStep: this, + currentExperimentRunId: this.id ?? null + }, + async () => { + try { + const output = await super.wrap(cb, updateStep); + return output; + } finally { + // Clear the currentExperimentRunId after execution + const updatedStore = this.client.store.getStore(); + if (updatedStore) { + updatedStore.currentExperimentRunId = null; + } + } + } + ); + + this.api.environment = originalEnvironment; + return output; + } +} + /** * Represents a user with optional metadata and identifier. */ @@ -638,6 +701,7 @@ class DatasetExperimentItemFields extends Utils { id?: string; datasetExperimentId!: string; datasetItemId?: string; + experimentRunId?: string; scores!: ScoreConstructor[]; input?: Record; output?: Record; @@ -668,9 +732,13 @@ export class DatasetExperiment extends Utils { 'id' | 'datasetExperimentId' > ) { + const currentStore = this.api.client.store.getStore(); + const experimentRunId = currentStore?.currentExperimentRunId; + const datasetExperimentItem = new DatasetExperimentItem({ ...itemFields, - datasetExperimentId: this.id + datasetExperimentId: this.id, + ...(experimentRunId && { experimentRunId }) }); const item = await this.api.createExperimentItem(datasetExperimentItem); diff --git a/tests/api.test.ts b/tests/api.test.ts index 0522925..3e63401 100644 --- a/tests/api.test.ts +++ b/tests/api.test.ts @@ -21,7 +21,7 @@ describe('End to end tests for the SDK', function () { throw new Error('Missing environment variables'); } - client = new LiteralClient(apiKey, url); + client = new LiteralClient({ apiKey, apiUrl: url }); }); it('should test user', async function () { diff --git a/tests/integration/llamaindex.test.ts b/tests/integration/llamaindex.test.ts index bb201c8..2c7112e 100644 --- a/tests/integration/llamaindex.test.ts +++ b/tests/integration/llamaindex.test.ts @@ -16,14 +16,14 @@ describe('Llama Index Instrumentation', () => { let client: LiteralClient; beforeAll(function () { - const url = process.env.LITERAL_API_URL; + const apiUrl = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; - if (!url || !apiKey) { + if (!apiUrl || !apiKey) { throw new Error('Missing environment variables'); } - client = new LiteralClient(apiKey, url); + client = new LiteralClient({ apiKey, apiUrl }); // Reset the callback manager Settings.callbackManager = new CallbackManager(); diff --git a/tests/integration/openai.test.ts b/tests/integration/openai.test.ts index 4391961..53cf4d3 100644 --- a/tests/integration/openai.test.ts +++ b/tests/integration/openai.test.ts @@ -11,10 +11,10 @@ import { Step } from '../../src'; -const url = process.env.LITERAL_API_URL; +const apiUrl = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; -if (!url || !apiKey) { +if (!apiUrl || !apiKey) { throw new Error('Missing environment variables'); } @@ -100,7 +100,7 @@ describe('OpenAI Instrumentation', () => { beforeAll(async () => { const testId = uuidv4(); - const client = new LiteralClient(apiKey, url); + const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ tags: [testId] }); await openai.chat.completions.create({ @@ -169,7 +169,7 @@ describe('OpenAI Instrumentation', () => { beforeAll(async () => { const testId = uuidv4(); - const client = new LiteralClient(apiKey, url); + const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ tags: [testId] }); @@ -233,7 +233,7 @@ describe('OpenAI Instrumentation', () => { it('should monitor image generation', async () => { const testId = uuidv4(); - const client = new LiteralClient(apiKey, url); + const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ tags: [testId] }); const response = await openai.images.generate({ @@ -269,7 +269,7 @@ describe('OpenAI Instrumentation', () => { it('logs the generation inside its thread and parent', async () => { const testId = uuidv4(); - const client = new LiteralClient(apiKey, url); + const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ tags: [testId] }); let threadId: Maybe; @@ -312,7 +312,7 @@ describe('OpenAI Instrumentation', () => { it("doesn't mix up threads and steps", async () => { const testId = uuidv4(); - const client = new LiteralClient(apiKey, url); + const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ tags: [testId] }); const firstThreadId = uuidv4(); diff --git a/tests/integration/vercel-sdk.test.ts b/tests/integration/vercel-sdk.test.ts index bc604b3..10274bc 100644 --- a/tests/integration/vercel-sdk.test.ts +++ b/tests/integration/vercel-sdk.test.ts @@ -9,14 +9,14 @@ describe('Vercel SDK Instrumentation', () => { let client: LiteralClient; beforeAll(function () { - const url = process.env.LITERAL_API_URL; + const apiUrl = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; - if (!url || !apiKey) { + if (!apiUrl || !apiKey) { throw new Error('Missing environment variables'); } - client = new LiteralClient(apiKey, url); + client = new LiteralClient({ apiKey, apiUrl }); }); // Skip for the CI diff --git a/tests/wrappers.test.ts b/tests/wrappers.test.ts index 35cbfc0..0702fbb 100644 --- a/tests/wrappers.test.ts +++ b/tests/wrappers.test.ts @@ -1,6 +1,6 @@ import 'dotenv/config'; -import { LiteralClient, Maybe, Step } from '../src'; +import { DatasetExperimentItem, LiteralClient, Maybe, Step } from '../src'; const url = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; @@ -9,7 +9,7 @@ if (!url || !apiKey) { throw new Error('Missing environment variables'); } -const client = new LiteralClient(apiKey, url); +const client = new LiteralClient({ apiKey, apiUrl: url }); function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); @@ -350,6 +350,38 @@ describe('Wrapper', () => { }); }); + describe('Wrapping experimentRun', () => { + it('wraps an experiment run', async () => { + const experiment = await client.api.createExperiment({ + name: 'Test Experiment Run' + }); + let persistedExperimentItem: DatasetExperimentItem | undefined = + undefined; + client.experimentRun().wrap(async () => { + const scores = [ + { + name: 'context_relevancy', + type: 'AI' as const, + value: 0.6 + } + ]; + client.step({ name: 'agent', type: 'run' }).wrap(async () => { + const experimentItem = { + scores: scores, // scores in the format above + input: { question: 'question' }, + output: { content: 'answer' } + }; + persistedExperimentItem = await experiment.log(experimentItem); + }); + }); + const experimentRunId = persistedExperimentItem!.experimentRunId; + expect(experimentRunId).toBeTruthy(); + const experimentRun = await client.api.getStep(experimentRunId!); + expect(experimentRun).toBeTruthy(); + expect(experimentRun?.environment).toEqual('experiment'); + }); + }); + describe('Concurrency', () => { it("doesn't mix up threads and steps", async () => { let firstThreadId: Maybe; From e2c70cc0ee716add1bccfa60e6f46496a46c15d6 Mon Sep 17 00:00:00 2001 From: Willy Douhard Date: Fri, 26 Jul 2024 12:01:15 +0200 Subject: [PATCH 3/4] fix: tests --- src/api.ts | 22 ++++++++-------------- src/index.ts | 2 +- src/types.ts | 14 +++++++++----- tests/api.test.ts | 2 -- tests/wrappers.test.ts | 9 +++++++-- 5 files changed, 25 insertions(+), 24 deletions(-) diff --git a/src/api.ts b/src/api.ts index 79bd4d3..f2f2084 100644 --- a/src/api.ts +++ b/src/api.ts @@ -110,11 +110,14 @@ const threadFields = ` id identifier metadata - } - steps { - ${stepFields} }`; +const threadFieldsWithSteps = ` +${threadFields} +steps { + ${stepFields} +}`; + /** * Serializes the step object with a suffix ID to each key. * @@ -880,7 +883,6 @@ export class API { * @param options.name - The name of the thread. (Optional) * @param options.metadata - Additional metadata for the thread as a key-value pair object. (Optional) * @param options.participantId - The unique identifier of the participant. (Optional) - * @param options.environment - The environment where the thread is being upserted. (Optional) * @param options.tags - An array of tags associated with the thread. (Optional) * @returns The upserted thread object. */ @@ -889,7 +891,6 @@ export class API { name?: Maybe; metadata?: Maybe>; participantId?: Maybe; - environment?: Maybe; tags?: Maybe; }): Promise; @@ -901,7 +902,6 @@ export class API { * @param name - The name of the thread. (Optional) * @param metadata - Additional metadata for the thread as a key-value pair object. (Optional) * @param participantId - The unique identifier of the participant. (Optional) - * @param environment - The environment where the thread is being upserted. (Optional) * @param tags - An array of tags associated with the thread. (Optional) * @returns The upserted thread object. */ @@ -910,7 +910,6 @@ export class API { name?: Maybe, metadata?: Maybe>, participantId?: Maybe, - environment?: Maybe, tags?: Maybe ): Promise; @@ -919,7 +918,6 @@ export class API { name?: Maybe, metadata?: Maybe>, participantId?: Maybe, - environment?: Maybe, tags?: Maybe ): Promise { let threadId = threadIdOrOptions; @@ -928,7 +926,6 @@ export class API { name = threadIdOrOptions.name; metadata = threadIdOrOptions.metadata; participantId = threadIdOrOptions.participantId; - environment = threadIdOrOptions.environment; tags = threadIdOrOptions.tags; } @@ -938,7 +935,6 @@ export class API { $name: String, $metadata: Json, $participantId: String, - $environment: String, $tags: [String!], ) { upsertThread( @@ -946,7 +942,6 @@ export class API { name: $name metadata: $metadata participantId: $participantId - environment: $environment tags: $tags ) { ${threadFields} @@ -959,7 +954,6 @@ export class API { name, metadata, participantId, - environment, tags }; @@ -1019,7 +1013,7 @@ export class API { edges { cursor node { - ${threadFields} + ${threadFieldsWithSteps} } } } @@ -1047,7 +1041,7 @@ export class API { const query = ` query GetThread($id: String!) { threadDetail(id: $id) { - ${threadFields} + ${threadFieldsWithSteps} } } `; diff --git a/src/index.ts b/src/index.ts index 191554b..da2bf51 100644 --- a/src/index.ts +++ b/src/index.ts @@ -20,7 +20,7 @@ export type * from './instrumentation'; type StoredContext = { currentThread: Thread | null; currentStep: Step | null; - currentExperimentRunId: string | null; + currentExperimentRunId?: string | null; }; const storage = new AsyncLocalStorage(); diff --git a/src/types.ts b/src/types.ts index fd66d3e..6065566 100644 --- a/src/types.ts +++ b/src/types.ts @@ -206,7 +206,6 @@ export class Thread extends ThreadFields { name: this.name, metadata: this.metadata, participantId: this.participantId, - environment: this.environment, tags: this.tags }); return this; @@ -225,8 +224,14 @@ export class Thread extends ThreadFields { | ((output: Output) => ThreadConstructor) | ((output: Output) => Promise) ) { + const currentStore = this.client.store.getStore(); + const output = await this.client.store.run( - { currentThread: this, currentStep: null, currentExperimentRunId: null }, + { + currentThread: this, + currentExperimentRunId: currentStore?.currentExperimentRunId ?? null, + currentStep: null + }, () => cb(this) ); @@ -386,7 +391,6 @@ export class Step extends StepFields { if (this.api.disabled) { return this; } - await new Promise((resolve) => setTimeout(resolve, 1)); if (!this.endTime) { this.endTime = new Date().toISOString(); } @@ -414,8 +418,8 @@ export class Step extends StepFields { const output = await this.client.store.run( { currentThread: currentStore?.currentThread ?? null, - currentStep: this, - currentExperimentRunId: null + currentExperimentRunId: currentStore?.currentExperimentRunId ?? null, + currentStep: this }, () => cb(this) ); diff --git a/tests/api.test.ts b/tests/api.test.ts index 699e9b2..34a09cb 100644 --- a/tests/api.test.ts +++ b/tests/api.test.ts @@ -97,7 +97,6 @@ describe('End to end tests for the SDK', function () { 'name', { foo: 'bar' }, undefined, - undefined, ['hello'] ); @@ -112,7 +111,6 @@ describe('End to end tests for the SDK', function () { 'test', { foo: 'baz' }, undefined, - undefined, ['hello:world'] ); expect(updatedThread.tags).toStrictEqual(['hello:world']); diff --git a/tests/wrappers.test.ts b/tests/wrappers.test.ts index 0702fbb..01deeb7 100644 --- a/tests/wrappers.test.ts +++ b/tests/wrappers.test.ts @@ -357,7 +357,8 @@ describe('Wrapper', () => { }); let persistedExperimentItem: DatasetExperimentItem | undefined = undefined; - client.experimentRun().wrap(async () => { + + await client.experimentRun().wrap(async () => { const scores = [ { name: 'context_relevancy', @@ -365,7 +366,7 @@ describe('Wrapper', () => { value: 0.6 } ]; - client.step({ name: 'agent', type: 'run' }).wrap(async () => { + await client.step({ name: 'agent', type: 'run' }).wrap(async () => { const experimentItem = { scores: scores, // scores in the format above input: { question: 'question' }, @@ -374,6 +375,10 @@ describe('Wrapper', () => { persistedExperimentItem = await experiment.log(experimentItem); }); }); + expect(persistedExperimentItem).toBeTruthy(); + + await sleep(1000); + const experimentRunId = persistedExperimentItem!.experimentRunId; expect(experimentRunId).toBeTruthy(); const experimentRun = await client.api.getStep(experimentRunId!); From 621504c3937fc0fe6aea383d965bdeda19f829f7 Mon Sep 17 00:00:00 2001 From: Damien BUTY Date: Fri, 26 Jul 2024 17:10:57 +0200 Subject: [PATCH 4/4] fix: tests --- src/instrumentation/openai.ts | 1 + tests/integration/openai.test.ts | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/src/instrumentation/openai.ts b/src/instrumentation/openai.ts index 0f04a68..f5edba6 100644 --- a/src/instrumentation/openai.ts +++ b/src/instrumentation/openai.ts @@ -508,6 +508,7 @@ const processOpenAIOutput = async ( outputTokenCount: output.usage?.completion_tokens, tokenCount: output.usage?.total_tokens }); + if (parent) { const step = parent.step({ name: generation.model || 'openai', diff --git a/tests/integration/openai.test.ts b/tests/integration/openai.test.ts index b1d2443..d68999d 100644 --- a/tests/integration/openai.test.ts +++ b/tests/integration/openai.test.ts @@ -370,6 +370,8 @@ describe('OpenAI Instrumentation', () => { }); }); + await new Promise((resolve) => setTimeout(resolve, 2000)); + const { data: [step] } = await client.api.getSteps({ @@ -411,6 +413,8 @@ describe('OpenAI Instrumentation', () => { }); }); + await new Promise((resolve) => setTimeout(resolve, 3000)); + const { data: [step] } = await client.api.getSteps({