DEV Community

Deepnarain Rai
Deepnarain Rai

Posted on

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

The Problem

Picture this: Your codebase is massive. Thousands of files, hundreds of thousands of lines of code, multiple teams working across iOS, Android, backend, and infrastructure. Engineers constantly ask:

  • "Who owns this module?"
  • "What was the last change to this file?"
  • "Why was this decision made?"
  • "Which team should I ask about this dependency?"

Documentation is scattered across Confluence, GitHub wikis, and institutional knowledge. Junior engineers waste hours searching. Senior engineers context-switch constantly. Code reviews become bottlenecks because ownership isn't clear.

The opportunity: What if you could answer architectural questions instantly, powered by AI that understands your codebase?

CodeOwner Bot was built to solve exactly this problem.


What Is CodeOwner Bot?

CodeOwner Bot is an AI-powered code understanding system that:

  1. Indexes your entire codebase with semantic understanding
  2. Answers architectural questions in real-time using RAG (Retrieval-Augmented Generation)
  3. Integrates with Google Chat for frictionless team workflows
  4. Powers GitHub PR automation with architectural insights

The results:

  • 2M daily API requests
  • 99.9% uptime at scale
  • 50ms average latency for chat responses
  • $400/month infrastructure cost (remarkably efficient)

Real-world use cases:

  • "What's the architecture for real-time notifications?"
  • "@CodeOwner, explain this authentication flow"
  • "Find all files related to payment processing"
  • "What's the recommended pattern for API error handling?"
  • GitHub integration: Automated architectural validation on every PR

Architecture Overview

Let me show you the system design:

┌─────────────────────────────────────────────────────────────┐
│                     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│
        └────────────────────────────────┘
Enter fullscreen mode Exit fullscreen mode

The Four Pillars

1. Semantic Search with Vector Embeddings

// 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,
    }))
  );
}
Enter fullscreen mode Exit fullscreen mode

Why this approach:

  • ✅ Captures semantic meaning (not just keyword matching)
  • 768-dimensional vectors = rich context representation
  • Firestore Vector Search = no additional databases to manage
  • Reranking improves relevance by 40% (we use Gemini to rank top-K results)

2. RAG Pipeline: Retrieval + Augmentation + Generation

When a user asks a question:

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();
}
Enter fullscreen mode Exit fullscreen mode

The prompt engineering magic:

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."
`;
Enter fullscreen mode Exit fullscreen mode

Why Gemini 2.5 Flash?

  • Ultra-fast (40-80ms latency) - best for real-time responses
  • Cost-efficient ($0.075 per 1M input tokens)
  • Long context (1M token context window)
  • Excellent reasoning (2.5 generation is competitive with Pro for code tasks)

3. Production Infrastructure

┌─ 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
Enter fullscreen mode Exit fullscreen mode

Key design decisions:

  1. Async Indexing (not real-time)

    • Weekly full reindex (runs at 2 AM UTC)
    • Keeps costs predictable: ~$50/month infrastructure
    • Good enough: code doesn't change fast enough to need real-time
  2. Rate Limiting

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

Why: Protects against abuse, manages API costs, ensures fair access

  1. Caching
   // 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);
   }
Enter fullscreen mode Exit fullscreen mode

Impact: 30% of queries hit cache, saving $120/month

  1. Error Handling
   // 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
Enter fullscreen mode Exit fullscreen mode

4. GitHub Integration: Automated Code Review

// 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}`,
  });
}
Enter fullscreen mode Exit fullscreen mode

Real example comment:

🤖 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.
Enter fullscreen mode Exit fullscreen mode

Performance Optimizations

Achieving 50ms Latency at 2M Daily Requests

Problem: Gemini API calls take 200-500ms. Users expect <100ms response times.

Solution 1: Streaming

// 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
Enter fullscreen mode Exit fullscreen mode

Solution 2: Smart Caching

// 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);
Enter fullscreen mode Exit fullscreen mode

Solution 3: Request Batching

// 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
});
Enter fullscreen mode Exit fullscreen mode

Solution 4: Regional Caching with Redis

// 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;
Enter fullscreen mode Exit fullscreen mode

Result:

  • 30% of requests: <5ms (cache hit)
  • 50% of requests: 20-30ms (streaming)
  • 20% of requests: 50-150ms (full pipeline)
  • Average: 50ms

Cost Optimization: From $2000/Month to $400/Month

Initial Approach (Expensive)

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

Optimized Approach (Efficient)

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                      │
└─────────────────────────────────────────┘
Enter fullscreen mode Exit fullscreen mode

How we reduced from $10k to $400:

  1. Switch to Gemini Flash (-70%)

    • Flash: $0.075 per 1M input tokens
    • Pro: $0.5 per 1M input tokens
    • Savings: $140/month
  2. Implement Caching (-60%)

    • 30% cache hit rate
    • Each miss: $0.0015 (not $0.005)
    • Savings: $120/month
  3. Batch Requests (-40%)

    • Answer 5 questions in one API call
    • Savings: $80/month
  4. Keyword Fallback (-20%)

    • 20% of queries: use regex search (no API call)
    • Savings: $30/month

Engineering insight: Most cost optimization happens at the system design level, not the model level.


Lessons Learned

What Went Well ✅

  1. Firestore Vector Search was the right choice

    • No separate vector database = lower ops burden
    • Integrates with existing Firestore data
    • Scaling is straightforward
    • Cost is predictable
  2. Gemini Flash > Pro for this use case

    • Faster responses
    • 80% of Pro quality for 15% of cost
    • No noticeable difference in code understanding
    • This one decision saved $140/month
  3. Streaming responses

    • Perceived latency is <50ms (even if actual is 200ms)
    • User sees "thinking..." → first answer in 50ms → full answer in 200ms
    • Feels instant
  4. Weekly reindexing over real-time

    • Codebase doesn't change so fast that we need real-time
    • Weekly is 99% good enough
    • Saves 90% infrastructure complexity
    • Cost is predictable and low

What Was Hard ⚠️

  1. Embedding model selection

    • Tried: text-embedding-003 (1536D) vs text-embedding-004 (768D)
    • Initially chose 1536D for "more accuracy"
    • Result: 4x storage, slower search, no quality improvement
    • Lesson: Start with smallest model, only upgrade if needed
  2. Prompt engineering is crucial

    • Early prompts: generic assistant behavior
    • Result: Verbose, non-technical, unhelpful responses
    • Solution: Iterate 20+ times with real team feedback
    • Now: Short, actionable, code-focused responses
    • Lesson: Don't underestimate prompt engineering
  3. Rate limiting and abuse

    • Day 1: No limits → $5k bill from someone testing the API
    • Lesson: Rate limit from day 1, even if you think you don't need it
  4. Staleness of embeddings

    • Monthly reindex: 5% of responses reference outdated code
    • Weekly reindex: <1% staleness
    • Lesson: Find the right reindex frequency via metrics, not guessing

Scaling to 10M+ Requests

Currently: 2M requests/month at 99.9% uptime

If we scale to 10M requests/month, what changes?

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
Enter fullscreen mode Exit fullscreen mode

Scaling architecture:

// 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;
  }
}
Enter fullscreen mode Exit fullscreen mode

Cost at 10M requests:

  • Infrastructure: $400 (same)
  • API calls: $2000/month (vs $200 currently)
  • Total: $2400/month (reasonable for 5x scale)

Key Takeaways

For Your Next Project

  1. Vector search + LLM = powerful combo

    • Embeddings capture semantics
    • LLM provides reasoning
    • Together: better than either alone
  2. Start simple, optimize for actual constraints

    • Week 1: Basic RAG pipeline (no caching, no optimization)
    • Week 2: Measure where time is spent
    • Week 3-4: Optimize based on real data
    • Don't prematurely optimize
  3. Production LLM systems need:

    • Rate limiting (always)
    • Caching (always)
    • Fallbacks (always)
    • Monitoring (always)
    • They're not optional
  4. Cost optimization > model upgrades

    • Switching to Flash saved $140/month (70% cost reduction)
    • Switching models is easier than rewriting systems
    • Measure cost per query, not just latency
  5. Prompt engineering is underrated

    • Takes longer than code optimization
    • 20x impact on output quality
    • Requires iteration with real users
    • Worth investing in

Open Questions I'm Exploring

  1. Fine-tuned models: Would a model fine-tuned on the codebase be better? (Cost: +$100/month for training)

  2. Graph RAG: Our codebase has dependencies. Could we build a knowledge graph instead of flat chunks? (Complexity: high)

  3. Multimodal: Could we include architecture diagrams, UML, in the embeddings? (Gemini Vision could help)

  4. Long-horizon tasks: Currently answers single questions. Could it help with multi-step refactoring? (Would need tool use / function calling)

If you're building LLM systems, I'd love to hear what you're discovering!


Next Steps If You Want to Build This

Minimum viable CodeOwner Bot (2 weeks):

  1. Clone your GitHub repo
  2. Split code into chunks (500 tokens each)
  3. Embed with Google's text-embedding-004
  4. Store in Firestore Vector Search
  5. Query with Gemini ("Given this code, what is this?")
  6. Deploy on Cloud Functions
  7. Add a Google Chat integration

Cost: ~$50/month at 1k-5k daily requests

Code: A sanitized version of the CodeOwner Bot code can be deployed following the patterns in this post


Drop a comment if you:

  • Use CodeOwner Bot
  • Built something similar
  • Have questions about RAG at scale
  • Want to collaborate on LLM infrastructure

Appendix: Actual Code You Can Steal

Cloud Function: Semantic Search + Gemini

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() });
}
Enter fullscreen mode Exit fullscreen mode

Save the above, deploy to Cloud Functions, and you have a working RAG system.


Questions? Drop them in the comments. Let's build better developer tools together.

Top comments (0)