CodeOwner Bot: Building a Production RAG System with Gemini at Scale

┌─────────────────────────────────────────────────────────────┐
│                     Google Chat / GitHub                     │
│                    (User Interface)                          │
└────────────────────────┬────────────────────────────────────┘
                         │
                         ▼
┌─────────────────────────────────────────────────────────────┐
│              Cloud Functions (TypeScript)                    │
│           - Request routing & validation                     │
│           - Authentication & rate limiting                  │
│           - Response formatting                             │
└────────────────────────┬────────────────────────────────────┘
                         │
        ┌────────────────┼────────────────┐
        │                │                │
        ▼                ▼                ▼
┌──────────────┐  ┌────────────────┐  ┌──────────────┐
│ Firestore    │  │ Vector Search  │  │ GitHub API   │
│ (Metadata)   │  │ (Embeddings)   │  │ (Source)     │
│              │  │                │  │              │
│ - File info  │  │ - Similarity   │  │ - Fetch code │
│ - Ownership  │  │   search       │  │ - PR context │
│ - Dates      │  │ - Reranking    │  │              │
└──────────────┘  └────────────────┘  └──────────────┘
        │                │                │
        └────────────────┼────────────────┘
                         │
                         ▼
        ┌────────────────────────────────┐
        │   Gemini 2.5 Flash             │
        │   (LLM Inference)              │
        │                                │
        │  - Context: Code chunks       │
        │  - Instructions: Architecture │
        │  - Examples: Few-shot prompts │
        │  - Reasoning: Chain-of-thought│
        └────────────────────────────────┘

// How we build the knowledge base:
const documents = await fetchCodebaseStructure();

for (const file of documents) {
  // Break into chunks (512 tokens each)
  const chunks = splitIntoChunks(file.content, 512);

  // Embed each chunk using Google's text-embedding-004
  const embeddings = await embeddingModel.embed({
    texts: chunks.map(c => `${file.path}:\n${c}`),
    outputDimensionality: 768, // 768-dimensional vectors
  });

  // Store in Firestore Vector Search
  await firestoreVectorSearch.addDocuments(
    chunks.map((chunk, idx) => ({
      id: `${file.id}_${idx}`,
      embedding: embeddings[idx],
      metadata: {
        filePath: file.path,
        owner: file.owner,
        lastModified: file.lastModified,
        chunkIndex: idx,
      },
      content: chunk,
    }))
  );
}

async function answerArchitecturalQuestion(userQuery: string) {
  // Step 1: RETRIEVAL - Find relevant code context
  const queryEmbedding = await embeddingModel.embed({
    texts: [userQuery],
    outputDimensionality: 768,
  });

  // Vector similarity search (Firestore)
  const topK = 20; // Get top 20 candidates
  const semanticResults = await firestoreVectorSearch.search(
    queryEmbedding,
    { limit: topK }
  );

  // Step 2: RERANKING - Improve result quality
  // Use Gemini to rank results by relevance
  const rerankedResults = await rerankResults(
    semanticResults,
    userQuery
  );

  // Step 3: AUGMENTATION - Build context
  const context = rerankedResults
    .slice(0, 5) // Top 5 after reranking
    .map(r => `File: ${r.metadata.filePath}\nOwner: ${r.metadata.owner}\n\n${r.content}`)
    .join("\n---\n");

  // Step 4: GENERATION - Ask Gemini with context
  const response = await genAI.generateContent({
    model: "gemini-2-5-flash",
    systemInstruction: ARCHITECTURAL_SYSTEM_PROMPT,
    contents: [
      {
        role: "user",
        parts: [
          {
            text: `Context from our codebase:\n\n${context}\n\n---\n\nQuestion: ${userQuery}`,
          },
        ],
      },
    ],
    generationConfig: {
      temperature: 0.3, // Lower = more focused, less creative
      topK: 40,
      topP: 0.95,
      maxOutputTokens: 500,
    },
  });

  return response.text();
}

