UNPKG

mongodb-rag-core

Version:

Common elements used by MongoDB Chatbot Framework components.

66 lines (61 loc) 3.28 kB
"use strict"; Object.defineProperty(exports, "__esModule", { value: true }); exports.makeGenerateAnnotatedCollectionSchema = void 0; const zod_1 = require("zod"); const getOpenAiFunctionResponse_1 = require("./getOpenAiFunctionResponse"); const prettyPrintMongoDbDocument_1 = require("../prettyPrintMongoDbDocument"); const braintrust_1 = require("braintrust"); const systemPrompt = `You are an expert MongoDB database architect. Your task is to analyze the provided database metadata and generate clear, concise descriptions and an annotated schema for the specified collection. The descriptions that you generate will be used in the prompt of a LLM for performing database-related tasks. In the annotated schema, create a TypeScript schema with the relevant types. Requirements: 1. Include descriptions in comments next to the fields of each field in the schema. 2. Do not include the top-level description of the collection. This is provided elsewhere. 3. Include raw TypeScript code. DO NOT wrap in a Markdown code block. For each index on the collection, include a comment describing what the index does. You can also briefly note potential use cases for this index. `; const functionName = "get_annotated_collection_schema"; const functionDescription = "Generate annotated schema and index definitions for the specified collection."; function makeDetailedCollectionDescriptionSchema(collectionInfo) { const indexNames = collectionInfo.indexes .map((metadata) => metadata.name) .filter((name) => name !== "undefined"); // Create a record schema where each key must be one of the collection names const indexDescriptionSchema = zod_1.z.array(zod_1.z.object({ name: zod_1.z.enum(indexNames), description: zod_1.z.string(), })); const DetailedCollectionDescriptionsSchema = zod_1.z.object({ typeScriptSchema: zod_1.z .string() .describe("Annotated TypeScript schema for the collection. Annotate with TypeDoc comments."), indexDescriptions: indexDescriptionSchema, }); return DetailedCollectionDescriptionsSchema; } const makeGenerateAnnotatedCollectionSchema = (openAiClient) => (0, braintrust_1.wrapTraced)(async function generateAnnotatedCollectionSchema({ collectionMetadata, databaseMetadata, llm: llmOptions, }) { const messages = [ { role: "system", content: systemPrompt, }, { role: "user", content: `Analyze the following collection: '${collectionMetadata.collectionName}'. Database metadata: ${(0, prettyPrintMongoDbDocument_1.prettyPrintMongoDbDocument)(databaseMetadata)} Again, analyze the collection named '${collectionMetadata.collectionName}'.`, }, ]; const result = await (0, getOpenAiFunctionResponse_1.getOpenAiFunctionResponse)({ messages, llmOptions, schema: makeDetailedCollectionDescriptionSchema(collectionMetadata), functionName, functionDescription, openAiClient, }); return result; }, { name: "generateAnnotatedCollectionSchema", }); exports.makeGenerateAnnotatedCollectionSchema = makeGenerateAnnotatedCollectionSchema; //# sourceMappingURL=generateAnnotatedCollectionSchema.js.map