[wip] Modal RAG example (#388)

jayhack · tkucar · commit e36d0b65719d · 2025-02-11T00:26:04.000+01:00
# Motivation

&lt;!-- Why is this change necessary? --&gt;

# Content

&lt;!-- Please include a summary of the change --&gt;

# Testing

&lt;!-- How was the change tested? --&gt;

# Please check the following before marking your PR as ready for review

- [ ] I have added tests for my changes
- [ ] I have updated the documentation or added new documentation as
needed
diff --git a/codegen-examples/examples/modal_repo_analytics/README.md b/codegen-examples/examples/modal_repo_analytics/README.md
@@ -0,0 +1,68 @@
+# Repository Analyzer API
+
+A simple Modal API endpoint that analyzes GitHub repositories using Codegen. The API returns basic metrics about any public GitHub repository including:
+
+- Total number of files
+- Number of functions
+- Number of classes
+
+## Running Locally
+
+1. Install dependencies:
+
+```bash
+uv add modal
+```
+
+2. Start the API server:
+
+```bash
+modal serve src/codegen/extensions/modal/api.py
+```
+
+3. Test with curl:
+
+```bash
+# Replace with your local Modal endpoint URL
+curl "{URL}?repo_name=fastapi/fastapi"
+```
+
+## Response Format
+
+The API returns JSON in this format:
+
+```json
+{
+  "status": "success",
+  "error": "",
+  "num_files": 123,
+  "num_functions": 456,
+  "num_classes": 78
+}
+```
+
+If there's an error, you'll get:
+
+```json
+{
+  "status": "error",
+  "error": "Error message here",
+  "num_files": 0,
+  "num_functions": 0,
+  "num_classes": 0
+}
+```
+
+## Development
+
+The API is built using:
+
+- Modal for serverless deployment
+- FastAPI for the web endpoint
+- Codegen for repository analysis
+
+To deploy changes:
+
+```bash
+modal deploy src/codegen/extensions/modal/api.py
+```
diff --git a/codegen-examples/examples/modal_repo_analytics/api.py b/codegen-examples/examples/modal_repo_analytics/api.py
@@ -0,0 +1,55 @@
+"""Modal API endpoint for repository analysis."""
+
+import modal  # deptry: ignore
+from codegen import Codebase
+from pydantic import BaseModel
+
+# Create image with dependencies
+image = modal.Image.debian_slim(python_version="3.13").apt_install("git").pip_install("fastapi[standard]", "codegen>=0.5.30")
+
+# Create Modal app
+app = modal.App("codegen-repo-analyzer")
+
+
+class RepoMetrics(BaseModel):
+    """Response model for repository metrics."""
+
+    num_files: int = 0
+    num_functions: int = 0
+    num_classes: int = 0
+    status: str = "success"
+    error: str = ""
+
+
+@app.function(image=image)
+@modal.web_endpoint(method="GET")
+def analyze_repo(repo_name: str) -> RepoMetrics:
+    """Analyze a GitHub repository and return metrics.
+
+    Args:
+        repo_name: Repository name in format 'owner/repo'
+
+    Returns:
+        RepoMetrics object containing repository metrics or error information
+    """
+    try:
+        # Validate input
+        if "/" not in repo_name:
+            return RepoMetrics(status="error", error="Repository name must be in format 'owner/repo'")
+
+        # Initialize codebase
+        codebase = Codebase.from_repo(repo_name)
+
+        # Calculate metrics
+        num_files = len(codebase.files(extensions="*"))  # Get all files
+        num_functions = len(codebase.functions)
+        num_classes = len(codebase.classes)
+
+        return RepoMetrics(
+            num_files=num_files,
+            num_functions=num_functions,
+            num_classes=num_classes,
+        )
+
+    except Exception as e:
+        return RepoMetrics(status="error", error=str(e))
diff --git a/codegen-examples/examples/modal_repo_analytics/pyproject.toml b/codegen-examples/examples/modal_repo_analytics/pyproject.toml
@@ -0,0 +1,6 @@
+[project]
+name = "codegen-repo-analyzer"
+version = "0.1.0"
+description = "Modal API endpoint for analyzing GitHub repositories using Codegen"
+requires-python = ">=3.13"
+dependencies = ["modal>=0.73.25", "fastapi[standard]", "codegen>=0.5.30"]
diff --git a/codegen-examples/examples/modal_repo_rag/README.md b/codegen-examples/examples/modal_repo_rag/README.md
@@ -0,0 +1,120 @@
+# Codegen RAG Q&A API
+
+<p align="center">
+  <a href="https://docs.codegen.com">
+    <img src="https://i.imgur.com/6RF9W0z.jpeg" />
+  </a>
+</p>
+
+<h2 align="center">
+  Answer questions about any GitHub repository using RAG
+</h2>
+
+<div align="center">
+
+[![Documentation](https://img.shields.io/badge/Docs-docs.codegen.com-purple?style=flat-square)](https://docs.codegen.com)
+[![License](https://img.shields.io/badge/Code%20License-Apache%202.0-gray?&color=gray)](https://github.com/codegen-sh/codegen-sdk/tree/develop?tab=Apache-2.0-1-ov-file)
+
+</div>
+
+This example demonstrates how to build a RAG-powered code Q&A API using Codegen's VectorIndex and Modal. The API can answer questions about any GitHub repository by:
+
+1. Creating embeddings for all files in the repository
+1. Finding the most relevant files for a given question
+1. Using GPT-4 to generate an answer based on the context
+
+## Quick Start
+
+1. Install dependencies:
+
+```bash
+pip install modal-client codegen openai
+```
+
+2. Create a Modal volume for storing indices:
+
+```bash
+modal volume create codegen-indices
+```
+
+3. Start the API server:
+
+```bash
+modal serve api.py
+```
+
+4. Test with curl:
+
+```bash
+curl -X POST "http://localhost:8000/answer_code_question" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "repo_name": "fastapi/fastapi",
+    "query": "How does FastAPI handle dependency injection?"
+  }'
+```
+
+## API Reference
+
+### POST /answer_code_question
+
+Request body:
+
+```json
+{
+  "repo_name": "owner/repo",
+  "query": "Your question about the code"
+}
+```
+
+Response format:
+
+```json
+{
+  "status": "success",
+  "error": "",
+  "answer": "Detailed answer based on the code...",
+  "context": [
+    {
+      "filepath": "path/to/file.py",
+      "snippet": "Relevant code snippet..."
+    }
+  ]
+}
+```
+
+## How It Works
+
+1. The API uses Codegen to clone and analyze the repository
+1. It creates/loads a VectorIndex of all files using OpenAI's embeddings
+1. For each question:
+   - Finds the most semantically similar files
+   - Extracts relevant code snippets
+   - Uses GPT-4 to generate an answer based on the context
+
+## Development
+
+The API is built using:
+
+- Modal for serverless deployment
+- Codegen for repository analysis
+- OpenAI for embeddings and Q&A
+- FastAPI for the web endpoint
+
+To deploy changes:
+
+```bash
+modal deploy api.py
+```
+
+## Environment Variables
+
+Required environment variables:
+
+- `OPENAI_API_KEY`: Your OpenAI API key
+
+## Learn More
+
+- [Codegen Documentation](https://docs.codegen.com)
+- [Modal Documentation](https://modal.com/docs)
+- [VectorIndex Tutorial](https://docs.codegen.com/building-with-codegen/semantic-code-search)
diff --git a/codegen-examples/examples/modal_repo_rag/api.py b/codegen-examples/examples/modal_repo_rag/api.py
@@ -0,0 +1,126 @@
+"""Modal API endpoint for RAG-based code Q&A using Codegen's VectorIndex."""
+
+import modal
+from codegen import Codebase
+from codegen.extensions import VectorIndex
+from pydantic import BaseModel
+
+# Create image with dependencies
+image = (
+    modal.Image.debian_slim(python_version="3.13")
+    .apt_install("git")
+    .pip_install(
+        "fastapi[standard]",
+        "codegen>=0.5.30",
+        "openai>=1.1.0",
+    )
+)
+
+# Create Modal app
+app = modal.App("codegen-rag-qa")
+
+# Create stub for persistent volume to store vector indices
+stub = modal.Stub("codegen-rag-qa")
+volume = modal.Volume.from_name("codegen-indices")
+
+
+class QARequest(BaseModel):
+    """Request model for code Q&A."""
+
+    repo_name: str
+    query: str
+
+
+class QAResponse(BaseModel):
+    """Response model for code Q&A."""
+
+    answer: str = ""
+    context: list[dict[str, str]] = []  # List of {filepath, snippet} used for answer
+    status: str = "success"
+    error: str = ""
+
+
+@stub.function(
+    image=image,
+    volumes={"/root/.codegen/indices": volume},
+    timeout=600,
+)
+@modal.web_endpoint(method="POST")
+async def answer_code_question(request: QARequest) -> QAResponse:
+    """Answer questions about code using RAG with Codegen's VectorIndex.
+
+    Args:
+        request: QARequest containing repository name and query
+
+    Returns:
+        QAResponse containing answer and context snippets
+    """
+    try:
+        # Validate input
+        if "/" not in request.repo_name:
+            return QAResponse(status="error", error="Repository name must be in format 'owner/repo'")
+
+        # Initialize codebase
+        codebase = Codebase.from_repo(request.repo_name)
+
+        # Initialize vector index
+        index = VectorIndex(codebase)
+
+        # Try to load existing index or create new one
+        try:
+            index.load(f"/root/.codegen/indices/{request.repo_name.replace('/', '_')}.pkl")
+        except FileNotFoundError:
+            # Create new index if none exists
+            index.create()
+            index.save(f"/root/.codegen/indices/{request.repo_name.replace('/', '_')}.pkl")
+
+        # Find relevant files
+        results = index.similarity_search(request.query, k=3)
+
+        # Collect context from relevant files
+        context = []
+        for filepath, score in results:
+            try:
+                file = codebase.get_file(filepath)
+                if file:
+                    context.append(
+                        {
+                            "filepath": filepath,
+                            "snippet": file.content[:1000],  # First 1000 chars as preview
+                            "score": f"{score:.3f}",
+                        }
+                    )
+            except Exception as e:
+                print(f"Error reading file {filepath}: {e}")
+
+        # Format context for prompt
+        context_str = "\n\n".join([f"File: {c['filepath']}\nScore: {c['score']}\n```\n{c['snippet']}\n```" for c in context])
+
+        # Create prompt for OpenAI
+        prompt = f"""Given the following code context and question, provide a clear and accurate answer.
+Focus on the specific code shown in the context.
+
+Question: {request.query}
+
+Relevant code context:
+{context_str}
+
+Answer:"""
+
+        # Get answer from OpenAI
+        from openai import OpenAI
+
+        client = OpenAI()
+        response = client.chat.completions.create(
+            model="gpt-4-turbo-preview",
+            messages=[
+                {"role": "system", "content": "You are a helpful code assistant. Answer questions about code accurately and concisely based on the provided context."},
+                {"role": "user", "content": prompt},
+            ],
+            temperature=0,
+        )
+
+        return QAResponse(answer=response.choices[0].message.content, context=[{"filepath": c["filepath"], "snippet": c["snippet"]} for c in context])
+
+    except Exception as e:
+        return QAResponse(status="error", error=str(e))
diff --git a/codegen-examples/examples/modal_repo_rag/pyproject.toml b/codegen-examples/examples/modal_repo_rag/pyproject.toml
@@ -0,0 +1,11 @@
+[project]
+name = "codegen-rag-qa"
+version = "0.1.0"
+description = "Modal API endpoint for embeddings-based RAG & Q&A on Codegen"
+requires-python = ">=3.13"
+dependencies = [
+  "modal>=0.73.25",
+  "fastapi[standard]",
+  "codegen>=0.5.30",
+  "openai>=1.1.0",
+]
