get-started.js

import "dotenv/config";
import { AzureCosmosDBNoSQLVectorStore, AzureCosmosDBNoSQLSearchType } from "@langchain/azure-cosmosdb";
import { AzureChatOpenAI, AzureOpenAIEmbeddings } from "@langchain/openai";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { RunnableSequence } from "@langchain/core/runnables";
import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
import { Document } from "@langchain/core/documents";
import { readFileSync } from "node:fs";

// Azure OpenAI configuration from environment variables
const azureOpenAIConfig = {
  azureOpenAIApiKey: process.env.AZURE_OPENAI_API_KEY,
  azureOpenAIBasePath: process.env.AZURE_OPENAI_ENDPOINT,
  azureOpenAIApiVersion: process.env.AZURE_OPENAI_API_VERSION,
};

// Cosmos DB connection string is read automatically by the SDK from env
process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING ??=
  process.env.AZURE_COSMOSDB_NOSQL_CONNECTION_STRING;

async function run() {
  // Validate required environment variables
  const required = [
    "AZURE_OPENAI_API_KEY",
    "AZURE_OPENAI_ENDPOINT",
    "AZURE_COSMOSDB_NOSQL_CONNECTION_STRING",
  ];
  const missing = required.filter((key) => !process.env[key]);
  if (missing.length > 0) {
    console.error(`Missing required environment variables: ${missing.join(", ")}`);
    console.error("Copy .env.example to .env and fill in your values.");
    process.exit(1);
  }

  const databaseName = process.env.AZURE_COSMOSDB_DATABASE_NAME || "langchain";
  const containerName = process.env.AZURE_COSMOSDB_CONTAINER_NAME || "documents";

  // Load and split documents
  const text = readFileSync("./sample-data.txt", "utf-8");
  const rawDocuments = [new Document({ pageContent: text })];

  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 1000,
    chunkOverlap: 100,
  });
  const documents = await splitter.splitDocuments(rawDocuments);
  console.log(`Split into ${documents.length} document chunks.`);

  // Create the Azure Cosmos DB vector store
  const store = await AzureCosmosDBNoSQLVectorStore.fromDocuments(
    documents,
    new AzureOpenAIEmbeddings({
      ...azureOpenAIConfig,
      azureOpenAIApiDeploymentName:
        process.env.AZURE_OPENAI_EMBEDDING_DEPLOYMENT || "text-embedding-3-small",
    }),
    { databaseName, containerName }
  );
  console.log("Documents stored in Azure Cosmos DB.");

  // 1. Semantic similarity search
  const searchResults = await store.similaritySearch(
    "What are the key features?",
    4
  );
  console.log("\n--- Semantic Search Results ---");
  for (const doc of searchResults) {
    console.log(`- ${doc.pageContent.substring(0, 100)}...`);
  }

  // 2. Vector search with score threshold
  const scoredResults = await store.similaritySearchWithScore(
    "What are the key features?",
    10,
    {
      searchType: AzureCosmosDBNoSQLSearchType.VectorScoreThreshold,
      threshold: 0.8,
    }
  );
  console.log("\n--- Results with Scores ---");
  for (const [doc, score] of scoredResults) {
    console.log(`  Score: ${score.toFixed(4)}, Content: ${doc.pageContent.substring(0, 80)}...`);
  }

  // 3. MMR search for diverse results
  const mmrResults = await store.maxMarginalRelevanceSearch(
    "What are the key features?",
    { k: 5, fetchK: 20, lambda: 0.5 }
  );
  console.log("\n--- MMR Search Results ---");
  for (const doc of mmrResults) {
    console.log(`- ${doc.pageContent.substring(0, 100)}...`);
  }

  // 4. RAG: Answer questions using retrieved documents
  const model = new AzureChatOpenAI({
    ...azureOpenAIConfig,
    azureOpenAIApiDeploymentName:
      process.env.AZURE_OPENAI_CHAT_DEPLOYMENT || "gpt-4.1",
    temperature: 0,
  });

  const retriever = store.asRetriever();
  const formatDocs = (docs) => docs.map((d) => d.pageContent).join("\n\n");

  const prompt = ChatPromptTemplate.fromMessages([
    [
      "system",
      "Answer the user's questions based on the below context:\n\n{context}",
    ],
    ["human", "{question}"],
  ]);

  const chain = RunnableSequence.from([
    {
      context: async (input) => {
        const docs = await retriever.invoke(input.question);
        return formatDocs(docs);
      },
      question: (input) => input.question,
    },
    prompt,
    model,
    new StringOutputParser(),
  ]);

  const question = "What are the main benefits described in the document?";
  const answer = await chain.invoke({ question });

  console.log(`\n--- RAG Response ---`);
  console.log(`Question: ${question}`);
  console.log(`Answer: ${answer}`);

  const sourceDocs = await retriever.invoke(question);
  console.log("\nSource documents:");
  for (const doc of sourceDocs) {
    console.log(`- ${doc.pageContent.substring(0, 100)}...`);
  }
}

run().catch((error) => {
  console.error("Error:", error.message);
  process.exit(1);
});