Saltar al contenido
Lección 9 de 12

RAG para Codigo

8 min read

Por que el Codigo Necesita Tratamiento Especial

El codigo es fundamentalmente diferente del texto en lenguaje natural. Tiene sintaxis rigida, estructura jerarquica (modulos contienen clases que contienen metodos), dependencias entre archivos (imports, herencia) y significado que proviene tanto del texto como de su estructura. Un pipeline RAG generico que fragmenta codigo por conteo de caracteres producira resultados terribles porque dividira funciones por la mitad, separara definiciones de clase de sus metodos y perdera el contexto de imports que hace el codigo comprensible.

Construir un sistema RAG para codigo requiere repensar cada etapa del pipeline: como fragmentas, como embedes, que metadatos extraes y como construyes el prompt al LLM.

Embeddings Especificos para Codigo

Los modelos de embedding de proposito general funcionan con codigo pero tienen rendimiento inferior comparado con modelos entrenados especificamente en lenguajes de programacion. Los modelos especificos para codigo entienden que def calculate_total(items) y function computeSum(products) son semanticamente similares a pesar de no tener palabras en comun.

Modelos recomendados para codigo:

  • OpenAI text-embedding-3-small/large -- Maneja codigo razonablemente bien ya que los datos de entrenamiento incluyen codigo. Suficientemente bueno para la mayoria de los casos.
  • Voyage Code 2 -- Modelo de embedding especializado en codigo de Voyage AI. Supera modelos generales en benchmarks de recuperacion de codigo.
  • CodeSage -- Modelo de embedding de codigo open-source optimizado para recuperacion.
  • StarEncoder -- De BigCode, entrenado especificamente en codigo. Bueno para despliegues autoalojados.
from langchain_openai import OpenAIEmbeddings

# Para la mayoria de proyectos, los embeddings de OpenAI manejan codigo bien
code_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# Para aplicaciones criticas de codigo, considera Voyage
# from langchain_voyageai import VoyageAIEmbeddings
# code_embeddings = VoyageAIEmbeddings(model="voyage-code-2")

Chunking Basado en AST

El Arbol de Sintaxis Abstracta (AST) es una representacion estructurada del codigo fuente que captura su estructura sintactica. En lugar de dividir codigo por conteo de caracteres, parseas el AST y divides en limites naturales: funciones, clases, metodos.

import ast
from dataclasses import dataclass

@dataclass
class CodeChunk:
    content: str
    chunk_type: str  # "function", "class", "module"
    name: str
    file_path: str
    start_line: int
    end_line: int
    docstring: str | None = None
    imports: list[str] | None = None

def extract_python_chunks(file_path: str) -> list[CodeChunk]:
    """Extraer funciones y clases de un archivo Python usando AST."""
    with open(file_path, "r", encoding="utf-8") as f:
        source = f.read()

    tree = ast.parse(source)
    chunks = []
    lines = source.split("\n")

    # Extraer imports a nivel de modulo
    imports = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            imports.append(ast.get_source_segment(source, node))

    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.FunctionDef):
            func_source = "\n".join(lines[node.lineno - 1:node.end_lineno])
            docstring = ast.get_docstring(node)
            chunks.append(CodeChunk(
                content=func_source,
                chunk_type="function",
                name=node.name,
                file_path=file_path,
                start_line=node.lineno,
                end_line=node.end_lineno,
                docstring=docstring,
                imports=imports,
            ))

        elif isinstance(node, ast.ClassDef):
            class_source = "\n".join(lines[node.lineno - 1:node.end_lineno])
            docstring = ast.get_docstring(node)
            chunks.append(CodeChunk(
                content=class_source,
                chunk_type="class",
                name=node.name,
                file_path=file_path,
                start_line=node.lineno,
                end_line=node.end_lineno,
                docstring=docstring,
                imports=imports,
            ))

    return chunks

# Uso
chunks = extract_python_chunks("src/auth/middleware.py")
for chunk in chunks:
    print(f"{chunk.chunk_type}: {chunk.name} ({chunk.start_line}-{chunk.end_line})")

Manejando Funciones y Clases Grandes