const ARCHITECTURAL_SYSTEM_PROMPT = `
You are CodeOwner, an expert architect. You have deep knowledge 
of the codebase, design patterns, and architectural decisions.

Your role:
1. Answer questions about code structure and architecture
2. Explain design decisions and rationale
3. Suggest improvements or alternative patterns
4. Help junior engineers understand complex systems

Style:
- Be concise (under 200 words unless asked for detail)
- Use code examples from context
- Explain the "why" not just the "what"
- Be honest about limitations (if you don't have enough context)

When referencing code:
- Always cite the file path: "In auth/biometric.ts, ..."
- Explain the pattern being used
- Connect to broader architecture

Example response format:
"The real-time notification system is built on Firebase Cloud Messaging. 
In notifications/fcm-handler.ts, we:
1. Register device tokens (FCM Service)
2. Batch notifications for efficiency
3. Use exponential backoff for retries

This approach gives us 99.8% delivery rate while keeping costs low."
`;

┌─ Cloud Scheduler (UTC) ─────────┐
│  "0 2 * * 0"  (Weekly sync)     │  Every Sunday 2 AM
└────────────────┬────────────────┘
                 │ Triggers
                 ▼
         Cloud Function
         (index-codebase)
                 │
         ┌───────┴────────┐
         │                │
         ▼                ▼
    Fetch from         Update
    GitHub             Firestore
    (Private repo)     Vector DB
         │                │
         └────────────────┘
              │
         ┌────┴─────────────┐
         │                  │
    Embed code        Index by
    chunks           ownership

   const rateLimiter = {
     // Per user per day
     maxRequests: 100,
     // Per minute globally
     globalThrottle: 1000,
     // Burst capacity
     burstCapacity: 50,
   };

   // Cache popular questions
   const cache = new Map<string, CacheEntry>();

   // Cache hit: return immediately (0ms latency)
   // Cache miss: run full pipeline (50-100ms)
   // TTL: 24 hours for common questions

   if (cache.has(queryHash)) {
     return cache.get(queryHash);
   }

   // Graceful degradation
   try {
     return await vectorSearch(...);
   } catch (error) {
     // Fall back to keyword search
     return await keywordFallback(...);
   }

   // Always return something useful
   // Even if Gemini is slow, return search results

// On every PR, we:
// 1. Fetch changed files
// 2. Analyze with CodeOwner
// 3. Post architectural insights as PR comment

async function analyzeGitHubPR(prNumber: number) {
  const changes = await github.getPRChanges(prNumber);

  const analysis = await codeOwnerBot.analyze({
    files: changes.map(c => ({
      path: c.filename,
      before: c.patch.before,
      after: c.patch.after,
    })),
    question: `This PR modifies ${changes.length} files. 
      Are there any architectural concerns? 
      Should we review with the code owner?`,
  });

  // Post as comment
  await github.createPRComment(prNumber, {
    body: `## 🤖 CodeOwner Analysis\n\n${analysis}`,
  });
}

🤖 CodeOwner Analysis

This PR modifies auth/biometric.ts and payment/razorpay.ts

⚠️ Architectural note: You're changing the biometric authentication flow. 
This impacts:
- App startup time (currently 1.2s)
- Security scanning in /security/biometric-validation.ts

Consider:
1. Running performance benchmarks
2. Reviewing with the security team (they own biometric-validation.ts)

Suggest mentioning this in the PR description.

// Don't wait for full response, stream chunks to client
const stream = await gemini.generateContentStream({...});

for await (const chunk of stream) {
  // Send each chunk to client immediately
  socket.emit('data', chunk.text());
}

// User sees first words in <50ms
// Full response arrives gradually

// Cache common questions (30% hit rate)
const hotQuestions = [
  "How do I add a new API endpoint?",
  "What's our payment flow?",
  "How does authentication work?",
  "What are our database schemas?",
];

// Pre-compute answers
await precomputeAnswers(hotQuestions);

// Batch multiple questions into single API call
// Reduces latency and cost by 40%

