genkit-flow-rag.ts

genkit-flow-rag.ts

import { NextRequest, NextResponse } from "next/server";

import { configureGenkit } from "@genkit-ai/core";
import { firebase } from "@genkit-ai/firebase";
import { retrieve } from "@genkit-ai/ai/retriever";
import { generate } from "@genkit-ai/ai";
import { defineFirestoreRetriever } from "@genkit-ai/firebase";
import { initializeApp } from "firebase-admin/app";
import { getFirestore } from "firebase-admin/firestore";
import { vertexAI } from "@genkit-ai/vertexai";
import { textEmbeddingGecko } from "@genkit-ai/vertexai";
import { googleCloud } from "@genkit-ai/google-cloud";
import { AlwaysOnSampler } from "@opentelemetry/sdk-trace-base";
import { geminiPro } from "@genkit-ai/vertexai";
import { defineFlow, runFlow } from "@genkit-ai/flow";
import * as z from "zod";

export const runtime = "nodejs";

configureGenkit({
  plugins: [
    firebase({ projectId: "cpet-sandbox" }),
    vertexAI({ projectId: "cpet-sandbox", location: "us-central1" }),
    googleCloud({
      forceDevExport: true,
      telemetryConfig: {
        sampler: new AlwaysOnSampler(),
        autoInstrumentation: true,
        autoInstrumentationConfig: {
          "@opentelemetry/instrumentation-fs": { enabled: false },
          "@opentelemetry/instrumentation-dns": { enabled: false },
          "@opentelemetry/instrumentation-net": { enabled: false },
        },
        metricExportIntervalMillis: 5_000,
      },
    }),
  ],
  flowStateStore: "firebase",
  traceStore: "firebase",
  enableTracingAndMetrics: true,
  logLevel: "debug",
  telemetry: {
    instrumentation: "googleCloud",
    logger: "googleCloud",
  },
});

const app = initializeApp();

const firestoreRetriever = defineFirestoreRetriever({
  name: "firestoreRetriever",
  firestore: getFirestore("powerpal"),
  collection: "documents",
  contentField: "name",
  vectorField: "embedding_field",
  embedder: textEmbeddingGecko,
  distanceMeasure: "COSINE",
});
console.log("✅ Initialized retriever.");

// GENKIT FLOW  - RAG

const powerpalFlow = defineFlow(
  {
    name: "powerpalFlow",
    inputSchema: z.string(),
    outputSchema: z.string(),
  },
  async (input: string) => {
    // retrieve relevant documents
    const docs = await retrieve({
      retriever: firestoreRetriever,
      query: input,
      options: { k: 3, limit: 3 },
    });

    // generate a response
    const llmResponse = await generate({
      model: geminiPro,
      prompt: `
      You are a home energy chatbot, designed to answer homeowners' questions about energy efficiency and sustainability. You will be given context pertaining to U.S. home energy policy.
      Answer the question based only on the context.
      
      The user's question is: ${input}
`,
      context: docs,
    });

    const output = llmResponse.text();
    return output;
  },
);

export async function POST(req: NextRequest) {
  try {
    const body = await req.json();
    const messages = body.messages ?? [];
    const question = messages[messages.length - 1].content;
    console.log("The question is: ", question);

    const result = await runFlow(powerpalFlow, question);
    console.log("🔮 FLOW RESPONSE: " + result);

    return NextResponse.json({ messages: result }, { status: 200 });
  } catch (e: any) {
    console.log("⚠️ GENKIT ERROR: ", e.message);
    return NextResponse.json({ error: e.message }, { status: 500 });
  }
}