Initial commit

templates/rag-agent.ts

@@ -0,0 +1,196 @@
+// RAG Agent with Vectorize + Workers AI
+
+import { Agent } from "agents";
+import { generateText } from "ai";
+import { openai } from "@ai-sdk/openai";
+
+interface Env {
+  AI: Ai;                      // Workers AI binding
+  VECTORIZE: Vectorize;        // Vectorize binding
+  OPENAI_API_KEY?: string;     // Optional: OpenAI API key
+}
+
+interface RAGState {
+  documentsIngested: number;
+  queriesProcessed: number;
+  lastQuery: string | null;
+}
+
+export class RAGAgent extends Agent<Env, RAGState> {
+  initialState: RAGState = {
+    documentsIngested: 0,
+    queriesProcessed: 0,
+    lastQuery: null
+  };
+
+  async onRequest(request: Request): Promise<Response> {
+    const url = new URL(request.url);
+
+    // POST /ingest - Ingest documents
+    if (url.pathname === "/ingest") {
+      const { documents } = await request.json();
+      // documents: Array<{ id: string, text: string, metadata: any }>
+
+      const count = await this.ingestDocuments(documents);
+
+      return Response.json({
+        ingested: count,
+        total: this.state.documentsIngested + count
+      });
+    }
+
+    // POST /query - Query knowledge base
+    if (url.pathname === "/query") {
+      const { query, topK } = await request.json();
+
+      const results = await this.queryKnowledge(query, topK || 5);
+
+      return Response.json({
+        query,
+        matches: results.matches,
+        context: results.context
+      });
+    }
+
+    // POST /chat - RAG-powered chat
+    if (url.pathname === "/chat") {
+      const { message } = await request.json();
+
+      const response = await this.chat(message);
+
+      return Response.json(response);
+    }
+
+    return new Response("Not Found", { status: 404 });
+  }
+
+  // Ingest documents into Vectorize
+  async ingestDocuments(documents: Array<{ id: string; text: string; metadata?: any }>) {
+    const vectors = [];
+
+    for (const doc of documents) {
+      // Generate embedding with Workers AI
+      const { data } = await this.env.AI.run('@cf/baai/bge-base-en-v1.5', {
+        text: [doc.text]
+      });
+
+      vectors.push({
+        id: doc.id,
+        values: data[0],
+        metadata: {
+          text: doc.text,
+          ...doc.metadata
+        }
+      });
+    }
+
+    // Upsert into Vectorize (batch operation)
+    await this.env.VECTORIZE.upsert(vectors);
+
+    // Update state
+    this.setState({
+      ...this.state,
+      documentsIngested: this.state.documentsIngested + vectors.length
+    });
+
+    return vectors.length;
+  }
+
+  // Query knowledge base
+  async queryKnowledge(userQuery: string, topK: number = 5) {
+    // Generate query embedding
+    const { data } = await this.env.AI.run('@cf/baai/bge-base-en-v1.5', {
+      text: [userQuery]
+    });
+
+    // Search Vectorize
+    const results = await this.env.VECTORIZE.query(data[0], {
+      topK,
+      returnMetadata: 'all'
+    });
+
+    // Extract context from matches
+    const context = results.matches
+      .map(match => match.metadata?.text)
+      .filter(Boolean)
+      .join('\n\n');
+
+    // Update state
+    this.setState({
+      ...this.state,
+      queriesProcessed: this.state.queriesProcessed + 1,
+      lastQuery: userQuery
+    });
+
+    return {
+      matches: results.matches.map(m => ({
+        id: m.id,
+        score: m.score,
+        metadata: m.metadata
+      })),
+      context
+    };
+  }
+
+  // RAG-powered chat
+  async chat(userMessage: string) {
+    // 1. Retrieve relevant context
+    const { context } = await this.queryKnowledge(userMessage);
+
+    // 2. Generate response with context
+    if (this.env.OPENAI_API_KEY) {
+      // Use OpenAI with AI SDK
+      const { text } = await generateText({
+        model: openai('gpt-4o-mini'),
+        messages: [
+          {
+            role: 'system',
+            content: `You are a helpful assistant. Use the following context to answer questions accurately:\n\n${context}`
+          },
+          {
+            role: 'user',
+            content: userMessage
+          }
+        ]
+      });
+
+      return { response: text, context };
+    } else {
+      // Use Workers AI
+      const response = await this.env.AI.run('@cf/meta/llama-3-8b-instruct', {
+        messages: [
+          {
+            role: 'system',
+            content: `You are a helpful assistant. Use the following context to answer questions:\n\n${context}`
+          },
+          {
+            role: 'user',
+            content: userMessage
+          }
+        ]
+      });
+
+      return {
+        response: response.response,
+        context
+      };
+    }
+  }
+
+  // Query with metadata filtering
+  async queryWithFilter(userQuery: string, filters: any) {
+    const { data } = await this.env.AI.run('@cf/baai/bge-base-en-v1.5', {
+      text: [userQuery]
+    });
+
+    const results = await this.env.VECTORIZE.query(data[0], {
+      topK: 5,
+      filter: filters,  // e.g., { category: { $eq: "docs" } }
+      returnMetadata: 'all'
+    });
+
+    return results;
+  }
+}
+
+export default RAGAgent;