const questions = [
  "What owns notifications?",
  "What owns payments?",
  "What owns auth?",
];

const response = await gemini.generateContent({
  contents: [{
    text: questions.map((q, i) => `${i+1}. ${q}`).join('\n'),
  }],
  // Get all answers at once
});

// Cache in multiple regions
// Users in Singapore hit SG cache (0-5ms)
// Cache miss goes to Firebase (50-100ms)

const cachedResult = await redis.get(`answer:${queryHash}`);
if (cachedResult) {
  return cachedResult; // <5ms
}

const result = await generateWithGemini(query);
await redis.set(`answer:${queryHash}`, result, { EX: 86400 }); // 24h TTL
return result;

Every request:
- Call Gemini Pro: $0.005 per request
- 2M requests/month = $10,000/month ❌

Month cost breakdown:
┌─────────────────────────────────────────┐
│ Infrastructure                          │
│ - Cloud Functions: $100                 │
│ - Firestore: $50                        │
│ - Vector Search: $50                    │
│ Subtotal: $200                          │
├─────────────────────────────────────────┤
│ Gemini API (after optimization)         │
│ - Cache hits (30%): $0 (no API call)    │
│ - Direct answers (50%): $50             │
│ - Full pipeline (20%): $150             │
│ Subtotal: $200                          │
├─────────────────────────────────────────┤
│ Total: ~$400/month                      │
└─────────────────────────────────────────┘

Current: 2M/month = 60k/day = 2.5k/hour = 42/sec
Target:  10M/month = 330k/day = 13.8k/hour = 230/sec

Problem: Gemini API has rate limits (10-30 RPS per key)
Solution: Multiple API keys + request routing

// Load balancing across API keys
class GeminiRateLimiter {
  private keys = [KEY_1, KEY_2, KEY_3, ...KEY_10]; // 10 keys
  private currentKey = 0;

  async call(request) {
    const key = this.keys[this.currentKey];
    this.currentKey = (this.currentKey + 1) % this.keys.length;

    try {
      return await callGeminiWithKey(key, request);
    } catch (error) {
      if (error.rateLimited) {
        return await this.queue.add(request); // Queue for retry
      }
      throw error;
    }
  }
}

// Queue: batch requests during high load
class RequestQueue {
  private queue = [];
  private processing = false;

  async add(request) {
    this.queue.push(request);
    if (!this.processing) this.process();
  }

  async process() {
    this.processing = true;
    while (this.queue.length > 0) {
      const batch = this.queue.splice(0, 10); // 10 at a time
      const results = await Promise.all(
        batch.map(r => callGemini(r))
      );
      batch.forEach((req, i) => req.resolve(results[i]));

      // Wait before next batch (avoid burst limits)
      await sleep(1000);
    }
    this.processing = false;
  }
}

import { GoogleGenerativeAI } from "@google/generative-ai";
import { Firestore } from "@google/cloud/firestore";

const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
const firestore = new Firestore();

export async function answerQuestion(req, res) {
  const { query } = req.body;

  // Step 1: Embed the query
  const embeddingModel = genAI.getGenerativeModel({
    model: "embedding-001",
  });

  const queryEmbedding = await embeddingModel.embedContent({
    content: { parts: [{ text: query }] },
  });

  // Step 2: Vector search
  const results = await firestore
    .collection("code_embeddings")
    .findNearest("embedding", queryEmbedding.embedding, {
      limit: 10,
      distanceMeasure: "COSINE",
    })
    .get();

  // Step 3: Rerank (optional but recommended)
  const context = results.docs
    .slice(0, 5)
    .map((doc) => doc.data().content)
    .join("\n---\n");

  // Step 4: Generate answer
  const model = genAI.getGenerativeModel({
    model: "gemini-2-5-flash",
    systemInstruction: "You are a code expert. Answer questions about this codebase concisely.",
  });

  const response = await model.generateContent([
    {
      text: `Context:\n${context}\n\nQuestion: ${query}`,
    },
  ]);

  res.json({ answer: response.response.text() });
}