Created
May 21, 2024 15:37
-
-
Save askmeegs/54bf804759b640fc3dc18b8d1f17e4dd to your computer and use it in GitHub Desktop.
genkit-flow-rag.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 }); | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment