mẫu pgvector retriever

Bạn có thể sử dụng PostgreSQL và pgvector làm cách triển khai trình truy xuất. Sử dụng ví dụ sau làm điểm xuất phát và sửa đổi để hoạt động với cơ sở dữ liệu của bạn giản đồ.

import { genkit, z, Document } from 'genkit';
import { googleAI, textEmbedding004 } from '@genkit-ai/googleai';
import { toSql } from 'pgvector';
import postgres from 'postgres';

const ai = genkit({
  plugins: [googleAI()],
});

const sql = postgres({ ssl: false, database: 'recaps' });

const QueryOptions = z.object({
  show: z.string(),
  k: z.number().optional(),
});

const sqlRetriever = ai.defineRetriever(
  {
    name: 'pgvector-myTable',
    configSchema: QueryOptions,
  },
  async (input, options) => {
    const embedding = (await ai.embed({
      embedder: textEmbedding004,
      content: input,
    }))[0].embedding;
    const results = await sql`
      SELECT episode_id, season_number, chunk as content
        FROM embeddings
        WHERE show_id = ${options.show}
        ORDER BY embedding <#> ${toSql(embedding)} LIMIT ${options.k ?? 3}
      `;
    return {
      documents: results.map((row) => {
        const { content, ...metadata } = row;
        return Document.fromText(content, metadata);
      }),
    };
  }
);

Và dưới đây là cách sử dụng chú chó trong một dòng dữ liệu:

// Simple flow to use the sqlRetriever
export const askQuestionsOnGoT = ai.defineFlow(
  {
    name: 'askQuestionsOnGoT',
    inputSchema: z.string(),
    outputSchema: z.string(),
  },
  async (inputQuestion) => {
    const docs = await ai.retrieve({
      retriever: sqlRetriever,
      query: inputQuestion,
      options: {
        show: 'Game of Thrones',
      },
    });
    console.log(docs);

    // Continue with using retrieved docs
    // in RAG prompts.
    //...
  }
);