Algunas funciones o clases son demasiado grandes para embeber como un solo fragmento. Para clases, divide en el docstring/firma de la clase mas metodos individuales. Para funciones largas, considera incluir la firma de la funcion y docstring como contexto con cada bloque logico.

def split_large_class(class_node, source_lines, file_path):
    """Dividir una clase grande en fragmentos a nivel de metodo."""
    chunks = []

    # Fragmento a nivel de clase (firma + docstring)
    class_header = f"class {class_node.name}:"
    docstring = ast.get_docstring(class_node)
    if docstring:
        class_header += f'\n    """{docstring}"""'

    for node in class_node.body:
        if isinstance(node, ast.FunctionDef):
            method_source = "\n".join(
                source_lines[node.lineno - 1:node.end_lineno]
            )
            # Anteponer contexto de clase a cada metodo
            contextual_chunk = f"# Clase: {class_node.name}\n{method_source}"
            chunks.append(CodeChunk(
                content=contextual_chunk,
                chunk_type="method",
                name=f"{class_node.name}.{node.name}",
                file_path=file_path,
                start_line=node.lineno,
                end_line=node.end_lineno,
            ))

    return chunks

Soporte Multi-Lenguaje

Para JavaScript/TypeScript, usa tree-sitter en lugar del modulo ast de Python. Tree-sitter soporta docenas de lenguajes de programacion con una API consistente.

# Usando el splitter consciente del lenguaje de LangChain
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter

# Division consciente de Python
python_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON,
    chunk_size=2000,
    chunk_overlap=200,
)

# Division consciente de JavaScript
js_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.JS,
    chunk_size=2000,
    chunk_overlap=200,
)

# TypeScript, Go, Java, Rust, etc. son todos soportados
ts_splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.TS,
    chunk_size=2000,
    chunk_overlap=200,
)

Enriqueciendo Fragmentos de Codigo con Metadatos

Los fragmentos de codigo se benefician enormemente de metadatos ricos. Cuanto mas contexto proporciones, mejor sera la recuperacion y mas util la respuesta del LLM.

def enrich_code_metadata(chunk: CodeChunk) -> dict:
    """Crear metadatos ricos para un fragmento de codigo."""
    metadata = {
        "source": chunk.file_path,
        "chunk_type": chunk.chunk_type,
        "name": chunk.name,
        "start_line": chunk.start_line,
        "end_line": chunk.end_line,
        "language": "python",
    }

    # Agregar docstring como campo buscable separado
    if chunk.docstring:
        metadata["docstring"] = chunk.docstring

    # Agregar contexto de imports
    if chunk.imports:
        metadata["imports"] = ", ".join(chunk.imports[:10])

    # Inferir ruta del modulo desde la ruta del archivo
    # src/auth/middleware.py -> auth.middleware
    module_path = chunk.file_path.replace("/", ".").replace(".py", "")
    if module_path.startswith("src."):
        module_path = module_path[4:]
    metadata["module"] = module_path

    return metadata

Construyendo un Chatbot de Repositorio

Aqui hay un ejemplo completo que ingesta un proyecto Python y crea una interfaz de Q&A:

import os
from pathlib import Path
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# --- Paso 1: Cargar archivos de codigo ---
loader = DirectoryLoader(
    "./src",
    glob="**/*.py",
    loader_cls=TextLoader,
    loader_kwargs={"encoding": "utf-8"},
)
code_docs = loader.load()

# Agregar metadatos de lenguaje
for doc in code_docs:
    doc.metadata["language"] = "python"
    doc.metadata["content_type"] = "code"

# --- Paso 2: Chunking consciente del lenguaje ---
splitter = RecursiveCharacterTextSplitter.from_language(
    language=Language.PYTHON,
    chunk_size=2000,
    chunk_overlap=200,
)
chunks = splitter.split_documents(code_docs)
print(f"Creados {len(chunks)} fragmentos de codigo")

# --- Paso 3: Embeber y almacenar ---
vectorstore = Chroma.from_documents(
    chunks,
    OpenAIEmbeddings(model="text-embedding-3-small"),
    persist_directory="./code_chroma_db",
)

