RAG para Código
Por qué el Código Necesita Tratamiento Especial
El código es fundamentalmente diferente del texto en lenguaje natural. Tiene sintaxis rígida, estructura jerárquica (módulos contienen clases que contienen métodos), dependencias entre archivos (imports, herencia) y significado que proviene tanto del texto como de su estructura. Un pipeline RAG genérico que fragmenta código por conteo de caracteres producirá resultados terribles porque dividirá funciones por la mitad, separará definiciones de clase de sus métodos y perderá el contexto de imports que hace el código comprensible.
Construir un sistema RAG para código requiere repensar cada etapa del pipeline: cómo fragmentas, cómo embedes, qué metadatos extraes y cómo construyes el prompt al LLM.
Embeddings Específicos para Código
Los modelos de embedding de propósito general funcionan con código pero tienen rendimiento inferior comparado con modelos entrenados específicamente en lenguajes de programación. Los modelos específicos para código entienden que def calculate_total(items) y function computeSum(products) son semánticamente similares a pesar de no tener palabras en común.
Modelos recomendados para código:
- OpenAI text-embedding-3-small/large -- Maneja código razonablemente bien ya que los datos de entrenamiento incluyen código. Suficientemente bueno para la mayoría de los casos.
- Voyage Code 2 -- Modelo de embedding especializado en código de Voyage AI. Supera modelos generales en benchmarks de recuperación de código.
- CodeSage -- Modelo de embedding de código open-source optimizado para recuperación.
- StarEncoder -- De BigCode, entrenado específicamente en código. 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 Árbol de Sintaxis Abstracta (AST) es una representación estructurada del código fuente que captura su estructura sintáctica. En lugar de dividir código por conteo de caracteres, parseas el AST y divides en límites naturales: funciones, clases, métodos.
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 más métodos individuales. Para funciones largas, considera incluir la firma de la función y docstring como contexto con cada bloque lógico.
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 módulo ast de Python. Tree-sitter soporta docenas de lenguajes de programación 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 Código con Metadatos
Los fragmentos de código se benefician enormemente de metadatos ricos. Cuanto más contexto proporciones, mejor será la recuperación y más útil 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
Aquí 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)
Integración 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)")
Generación de Documentación Desde Código
RAG para código también puede funcionar a la inversa: dado código, generar documentación.
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 Código
- Incluye la ruta del archivo en cada fragmento. Código sin contexto de ruta de archivo es difícil de navegar.
- Preserva las sentencias de import. Incluye imports en cada fragmento o almacénalos como metadatos. Le dicen al LLM qué bibliotecas y módulos usa el código.
- Usa fragmentos más grandes para código. El código es más denso que la prosa. Un fragmento de código de 2000 caracteres es a menudo una sola función, mientras que un fragmento de prosa de 2000 caracteres podría ser varios párrafos. Inclínate hacia fragmentos más grandes.
- Indexa documentación junto con código. Archivos README, docstrings, comentarios y documentación en línea deberían indexarse como fragmentos separados con referencias al código que describen.
- Prueba con preguntas reales de desarrolladores. "¿Cómo funciona X?", "¿Dónde está implementado Y?", "¿Qué llama a la función Z?" -- estas son las consultas que tu sistema necesita manejar bien.
En la siguiente lección, aprenderás cómo medir si tu sistema RAG -- para código o cualquier otro contenido -- está realmente produciendo buenos resultados.