diff --git a/.gitignore b/.gitignore index c177bb3..95bf6fe 100644 --- a/.gitignore +++ b/.gitignore @@ -15,7 +15,6 @@ dist-ssr *.local .env -examples # Editor directories and files .vscode/* diff --git a/examples/openai.ts b/examples/openai.ts new file mode 100644 index 0000000..7987aea --- /dev/null +++ b/examples/openai.ts @@ -0,0 +1,141 @@ +import 'dotenv/config'; +import OpenAI from 'openai'; + +import { LiteralClient } from '../src'; + +const literalClient = new LiteralClient(); + +const openai = new OpenAI(); + +// Instrument the OpenAI client +literalClient.instrumentation.openai(); + +async function main() { + // This will be automagically logged by Literal AI + await openai.chat.completions.create({ + model: 'gpt-4', + messages: [{ role: 'user', content: 'Say this is a test' }] + }); + + await openai.chat.completions.create({ + model: 'gpt-4', + messages: [ + { + role: 'user', + content: + 'Write a README document in markdown explaining the basic usage of the `with-csv` npm library.' + } + ] + }); + + await literalClient + .run({ + name: 'Test run', + input: { + content: + '{"question": ["how to code chat with pdf and use pdfplumber and ollama local",[]],"images_content": {}}' + }, + output: { + content: `To create a chat application that interacts with PDF files using \`pdfplumber\` and the Ollama local model in Chainlit, you can follow the example below. This example demonstrates how to upload a PDF, extract text using \`pdfplumber\`, and allow users to ask questions about its content using the Ollama model. + + ### Step 1: Install Required Packages + Make sure you have the necessary packages installed: + \`\`\`bash + pip install chainlit pdfplumber langchain + \`\`\` + + ### Step 2: Create the Application Script + Create a Python file named \`app.py\` and add the following code: + + \`\`\`python + import os + import pdfplumber + import chainlit as cl + from langchain.llms.ollama import Ollama + from langchain.prompts import ChatPromptTemplate + from langchain.schema import StrOutputParser + from langchain.schema.runnable import Runnable + from langchain.schema.runnable.config import RunnableConfig + from langchain.text_splitter import RecursiveCharacterTextSplitter + from langchain.memory import ChatMessageHistory, ConversationBufferMemory + from langchain.schema import Document + + # Initialize the Ollama model + model = Ollama(model="llama2") + + @cl.on_chat_start + async def on_chat_start(): + files = await cl.AskFileMessage( + content="Please upload a PDF file to start asking questions.", + accept=["application/pdf"], + max_size_mb=20, + timeout=180, + ).send() + + file = files[0] + docs = process_pdf(file) + + message_history = ChatMessageHistory() + memory = ConversationBufferMemory(memory_key="chat_history", output_key="answer", chat_memory=message_history, return_messages=True) + + prompt = ChatPromptTemplate.from_messages( + [ + ("system", "You're a knowledgeable assistant who provides accurate answers based on the PDF content."), + ("human", "{question}"), + ] + ) + + runnable = prompt | model | StrOutputParser() + cl.user_session.set("runnable", runnable) + + await cl.Message(content="You can now ask questions about the PDF!").send() + + def process_pdf(file): + text = "" + with pdfplumber.open(file.path) as pdf: + for page in pdf.pages: + text += page.extract_text() + "\ + " + + # Split the text into chunks + text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) + docs = text_splitter.split_text(text) + documents = [Document(page_content=chunk) for chunk in docs] + return documents + + @cl.on_message + async def on_message(message: cl.Message): + runnable = cl.user_session.get("runnable") # type: Runnable + msg = cl.Message(content="") + + for chunk in await cl.make_async(runnable.stream)( + {"question": message.content}, + config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]), + ): + await msg.stream_token(chunk) + + await msg.send() + \`\`\` + + ### Step 3: Run the Application + To start the Chainlit application, run the following command in your terminal: + \`\`\`bash + chainlit run app.py + \`\`\` + + ### Step 4: Interact with the Application + Open your browser and navigate to \`http://localhost:8000\`. You can upload a PDF file and start asking questions about its content. + + ### Explanation + - **PDF Processing**: The \`process_pdf\` function uses \`pdfplumber\` to extract text from the uploaded PDF file. + - **Text Splitting**: The extracted text is split into manageable chunks using \`RecursiveCharacterTextSplitter\`. + - **Ollama Model**: The Ollama model is used to generate responses based on the extracted text. + - **Conversational Interface**: Users can ask questions, and the model will respond based on the content of the PDF. + + This setup allows you to create a conversational interface that can answer questions based on the content of a PDF file using \`pdfplumber\` and the Ollama local model.` + } + }) + .send(); +} + +main(); diff --git a/examples/prompt.ts b/examples/prompt.ts new file mode 100644 index 0000000..60b5f24 --- /dev/null +++ b/examples/prompt.ts @@ -0,0 +1,13 @@ +import 'dotenv/config'; + +import { LiteralClient } from '../src'; + +const literalClient = new LiteralClient(); + +async function main() { + const prompt = await literalClient.api.getPrompt(''); + + console.log(prompt); +} + +main(); diff --git a/jest.config.ts b/jest.config.ts index ff38151..8f42aa3 100644 --- a/jest.config.ts +++ b/jest.config.ts @@ -4,6 +4,7 @@ const config: Config.InitialOptions = { verbose: true, transform: { '^.+\\.ts?$': 'ts-jest' - } + }, + testTimeout: 30_000 }; export default config; diff --git a/package-lock.json b/package-lock.json index 14d7e9a..ceef4c7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -43,6 +43,7 @@ }, "peerDependencies": { "@ai-sdk/openai": "0.0.x", + "@langchain/openai": "^0.2.7", "ai": "3.x", "langchain": "0.1.x", "llamaindex": "0.3.x", @@ -4204,17 +4205,16 @@ } }, "node_modules/@langchain/core": { - "version": "0.2.23", - "resolved": "https://registry.npmjs.org/@langchain/core/-/core-0.2.23.tgz", - "integrity": "sha512-elPg6WpAkxWEIGC9u38F2anbzqfYYEy32lJdsd9dtChcHSFmFLlXqa+SnpO3R772gUuJmcu+Pd+fCvmRFy029w==", - "optional": true, + "version": "0.2.28", + "resolved": "https://registry.npmjs.org/@langchain/core/-/core-0.2.28.tgz", + "integrity": "sha512-xN3+UdfxFaBcm29auMHFHGEYRh+3HwBc/dICHtwfk2wTSmw4HzWmBtZMx3BG+TOgh5Et7+mT6eF6E3omDLfk+A==", "peer": true, "dependencies": { "ansi-styles": "^5.0.0", "camelcase": "6", "decamelize": "1.2.0", "js-tiktoken": "^1.0.12", - "langsmith": "~0.1.39", + "langsmith": "^0.1.43", "mustache": "^4.2.0", "p-queue": "^6.6.2", "p-retry": "4", @@ -4230,7 +4230,6 @@ "version": "5.2.0", "resolved": "https://registry.npmjs.org/ansi-styles/-/ansi-styles-5.2.0.tgz", "integrity": "sha512-Cxwpt2SfTzTtXcfOlzGEee8O+c+MmUgGrNiBcXnuWxuFJHe6a5Hz7qwhwe5OgaSYI0IJvkLqWX1ASG+cJOkEiA==", - "optional": true, "peer": true, "engines": { "node": ">=10" @@ -4243,7 +4242,6 @@ "version": "6.3.0", "resolved": "https://registry.npmjs.org/camelcase/-/camelcase-6.3.0.tgz", "integrity": "sha512-Gmy6FhYlCY7uOElZUSbxo2UCDH8owEk996gkbrpsgGtrJLM3J7jGxl9Ic7Qwwj4ivOE5AWZWRMecDdF7hqGjFA==", - "optional": true, "peer": true, "engines": { "node": ">=10" @@ -4260,12 +4258,27 @@ "https://github.com/sponsors/broofa", "https://github.com/sponsors/ctavan" ], - "optional": true, "peer": true, "bin": { "uuid": "dist/bin/uuid" } }, + "node_modules/@langchain/openai": { + "version": "0.2.7", + "resolved": "https://registry.npmjs.org/@langchain/openai/-/openai-0.2.7.tgz", + "integrity": "sha512-f2XDXbExJf4SYsy17QSiq0YY/UWJXhJwoiS8uRi/gBa20zBQ8+bBFRnb9vPdLkOkGiaTy+yXZVFro3a9iW2r3w==", + "peer": true, + "dependencies": { + "@langchain/core": ">=0.2.26 <0.3.0", + "js-tiktoken": "^1.0.12", + "openai": "^4.55.0", + "zod": "^3.22.4", + "zod-to-json-schema": "^3.22.3" + }, + "engines": { + "node": ">=18" + } + }, "node_modules/@llamaindex/env": { "version": "0.1.3", "resolved": "https://registry.npmjs.org/@llamaindex/env/-/env-0.1.3.tgz", @@ -10611,9 +10624,9 @@ "peer": true }, "node_modules/langsmith": { - "version": "0.1.41", - "resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.1.41.tgz", - "integrity": "sha512-8R7s/225Pxmv0ipMfd6sqmWVsfHLQivYlQZ0vx5K+ReoknummTenQlVK8gapk3kqRMnzkrouuRHMhWjMR6RgUA==", + "version": "0.1.43", + "resolved": "https://registry.npmjs.org/langsmith/-/langsmith-0.1.43.tgz", + "integrity": "sha512-+IL59ye/je9HmMttJU50epJneEbEwlMJ8i5tEFjJC6l2+SWPtedT0UPuAnPEybMhfjU3ziNfqAxck7WTEncL8w==", "peer": true, "dependencies": { "@types/uuid": "^9.0.1", diff --git a/package.json b/package.json index 13c2dc9..2ee035a 100644 --- a/package.json +++ b/package.json @@ -59,6 +59,7 @@ }, "peerDependencies": { "@ai-sdk/openai": "0.0.x", + "@langchain/openai": "^0.2.7", "ai": "3.x", "langchain": "0.1.x", "llamaindex": "0.3.x", diff --git a/src/api.ts b/src/api.ts index 470a0e5..8e3ac73 100644 --- a/src/api.ts +++ b/src/api.ts @@ -410,7 +410,6 @@ export class API { variables: variables } }); - if (response.data.errors) { throw new Error(JSON.stringify(response.data.errors)); } @@ -851,21 +850,22 @@ export class API { * @returns A Promise resolving to the newly created `Generation` object. */ async createGeneration(generation: Generation) { - const mutation = ` - mutation CreateGeneration($generation: GenerationPayloadInput!) { - createGeneration(generation: $generation) { - id, - type - } - } - `; - - const variables = { - generation - }; + const stepId = generation.id; + const stepMetadata = generation.metadata; + const stepTags = generation.tags; + + delete generation.id; + + const generationAsStep = this.client.step({ + id: stepId, + metadata: stepMetadata, + tags: stepTags, + generation, + name: generation.type ?? '', + type: 'llm' + }); - const response = await this.makeGqlCall(mutation, variables); - return response.data.createGeneration as PersistedGeneration; + return generationAsStep.send(); } /** @@ -929,13 +929,13 @@ export class API { $metadata: Json, $participantId: String, $tags: [String!], - ) { + ) { upsertThread( - id: $threadId - name: $name - metadata: $metadata - participantId: $participantId - tags: $tags + id: $threadId + name: $name + metadata: $metadata + participantId: $participantId + tags: $tags ) { ${threadFields} } diff --git a/src/evaluation/experiment-item-run.ts b/src/evaluation/experiment-item-run.ts index f71ab0f..565d7a5 100644 --- a/src/evaluation/experiment-item-run.ts +++ b/src/evaluation/experiment-item-run.ts @@ -44,7 +44,10 @@ export class ExperimentItemRun extends Step { ? currentStore?.rootRun : this.type === 'run' ? this - : null + : null, + metadata: currentStore?.metadata ?? null, + tags: currentStore?.tags ?? null, + stepId: currentStore?.stepId ?? null }, async () => { try { diff --git a/src/index.ts b/src/index.ts index 7bde18b..8426fe9 100644 --- a/src/index.ts +++ b/src/index.ts @@ -18,6 +18,9 @@ type StoredContext = { currentStep: Step | null; currentExperimentItemRunId?: string | null; rootRun: Step | null; + metadata: Record | null; + tags: string[] | null; + stepId: string | null; }; /** @@ -217,4 +220,30 @@ export class LiteralClient { return store.rootRun; } + + decorate(options: { + metadata?: Record; + tags?: string[]; + stepId?: string; + }) { + return { + wrap: async (cb: () => T) => { + const currentStore = this.store.getStore(); + + return this.store.run( + { + currentThread: currentStore?.currentThread ?? null, + currentExperimentItemRunId: + currentStore?.currentExperimentItemRunId ?? null, + currentStep: currentStore?.currentStep ?? null, + rootRun: currentStore?.rootRun ?? null, + metadata: options?.metadata ?? null, + tags: options?.tags ?? null, + stepId: options?.stepId ?? null + }, + () => cb() + ); + } + }; + } } diff --git a/src/instrumentation/index.ts b/src/instrumentation/index.ts index 583b040..1a66344 100644 --- a/src/instrumentation/index.ts +++ b/src/instrumentation/index.ts @@ -1,3 +1,5 @@ +import OpenAI from 'openai'; + import { LiteralClient, Maybe } from '..'; import { LiteralCallbackHandler } from './langchain'; import { instrumentLlamaIndex, withThread } from './llamaindex'; @@ -7,6 +9,7 @@ import { makeInstrumentVercelSDK } from './vercel-sdk'; export type OpenAIGlobalOptions = { tags?: Maybe; metadata?: Maybe>; + client?: OpenAI; }; export default (client: LiteralClient) => ({ diff --git a/src/instrumentation/langchain.ts b/src/instrumentation/langchain.ts index 9f73ac2..737228d 100644 --- a/src/instrumentation/langchain.ts +++ b/src/instrumentation/langchain.ts @@ -23,6 +23,7 @@ import { import { ChainValues, InputValues } from '@langchain/core/utils/types'; import mustache from 'mustache'; import { v5 as uuidv5 } from 'uuid'; +import { validate as uuidValidate } from 'uuid'; import { ChatGeneration, @@ -121,6 +122,8 @@ interface ChatGenerationStart { start: number; outputTokenCount: number; ttFirstToken?: number; + metadata?: Record; + tags?: string[]; } interface CompletionGenerationStart { @@ -131,6 +134,8 @@ interface CompletionGenerationStart { start: number; outputTokenCount: number; ttFirstToken?: number; + metadata?: Record; + tags?: string[]; } function convertMessageRole(role: string) { @@ -289,6 +294,29 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { } } + getGenerationStepId(metadata?: Record | undefined) { + const generationStepIdFromMetadata = metadata?.literalaiStepId; + + if (typeof generationStepIdFromMetadata !== 'string') { + return null; + } + + if (!uuidValidate(generationStepIdFromMetadata)) { + return null; + } + + // The stepId from metadata can only be used on one generation + if ( + Object.values(this.steps).find( + (step) => step.id === generationStepIdFromMetadata + ) + ) { + return null; + } + + return generationStepIdFromMetadata; + } + /** * LLM Callbacks */ @@ -317,7 +345,9 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { settings, prompt: prompts[0], start: Date.now(), - outputTokenCount: 0 + outputTokenCount: 0, + metadata, + tags }; const parentId = this.getParentId(parentRunId); @@ -329,7 +359,7 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { type: 'llm', tags: tags, threadId: this.threadId, - id: runId, + id: this.getGenerationStepId(metadata), startTime: new Date().toISOString(), parentId: this.getParentId(parentRunId), metadata: metadata, @@ -375,97 +405,115 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { const completionGeneration = this.completionGenerations[runId]; const chatGeneration = this.chatGenerations[runId]; - if (completionGeneration) { - const { - start, - outputTokenCount, - ttFirstToken, - prompt, - model, - provider, - settings - } = this.completionGenerations[runId]; - const duration = (Date.now() - start) / 1000; - const tokenThroughputInSeconds = - duration && outputTokenCount - ? outputTokenCount / (duration / 1000) - : undefined; - - const generation = new CompletionGeneration({ - provider, - model, - settings, - completion: output.generations[0][0].text, - prompt: prompt, - duration, - ttFirstToken, - outputTokenCount, - tokenThroughputInSeconds: tokenThroughputInSeconds - }); - - if (this.steps[runId]) { - this.steps[runId].generation = generation; - this.steps[runId].output = output.generations[0][0]; - this.steps[runId].endTime = new Date().toISOString(); - - await this.steps[runId].send(); - } else { - await this.client.api.createGeneration(generation); - } - } else if (chatGeneration) { - const { - promptId, - variables, - start, - outputTokenCount, - ttFirstToken, - inputMessages, - model, - provider, - settings, - tools - } = this.chatGenerations[runId]; - - const duration = (Date.now() - start) / 1000; - const tokenThroughputInSeconds = - duration && outputTokenCount - ? outputTokenCount / (duration / 1000) - : undefined; - const messageCompletion = convertMessage( - (output.generations[0][0] as any).message - ); - - const generation = new ChatGeneration({ - promptId, - variables, - provider, - model, - settings, - tools, - messageCompletion, - messages: addToolCallIdToMessages(inputMessages), - duration, - ttFirstToken, - outputTokenCount, - tokenThroughputInSeconds: tokenThroughputInSeconds - }); - - if (this.steps[runId]) { - this.steps[runId].generation = generation; - this.steps[runId].generation!.inputTokenCount = - output.llmOutput?.estimatedTokenUsage?.promptTokens; - this.steps[runId].generation!.outputTokenCount = - output.llmOutput?.estimatedTokenUsage?.completionTokens; - this.steps[runId].generation!.tokenCount = - output.llmOutput?.estimatedTokenUsage?.totalTokens; - - this.steps[runId].output = messageCompletion; - this.steps[runId].endTime = new Date().toISOString(); - - await this.steps[runId].send(); - } else { - await this.client.api.createGeneration(generation); + try { + if (completionGeneration) { + const { + start, + outputTokenCount, + ttFirstToken, + prompt, + model, + provider, + settings, + metadata, + tags + } = this.completionGenerations[runId]; + const duration = (Date.now() - start) / 1000; + const tokenThroughputInSeconds = + duration && outputTokenCount + ? outputTokenCount / (duration / 1000) + : undefined; + + const generation = new CompletionGeneration({ + metadata, + tags, + provider, + model, + settings, + completion: output.generations[0][0].text, + prompt: prompt, + duration, + ttFirstToken, + outputTokenCount, + tokenThroughputInSeconds: tokenThroughputInSeconds + }); + + if (this.steps[runId]) { + this.steps[runId].generation = generation; + this.steps[runId].output = output.generations[0][0]; + this.steps[runId].endTime = new Date().toISOString(); + + await this.steps[runId].send(); + } else { + await this.client.api.createGeneration({ + ...generation, + id: this.getGenerationStepId(metadata) + }); + } + } else if (chatGeneration) { + const { + promptId, + variables, + start, + outputTokenCount, + ttFirstToken, + inputMessages, + model, + provider, + settings, + tools, + metadata, + tags + } = this.chatGenerations[runId]; + + const duration = (Date.now() - start) / 1000; + const tokenThroughputInSeconds = + duration && outputTokenCount + ? outputTokenCount / (duration / 1000) + : undefined; + const messageCompletion = convertMessage( + (output.generations[0][0] as any).message + ); + + const generation = new ChatGeneration({ + metadata, + tags, + promptId, + variables, + provider, + model, + settings, + tools, + messageCompletion, + messages: addToolCallIdToMessages(inputMessages), + duration, + ttFirstToken, + outputTokenCount, + tokenThroughputInSeconds: tokenThroughputInSeconds + }); + + if (this.steps[runId]) { + this.steps[runId].generation = generation; + this.steps[runId].generation!.inputTokenCount = + output.llmOutput?.estimatedTokenUsage?.promptTokens; + this.steps[runId].generation!.outputTokenCount = + output.llmOutput?.estimatedTokenUsage?.completionTokens; + this.steps[runId].generation!.tokenCount = + output.llmOutput?.estimatedTokenUsage?.totalTokens; + + this.steps[runId].output = messageCompletion; + this.steps[runId].endTime = new Date().toISOString(); + + await this.steps[runId].send(); + } else { + await this.client.api.createGeneration({ + ...generation, + id: this.getGenerationStepId(metadata) + }); + } } + } catch (e) { + console.log('Error in handleLLMEnd', e); } } @@ -505,7 +553,9 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { tools, inputMessages: messageList.map(convertMessage), start: Date.now(), - outputTokenCount: 0 + outputTokenCount: 0, + metadata, + tags }; const parentId = this.getParentId(parentRunId); @@ -523,10 +573,10 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { type: 'llm', tags: tags, threadId: this.threadId, - id: runId, + id: this.getGenerationStepId(metadata), startTime: new Date().toISOString(), parentId: parentId, - metadata: metadata, + metadata, input: { content: messages[0] } }) .send(); @@ -598,12 +648,12 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { const step = await this.client .run({ name: name || chainType, - tags: tags, threadId: this.threadId, id: runId, input: stepInput, startTime: new Date().toISOString(), - metadata: metadata + metadata, + tags }) .send(); @@ -623,6 +673,7 @@ export class LiteralCallbackHandler extends BaseCallbackHandler { type: 'tool', parentId, tags: tags, + metadata, threadId: this.threadId, id: runId, input: stepInput, diff --git a/src/instrumentation/openai.ts b/src/instrumentation/openai.ts index 0541fae..82e874f 100644 --- a/src/instrumentation/openai.ts +++ b/src/instrumentation/openai.ts @@ -26,6 +26,7 @@ type OriginalFunction = ( type OpenAICallOptions = { literalaiTags?: Maybe; literalaiMetadata?: Maybe>; + literalaiStepId?: Maybe; }; function cleanOpenAIArgs( @@ -87,8 +88,11 @@ function instrumentOpenAI( client: LiteralClient, options: OpenAIGlobalOptions = {} ) { + const originalMethods = (OpenAI.prototype as any).__literalai_originalMethods; + // Patching the chat.completions.create function const originalChatCompletionsCreate = + originalMethods?.originalChatCompletionsCreate ?? OpenAI.Chat.Completions.prototype.create; const wrappedChatCompletionsCreate = wrapFunction( originalChatCompletionsCreate, @@ -97,7 +101,9 @@ function instrumentOpenAI( ); // Patching the completions.create function - const originalCompletionsCreate = OpenAI.Completions.prototype.create; + const originalCompletionsCreate = + originalMethods?.originalCompletionsCreate ?? + OpenAI.Completions.prototype.create; const wrappedCompletionsCreate = wrapFunction( originalCompletionsCreate, client, @@ -105,7 +111,8 @@ function instrumentOpenAI( ); // Patching the images.generate function - const originalImagesGenerate = OpenAI.Images.prototype.generate; + const originalImagesGenerate = + originalMethods?.originalImagesGenerate ?? OpenAI.Images.prototype.generate; const wrappedImagesGenerate = wrapFunction( originalImagesGenerate, client, @@ -117,16 +124,26 @@ function instrumentOpenAI( OpenAI.Completions.prototype.create = wrappedCompletionsCreate as any; OpenAI.Images.prototype.generate = wrappedImagesGenerate as any; + (OpenAI.prototype as any).__literalai_originalMethods = { + originalChatCompletionsCreate, + originalCompletionsCreate, + originalImagesGenerate + }; + return { + ...options.client, chat: { completions: { + _client: options.client, create: wrappedChatCompletionsCreate } }, completions: { + _client: options.client, create: wrappedCompletionsCreate }, images: { + _client: options.client, generate: wrappedImagesGenerate } }; @@ -357,10 +374,12 @@ const processOpenAIOutput = async ( }; const baseGeneration = { + ...(callOptions?.literalaiStepId && { id: callOptions.literalaiStepId }), provider: 'openai', model: inputs.model, settings: getSettings(inputs), - tags + tags, + metadata }; const threadFromStore = client._currentThread(); diff --git a/src/instrumentation/vercel-sdk.ts b/src/instrumentation/vercel-sdk.ts index ec68074..01409ca 100644 --- a/src/instrumentation/vercel-sdk.ts +++ b/src/instrumentation/vercel-sdk.ts @@ -15,10 +15,9 @@ import { IGenerationMessage, ILLMSettings, ITool, - LiteralClient + LiteralClient, + Maybe } from '..'; -import { Step } from '../observability/step'; -import { Thread } from '../observability/thread'; export type VercelLanguageModel = LanguageModel; @@ -188,6 +187,10 @@ const computeMetricsStream = async ( let outputTokenCount = 0; let ttFirstToken: number | undefined = undefined; + + let accumulatedStreamObjectResponse: IGenerationMessage | undefined = + undefined; + for await (const chunk of stream as unknown as AsyncIterable) { if (typeof chunk === 'string') { textMessage.content += chunk; @@ -222,6 +225,14 @@ const computeMetricsStream = async ( }); break; } + case 'object': { + const { object } = chunk as any; + accumulatedStreamObjectResponse = { + role: 'assistant', + content: JSON.stringify(object) + }; + break; + } } } @@ -231,6 +242,10 @@ const computeMetricsStream = async ( outputTokenCount += 1; } + if (accumulatedStreamObjectResponse) { + messages.push(accumulatedStreamObjectResponse); + } + const duration = (Date.now() - startTime) / 1000; const tokenThroughputInSeconds = duration && outputTokenCount @@ -260,7 +275,9 @@ const computeMetricsStream = async ( }; type VercelExtraOptions = { - literalAiParent?: Step | Thread; + literalaiTags?: Maybe; + literalaiMetadata?: Maybe>; + literalaiStepId?: Maybe; }; export type InstrumentationVercelMethod = { @@ -288,7 +305,7 @@ export const makeInstrumentVercelSDK = ( type TOptions = Options; type TResult = Result; - return async (options: TOptions): Promise => { + return async (options: TOptions & VercelExtraOptions): Promise => { const startTime = Date.now(); const result: TResult = await (fn as any)(options); @@ -318,6 +335,9 @@ export const makeInstrumentVercelSDK = ( ); const generation = new ChatGeneration({ + ...(options.literalaiStepId && { id: options.literalaiStepId }), + metadata: options.literalaiMetadata, + tags: options.literalaiTags, provider: options.model.provider, model: options.model.modelId, settings: extractSettings(options), @@ -346,6 +366,9 @@ export const makeInstrumentVercelSDK = ( const metrics = computeMetricsSync(options, result, startTime); const generation = new ChatGeneration({ + ...(options.literalaiStepId && { id: options.literalaiStepId }), + metadata: options.literalaiMetadata, + tags: options.literalaiTags, provider: options.model.provider, model: options.model.modelId, settings: extractSettings(options), diff --git a/src/observability/generation.ts b/src/observability/generation.ts index 07f6a09..297b233 100644 --- a/src/observability/generation.ts +++ b/src/observability/generation.ts @@ -52,6 +52,7 @@ export class BaseGeneration extends Utils { model?: Maybe; id?: Maybe; tags?: Maybe; + metadata?: Maybe>; error?: Maybe; variables?: Maybe>; settings?: Maybe; @@ -72,6 +73,7 @@ export class CompletionGeneration extends BaseGeneration { constructor(data: OmitUtils) { super(); this.type = 'COMPLETION'; + Object.assign(this, data); } } diff --git a/src/observability/step.ts b/src/observability/step.ts index ff5e4f8..e164753 100644 --- a/src/observability/step.ts +++ b/src/observability/step.ts @@ -3,7 +3,14 @@ import { v4 as uuidv4 } from 'uuid'; import { LiteralClient } from '..'; import { API } from '../api'; import { Score } from '../evaluation/score'; -import { Environment, Maybe, OmitUtils, Utils, isPlainObject } from '../utils'; +import { + Environment, + Maybe, + OmitUtils, + Utils, + isPlainObject, + omitLiteralAiMetadata +} from '../utils'; import { Attachment } from './attachment'; import { Generation } from './generation'; @@ -64,21 +71,13 @@ export class Step extends StepFields { this.client = client; Object.assign(this, data); + this.enrichFromStore(ignoreContext); // Automatically generate an ID if not provided. if (!this.id) { this.id = uuidv4(); } - if (ignoreContext) { - return; - } - - // Automatically assign parent thread & step & rootRun if there are any in the store. - this.threadId = this.threadId ?? this.client._currentThread()?.id; - this.parentId = this.parentId ?? this.client._currentStep()?.id; - this.rootRunId = this.rootRunId ?? this.client._rootRun()?.id; - // Set the creation and start time to the current time if not provided. if (!this.createdAt) { this.createdAt = new Date().toISOString(); @@ -93,6 +92,37 @@ export class Step extends StepFields { } } + private enrichFromStore(ignoreContext?: true) { + if (ignoreContext) { + return; + } + + const currentStore = this.client.store.getStore(); + + if (currentStore) { + if (currentStore.metadata) { + this.metadata = omitLiteralAiMetadata({ + ...this.metadata, + ...currentStore.metadata + }); + } + + if (currentStore.tags) { + this.tags = [...(this.tags ?? []), ...currentStore.tags]; + } + + if (currentStore.stepId && !this.id) { + this.id = currentStore.stepId; + currentStore.stepId = null; + } + } + + // Automatically assign parent thread & step & rootRun if there are any in the store. + this.threadId = this.threadId ?? this.client._currentThread()?.id; + this.parentId = this.parentId ?? this.client._currentStep()?.id; + this.rootRunId = this.rootRunId ?? this.client._rootRun()?.id; + } + /** * Serializes the step instance, converting complex objects to strings as necessary. * @returns A serialized representation of the step. @@ -174,7 +204,10 @@ export class Step extends StepFields { ? currentStore?.rootRun : this.type === 'run' ? this - : null + : null, + metadata: currentStore?.metadata ?? null, + tags: currentStore?.tags ?? null, + stepId: currentStore?.stepId ?? null }, () => cb(this) ); diff --git a/src/observability/thread.ts b/src/observability/thread.ts index 4defc13..05f7bc4 100644 --- a/src/observability/thread.ts +++ b/src/observability/thread.ts @@ -2,7 +2,13 @@ import { v4 as uuidv4 } from 'uuid'; import { LiteralClient } from '..'; import { API } from '../api'; -import { Environment, Maybe, OmitUtils, Utils } from '../utils'; +import { + Environment, + Maybe, + OmitUtils, + Utils, + omitLiteralAiMetadata +} from '../utils'; import { Step, StepConstructor } from './step'; /** @@ -55,6 +61,24 @@ export class Thread extends ThreadFields { } Object.assign(this, data); + + this.enrichFromStore(); + } + + private enrichFromStore() { + const currentStore = this.client.store.getStore(); + + if (currentStore) { + if (currentStore.metadata) { + this.metadata = omitLiteralAiMetadata({ + ...this.metadata, + ...currentStore.metadata + }); + } + if (currentStore.tags) { + this.tags = [...(this.tags ?? []), ...currentStore.tags]; + } + } } /** @@ -117,7 +141,10 @@ export class Thread extends ThreadFields { currentExperimentItemRunId: currentStore?.currentExperimentItemRunId ?? null, currentStep: null, - rootRun: null + rootRun: null, + metadata: currentStore?.metadata ?? null, + tags: currentStore?.tags ?? null, + stepId: currentStore?.stepId ?? null }, () => cb(this) ); diff --git a/src/utils.ts b/src/utils.ts index d04e5ff..b6df629 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -79,3 +79,8 @@ export class User extends Utils { Object.assign(this, data); } } + +export function omitLiteralAiMetadata(obj: Record) { + const { literalaiTags, literalaiMetadata, literalaiStepId, ...rest } = obj; + return rest; +} diff --git a/tests/api.test.ts b/tests/api.test.ts index ede3b50..97cf2fc 100644 --- a/tests/api.test.ts +++ b/tests/api.test.ts @@ -4,6 +4,7 @@ import { v4 as uuidv4 } from 'uuid'; import { ChatGeneration, LiteralClient } from '../src'; import { Dataset } from '../src/evaluation/dataset'; import { Score } from '../src/evaluation/score'; +import { sleep } from './utils'; describe('End to end tests for the SDK', function () { let client: LiteralClient; @@ -44,7 +45,7 @@ describe('End to end tests for the SDK', function () { const deletedUser = await client.api.getUser(identifier); expect(deletedUser).toBeUndefined(); - }, 30000); + }); it('should test generation', async function () { const generation = await client.api.createGeneration({ @@ -91,7 +92,7 @@ describe('End to end tests for the SDK', function () { const deletedThread = await client.api.getThread(thread.id); expect(deletedThread).toBeNull(); - }, 30000); + }); it('should test thread (deprecated)', async function () { const thread = await client.api.upsertThread( @@ -119,6 +120,9 @@ describe('End to end tests for the SDK', function () { await client.api.deleteThread(thread.id); + // We have to await 5 seconds for the thread to disappear from the cache + await sleep(5000); + const deletedThread = await client.api.getThread(thread.id); expect(deletedThread).toBeNull(); }); @@ -175,7 +179,7 @@ describe('End to end tests for the SDK', function () { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const fetchedStep = await client.api.getStep(step.id!); expect(fetchedStep?.id).toBe(step.id); @@ -209,7 +213,7 @@ describe('End to end tests for the SDK', function () { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const fetchedStep = await client.api.getStep(step.id!); expect(fetchedStep?.id).toBe(step.id); @@ -222,7 +226,7 @@ describe('End to end tests for the SDK', function () { const deletedStep = await client.api.getStep(step.id!); expect(deletedStep).toBeNull(); - }, 30000); + }); it('should test steps', async function () { const thread = await client.thread({ id: uuidv4() }); @@ -237,7 +241,7 @@ describe('End to end tests for the SDK', function () { expect(step.id).not.toBeNull(); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const steps = await client.api.getSteps({ filters: [ @@ -269,7 +273,7 @@ describe('End to end tests for the SDK', function () { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const score = await client.api.createScore({ stepId: step.id!, @@ -309,7 +313,7 @@ describe('End to end tests for the SDK', function () { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 1000)); + await sleep(1000); const firstScoreValue = 0.9234; const scores = await client.api.createScores([ @@ -527,7 +531,7 @@ describe('End to end tests for the SDK', function () { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 1000)); + await sleep(1000); const datasetItem = await dataset.addStep(step.id!); diff --git a/tests/attachments.test.ts b/tests/attachments.test.ts index 3db3135..44ccce2 100644 --- a/tests/attachments.test.ts +++ b/tests/attachments.test.ts @@ -3,6 +3,7 @@ import { createReadStream, readFileSync } from 'fs'; import { LiteralClient, Maybe } from '../src'; import { Attachment } from '../src/observability/attachment'; +import { sleep } from './utils'; const apiUrl = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; @@ -48,7 +49,7 @@ describe('Attachments', () => { }) .send(); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const fetchedStep = await client.api.getStep(step.id!); @@ -86,7 +87,7 @@ describe('Attachments', () => { }); }); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const fetchedStep = await client.api.getStep(stepId!); diff --git a/tests/decorate.test.ts b/tests/decorate.test.ts new file mode 100644 index 0000000..46b2820 --- /dev/null +++ b/tests/decorate.test.ts @@ -0,0 +1,174 @@ +import { openai } from '@ai-sdk/openai'; +import { ChatOpenAI } from '@langchain/openai'; +import { generateText } from 'ai'; +import 'dotenv/config'; +import { SimpleChatEngine } from 'llamaindex'; +import OpenAI from 'openai'; +import { v4 as uuidv4 } from 'uuid'; + +import { LiteralClient, Maybe } from '../src'; +import { sleep } from './utils'; + +const url = process.env.LITERAL_API_URL; +const apiKey = process.env.LITERAL_API_KEY; + +if (!url || !apiKey) { + throw new Error('Missing environment variables'); +} + +const client = new LiteralClient({ apiKey, apiUrl: url }); + +describe('Decorator', () => { + describe('Manual logging', () => { + it('adds metadata and tags to everything logged inside the wrapper', async () => { + let threadId: Maybe; + let stepId: Maybe; + const metadata = { key: 'value' }; + const tags = ['tag1', 'tag2']; + + await client.decorate({ metadata, tags }).wrap(async () => { + const createdThread = await client + .thread({ name: 'Test thread' }) + .upsert(); + + const createdStep = await createdThread + .step({ name: 'Test step', type: 'assistant_message' }) + .send(); + + threadId = createdThread.id; + stepId = createdStep.id; + }); + + await sleep(1000); + + const thread = await client.api.getThread(threadId!); + const step = await client.api.getStep(stepId!); + + expect(thread?.metadata).toEqual(expect.objectContaining(metadata)); + expect(thread?.tags).toEqual(expect.arrayContaining(tags)); + expect(step?.metadata).toEqual(expect.objectContaining(metadata)); + expect(step?.tags).toEqual(expect.arrayContaining(tags)); + }); + + it('creates the first step with the provided ID', async () => { + const stepId = uuidv4(); + let generatedFirstStepId: Maybe; + let generatedSecondStepId: Maybe; + + await client.decorate({ stepId }).wrap(async () => { + const firstStep = await client.run({ name: 'First step' }).send(); + generatedFirstStepId = firstStep.id; + + const secondStep = await client.run({ name: 'Second step' }).send(); + generatedSecondStepId = secondStep.id; + }); + + expect(generatedFirstStepId).toBe(stepId); + expect(generatedSecondStepId).not.toBe(stepId); + }); + }); + + // Skip for the CI + describe.skip('Integrations', () => { + it('logs Langchain generations with the given ID, metadata and tags', async () => { + const cb = client.instrumentation.langchain.literalCallback(); + const model = new ChatOpenAI({}); + + const stepId = uuidv4(); + const metadata = { key: 'value' }; + const tags = ['tag1', 'tag2']; + + await client.decorate({ stepId, metadata, tags }).wrap(async () => { + await model.invoke('Hello, how are you?', { + callbacks: [cb] + }); + }); + + await sleep(2000); + + const step = await client.api.getStep(stepId); + + expect(step?.type).toBe('llm'); + expect(step?.id).toBe(stepId); + expect(step?.metadata).toEqual(expect.objectContaining(metadata)); + expect(step?.tags).toEqual(expect.arrayContaining(tags)); + }); + + it('logs LlamaIndex generations with the given ID, metadata and tags', async () => { + client.instrumentation.llamaIndex.instrument(); + const engine = new SimpleChatEngine(); + + const stepId = uuidv4(); + const metadata = { key: 'value' }; + const tags = ['tag1', 'tag2']; + + await client.decorate({ stepId, metadata, tags }).wrap(async () => { + await engine.chat({ + message: 'Write a vegetarian lasagna recipe for 4 people.' + }); + }); + + await sleep(1000); + + const step = await client.api.getStep(stepId); + + expect(step?.type).toBe('llm'); + expect(step?.id).toBe(stepId); + expect(step?.metadata).toEqual(expect.objectContaining(metadata)); + expect(step?.tags).toEqual(expect.arrayContaining(tags)); + }); + + it('logs OpenAI generations with the given ID, metadata and tags', async () => { + const openai = new OpenAI(); + client.instrumentation.openai(); + + const stepId = uuidv4(); + const metadata = { key: 'value' }; + const tags = ['tag1', 'tag2']; + + await client.decorate({ stepId, metadata, tags }).wrap(async () => { + await openai.chat.completions.create({ + model: 'gpt-3.5-turbo', + messages: [ + { role: 'system', content: 'You are a helpful assistant.' }, + { role: 'user', content: 'What is the capital of Canada?' } + ] + }); + }); + + await sleep(1000); + + const step = await client.api.getStep(stepId); + + expect(step?.type).toBe('llm'); + expect(step?.id).toBe(stepId); + expect(step?.metadata).toEqual(expect.objectContaining(metadata)); + expect(step?.tags).toEqual(expect.arrayContaining(tags)); + }); + + it('logs Vercel AI SDK generations with the given ID, metadata and tags', async () => { + const generateTextWithLiteralAI = + client.instrumentation.vercel.instrument(generateText); + + const stepId = uuidv4(); + const metadata = { key: 'value' }; + const tags = ['tag1', 'tag2']; + + await client.decorate({ stepId, metadata, tags }).wrap(async () => { + await generateTextWithLiteralAI({ + model: openai('gpt-3.5-turbo'), + prompt: 'Write a vegetarian lasagna recipe for 4 people.' + }); + }); + + await sleep(1000); + + const step = await client.api.getStep(stepId); + + expect(step?.type).toBe('llm'); + expect(step?.id).toBe(stepId); + expect(step?.metadata).toEqual(expect.objectContaining(metadata)); + expect(step?.tags).toEqual(expect.arrayContaining(tags)); + }); + }); +}); diff --git a/tests/integration/langchain.test.ts b/tests/integration/langchain.test.ts new file mode 100644 index 0000000..95be837 --- /dev/null +++ b/tests/integration/langchain.test.ts @@ -0,0 +1,60 @@ +import { ChatOpenAI } from '@langchain/openai'; +import 'dotenv/config'; +import { v4 as uuidv4 } from 'uuid'; + +import { LiteralClient } from '../../src'; +import { sleep } from '../utils'; + +const url = process.env.LITERAL_API_URL; +const apiKey = process.env.LITERAL_API_KEY; + +if (!url || !apiKey) { + throw new Error('Missing environment variables'); +} + +const client = new LiteralClient({ apiKey, apiUrl: url }); +const cb = client.instrumentation.langchain.literalCallback(); + +describe.skip('Langchain integration', function () { + it('should create a generation with the provided id', async function () { + const literalaiStepId = uuidv4(); + const model = new ChatOpenAI({}); + + await model.invoke('Hello, how are you?', { + callbacks: [cb], + metadata: { literalaiStepId } + }); + + await sleep(1000); + + const step = await client.api.getStep(literalaiStepId); + + expect(step!.type).toBe('llm'); + }); + + it('should copy tags and metadata to the generation', async function () { + const literalaiStepId = uuidv4(); + const model = new ChatOpenAI({}); + + const metadata = { + framework: 'Langchain', + awesome: 'yes', + literalaiStepId + }; + + const tags = ['bim', 'bam', 'boom']; + + await model.invoke('Hello, how are you?', { + callbacks: [cb], + metadata, + tags + }); + + await sleep(1000); + + const step = await client.api.getStep(literalaiStepId); + + expect(step!.metadata).toEqual(expect.objectContaining(metadata)); + expect(step!.tags).toEqual(expect.arrayContaining(tags)); + }); +}); diff --git a/tests/integration/llamaindex.test.ts b/tests/integration/llamaindex.test.ts index 2c7112e..58fb55b 100644 --- a/tests/integration/llamaindex.test.ts +++ b/tests/integration/llamaindex.test.ts @@ -11,6 +11,7 @@ import { import { resolve } from 'path'; import { LiteralClient } from '../../src'; +import { sleep } from '../utils'; describe('Llama Index Instrumentation', () => { let client: LiteralClient; @@ -119,7 +120,7 @@ describe('Llama Index Instrumentation', () => { expect(response).toBeTruthy(); // Sending message is done asynchronously - await new Promise((resolve) => setTimeout(resolve, 10)); + await sleep(1000); expect(spy).toHaveBeenCalledWith([ expect.objectContaining({ @@ -168,7 +169,7 @@ describe('Llama Index Instrumentation', () => { expect(response).toBeTruthy(); // Sending message is done asynchronously - await new Promise((resolve) => setTimeout(resolve, 10)); + await sleep(1000); expect(spy).toHaveBeenCalledWith([ expect.objectContaining({ diff --git a/tests/integration/openai.test.ts b/tests/integration/openai.test.ts index 59fdad7..9dfaa42 100644 --- a/tests/integration/openai.test.ts +++ b/tests/integration/openai.test.ts @@ -5,6 +5,7 @@ import { v4 as uuidv4 } from 'uuid'; import { ChatGeneration, LiteralClient, Maybe, OmitUtils } from '../../src'; import { Step } from '../../src/observability/step'; +import { sleep } from '../utils'; const apiUrl = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; @@ -16,7 +17,7 @@ if (!apiUrl || !apiKey) { const openai = new OpenAI({ apiKey: 'an-ocean-of-noise' }); // Skip for the CI -describe('OpenAI Instrumentation', () => { +describe.skip('OpenAI Instrumentation', () => { // Mock OpenAI Calls beforeAll(() => { /* @ts-expect-error the mock is incomplete but that's OK */ @@ -214,7 +215,7 @@ describe('OpenAI Instrumentation', () => { n: 1 }); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(1000); const { data: [step] } = await client.api.getSteps({ @@ -227,7 +228,7 @@ describe('OpenAI Instrumentation', () => { expect(step?.type).toBe('run'); expect(step?.output?.data[0].url).toEqual(response.data[0].url); - }, 30000); + }); }); }); @@ -256,7 +257,7 @@ describe('OpenAI Instrumentation', () => { }); }); - await new Promise((resolve) => setTimeout(resolve, 2000)); + await sleep(2000); const { data: [step] @@ -267,7 +268,7 @@ describe('OpenAI Instrumentation', () => { expect(step?.threadId).toBe(threadId); expect(step?.parentId).toBe(parentId); - }, 30_000); + }); it("doesn't mix up threads and steps", async () => { const testId = uuidv4(); @@ -341,10 +342,39 @@ describe('OpenAI Instrumentation', () => { expect(firstGeneration?.parentId).toEqual(firstStep?.id); expect(secondGeneration?.threadId).toEqual(secondThreadId); expect(secondGeneration?.parentId).toEqual(secondStep?.id); - }, 30_000); + }); }); describe('Handling tags and metadata', () => { + it('should assign a specific ID to the generation if provided', async () => { + const openai_ = new OpenAI({ apiKey: 'an-ocean-of-noise' }); + + const client = new LiteralClient({ apiKey, apiUrl }); + const literalaiStepId = uuidv4(); + + const instrumentedOpenAi = client.instrumentation.openai({ + client: openai_ + }); + + await instrumentedOpenAi.chat.completions.create( + { + model: 'gpt-3.5-turbo', + messages: [ + { role: 'system', content: 'You are a helpful assistant.' }, + { role: 'user', content: 'What is the capital of Canada?' } + ] + }, + { literalaiStepId } + ); + + await sleep(1000); + + const step = await client.api.getStep(literalaiStepId); + + expect(step!.id).toEqual(literalaiStepId); + expect(step!.type).toEqual('llm'); + }); + it('handles tags and metadata on the instrumentation call', async () => { const client = new LiteralClient({ apiKey, apiUrl }); client.instrumentation.openai({ @@ -368,7 +398,7 @@ describe('OpenAI Instrumentation', () => { }); }); - await new Promise((resolve) => setTimeout(resolve, 4000)); + await sleep(4000); const { data: [step] @@ -379,12 +409,13 @@ describe('OpenAI Instrumentation', () => { expect(step!.tags).toEqual(expect.arrayContaining(['tag1', 'tag2'])); expect(step!.metadata).toEqual({ key: 'value' }); - }, 30_000); + }); it('handles tags and metadata on the LLM call', async () => { const client = new LiteralClient({ apiKey, apiUrl }); const instrumentedOpenAi = client.instrumentation.openai({ + client: openai, tags: ['tag1', 'tag2'], metadata: { key: 'value' } }); @@ -411,7 +442,7 @@ describe('OpenAI Instrumentation', () => { }); }); - await new Promise((resolve) => setTimeout(resolve, 5000)); + await sleep(5000); const { data: [step] @@ -425,6 +456,6 @@ describe('OpenAI Instrumentation', () => { ); expect(step!.metadata!.key).toEqual('value'); expect(step!.metadata!.otherKey).toEqual('otherValue'); - }, 30_000); + }); }); }); diff --git a/tests/integration/vercel-sdk.test.ts b/tests/integration/vercel-sdk.test.ts index e208dd0..76e88b6 100644 --- a/tests/integration/vercel-sdk.test.ts +++ b/tests/integration/vercel-sdk.test.ts @@ -1,26 +1,24 @@ import { openai } from '@ai-sdk/openai'; import { generateObject, generateText, streamObject, streamText } from 'ai'; import 'dotenv/config'; +import { v4 as uuidv4 } from 'uuid'; import { z } from 'zod'; import { LiteralClient } from '../../src'; +import { sleep } from '../utils'; -describe('Vercel SDK Instrumentation', () => { - let client: LiteralClient; +const apiUrl = process.env.LITERAL_API_URL; +const apiKey = process.env.LITERAL_API_KEY; - beforeAll(function () { - const apiUrl = process.env.LITERAL_API_URL; - const apiKey = process.env.LITERAL_API_KEY; +if (!apiUrl || !apiKey) { + throw new Error('Missing environment variables'); +} - if (!apiUrl || !apiKey) { - throw new Error('Missing environment variables'); - } +const client = new LiteralClient({ apiKey, apiUrl }); - client = new LiteralClient({ apiKey, apiUrl }); - }); - - // Skip for the CI - describe.skip('With OpenAI', () => { +// Skip for the CI +describe.skip('Vercel SDK Instrumentation', () => { + describe('With OpenAI', () => { afterEach(() => jest.restoreAllMocks()); it('should work a simple text generation', async () => { @@ -216,7 +214,7 @@ describe('Vercel SDK Instrumentation', () => { expect(result.text).toBeTruthy(); // Sending message is done asynchronously - await new Promise((resolve) => setTimeout(resolve, 10)); + await sleep(1000); expect(spy).toHaveBeenCalledWith([ expect.objectContaining({ @@ -431,4 +429,47 @@ describe('Vercel SDK Instrumentation', () => { ); }); }); + + describe('Literal AI metadata', () => { + const generateTextWithLiteralAI = + client.instrumentation.vercel.instrument(generateText); + + it('should create a generation with the provided ID', async () => { + const literalaiStepId = uuidv4(); + + await generateTextWithLiteralAI({ + model: openai('gpt-3.5-turbo'), + prompt: 'Write a vegetarian lasagna recipe for 4 people.', + literalaiStepId + }); + + await sleep(1000); + + const step = await client.api.getStep(literalaiStepId); + + expect(step!.id).toEqual(literalaiStepId); + expect(step!.type).toEqual('llm'); + }); + + it('should create a generation with the provided tags and metadata', async () => { + const literalaiStepId = uuidv4(); + + await generateTextWithLiteralAI({ + model: openai('gpt-3.5-turbo'), + prompt: 'Write a vegetarian lasagna recipe for 4 people.', + literalaiStepId, + literalaiTags: ['tag1', 'tag2'], + literalaiMetadata: { otherKey: 'otherValue' } + }); + + await sleep(1000); + + const step = await client.api.getStep(literalaiStepId); + + expect(step!.metadata).toEqual( + expect.objectContaining({ otherKey: 'otherValue' }) + ); + expect(step!.tags).toEqual(expect.arrayContaining(['tag1', 'tag2'])); + }); + }); }); diff --git a/tests/utils.ts b/tests/utils.ts new file mode 100644 index 0000000..421bda0 --- /dev/null +++ b/tests/utils.ts @@ -0,0 +1,3 @@ +export function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} diff --git a/tests/wrappers.test.ts b/tests/wrappers.test.ts index d4c28ee..16f1bb8 100644 --- a/tests/wrappers.test.ts +++ b/tests/wrappers.test.ts @@ -3,6 +3,7 @@ import 'dotenv/config'; import { LiteralClient, Maybe } from '../src'; import { DatasetExperimentItem } from '../src/evaluation/dataset'; import { Step } from '../src/observability/step'; +import { sleep } from './utils'; const url = process.env.LITERAL_API_URL; const apiKey = process.env.LITERAL_API_KEY; @@ -13,10 +14,6 @@ if (!url || !apiKey) { const client = new LiteralClient({ apiKey, apiUrl: url }); -function sleep(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - describe('Wrapper', () => { it('handles failing step', async () => { let threadId: Maybe;