# --- Paso 4: Prompt especifico para codigo ---
CODE_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """Eres un asistente desarrollador senior que responde preguntas
sobre un repositorio de codigo. Tienes acceso a los siguientes fragmentos de codigo.

Reglas:
1. Responde basandote en el codigo proporcionado en el contexto.
2. Al referenciar codigo, cita la ruta del archivo y numeros de linea.
3. Si el codigo no contiene suficiente informacion, dilo.
4. Explica no solo que hace el codigo, sino por que probablemente
   lo hace de esa manera.
5. Sugiere mejoras solo cuando te lo pidan.

Contexto de Codigo:
{context}"""),
    ("human", "{question}"),
])

# --- Paso 5: Construir cadena ---
retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={"k": 8, "fetch_k": 25},
)

def format_code_docs(docs):
    formatted = []
    for doc in docs:
        source = doc.metadata.get("source", "desconocido")
        formatted.append(f"# Archivo: {source}\n{doc.page_content}")
    return "\n\n---\n\n".join(formatted)

chain = (
    {"context": retriever | format_code_docs, "question": RunnablePassthrough()}
    | CODE_PROMPT
    | ChatOpenAI(model="gpt-4o", temperature=0)
    | StrOutputParser()
)

# --- Usarlo ---
answer = chain.invoke("Como valida tokens el middleware de autenticacion?")
print(answer)

Integracion con Repositorios Git

Para sistemas RAG que necesitan mantenerse actualizados con un repositorio, integra con git para detectar cambios y re-indexar solo archivos modificados:

import subprocess

def get_changed_files(since_commit: str = "HEAD~1") -> list[str]:
    """Obtener archivos cambiados desde un commit especifico."""
    result = subprocess.run(
        ["git", "diff", "--name-only", since_commit],
        capture_output=True, text=True
    )
    return [f for f in result.stdout.strip().split("\n") if f.endswith(".py")]

def incremental_index(vectorstore, changed_files: list[str]):
    """Re-indexar solo archivos cambiados."""
    for file_path in changed_files:
        # Eliminar fragmentos viejos de este archivo
        vectorstore.delete(where={"source": file_path})

        # Re-cargar y re-fragmentar el archivo
        loader = TextLoader(file_path, encoding="utf-8")
        docs = loader.load()
        chunks = splitter.split_documents(docs)

        # Agregar nuevos fragmentos
        vectorstore.add_documents(chunks)
        print(f"Re-indexado: {file_path} ({len(chunks)} fragmentos)")

Generacion de Documentacion Desde Codigo

RAG para codigo tambien puede funcionar a la inversa: dado codigo, generar documentacion.

DOC_GEN_PROMPT = ChatPromptTemplate.from_messages([
    ("system", """Genera documentacion clara y concisa para el codigo dado.
Incluye:
- Un resumen de una linea
- Parametros y tipos de retorno
- Ejemplos de uso
- Cualquier nota importante sobre comportamiento o casos limite

Codigo:
{code}"""),
    ("human", "Genera documentacion para este codigo."),
])

Consejos para RAG de Codigo

  • Incluye la ruta del archivo en cada fragmento. Codigo sin contexto de ruta de archivo es dificil de navegar.
  • Preserva las sentencias de import. Incluye imports en cada fragmento o almacenalos como metadatos. Le dicen al LLM que bibliotecas y modulos usa el codigo.
  • Usa fragmentos mas grandes para codigo. El codigo es mas denso que la prosa. Un fragmento de codigo de 2000 caracteres es a menudo una sola funcion, mientras que un fragmento de prosa de 2000 caracteres podria ser varios parrafos. Inclinate hacia fragmentos mas grandes.
  • Indexa documentacion junto con codigo. Archivos README, docstrings, comentarios y documentacion en linea deberian indexarse como fragmentos separados con referencias al codigo que describen.
  • Prueba con preguntas reales de desarrolladores. "Como funciona X?", "Donde esta implementado Y?", "Que llama a la funcion Z?" -- estas son las consultas que tu sistema necesita manejar bien.

En la siguiente leccion, aprenderas como medir si tu sistema RAG -- para codigo o cualquier otro contenido -- esta realmente produciendo buenos resultados.