Building a RAG-powered Slack Bot

from codegen import Codebase
from codegen.extensions import VectorIndex

def initialize_codebase():
    """Initialize and index the codebase."""
    # Initialize codebase with smart caching
    codebase = Codebase.from_repo(
        "codegen-sh/codegen-sdk",
        language="python",
        tmp_dir="/root"
    )

    # Initialize vector index
    index = VectorIndex(codebase)

    # Try to load existing index or create new one
    index_path = "/root/E.pkl"
    try:
        index.load(index_path)
    except FileNotFoundError:
        # Create new index if none exists
        index.create()
        index.save(index_path)

    return codebase, index

def find_relevant_code(index: VectorIndex, query: str) -> list[tuple[str, float]]:
    """Find code snippets relevant to the query."""
    # Get top 10 most relevant files
    results = index.similarity_search(query, k=10)

    # Clean up chunk references from index
    cleaned_results = []
    for filepath, score in results:
        if "#chunk" in filepath:
            filepath = filepath.split("#chunk")[0]
        cleaned_results.append((filepath, score))

    return cleaned_results

from openai import OpenAI

def generate_answer(query: str, context: str) -> str:
    """Generate an answer using RAG."""
    prompt = f"""You are a code expert. Given the following code context and question,
provide a clear and accurate answer.

Note: Keep it short and sweet - 2 paragraphs max.

Question: {query}

Relevant code:
{context}

Answer:"""

    client = OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a code expert. Answer questions about the given repo based on RAG'd results."},
            {"role": "user", "content": prompt},
        ],
        temperature=0,
    )

    return response.choices[0].message.content

def answer_question(query: str) -> tuple[str, list[tuple[str, float]]]:
    """Answer a question about the codebase using RAG."""
    # Initialize or load codebase and index
    codebase, index = initialize_codebase()

    # Find relevant files
    results = find_relevant_code(index, query)

    # Collect context from relevant files
    context = ""
    for filepath, score in results:
        file = codebase.get_file(filepath)
        context += f"File: {filepath}\n```\n{file.content}\n```\n\n"

    # Generate answer
    answer = generate_answer(query, context)

    return answer, results

answer, files = answer_question("How does VectorIndex handle large files?")

print("Answer:", answer)
print("\nRelevant files:")
for filepath, score in files:
    print(f"• {filepath} (score: {score:.2f})")

Answer:
VectorIndex handles large files by automatically chunking them into smaller pieces
using tiktoken. Each chunk is embedded separately and can be searched independently,
allowing for more precise semantic search results.

Relevant files:
• src/codegen/extensions/vector_index.py (score: 0.92)
• src/codegen/extensions/tools/semantic_search.py (score: 0.85)