¿Quieres que la IA realmente entienda tu código? Esta herramienta afirma que puede ayudarte.

En este blog, te mostraremos cómo indexar código base para RAG con CocoIndex. CocoIndex es una herramienta que te ayuda a indexar y consultar tus datos. Está diseñada para usarse como marco de trabajo para crear tu propia canalización de datos. CocoIndex ofrece compatibilidad integrada con la fragmentación de código base, con compatibilidad nativa con Tree-sitter.

Cuidador de árboles

Tree-sitter es una herramienta generadora de analizadores sintácticos y una biblioteca de análisis incremental. Está disponible en Rust 🦀 - GitHub . CocoIndex se integra con Rust y Tree-sitter para analizar código eficientemente y extraer árboles de sintaxis para varios lenguajes de programación.

La fragmentación de código base consiste en descomponer un código base en fragmentos más pequeños y semánticamente significativos. CocoIndex aprovecha las capacidades de Tree-sitter para fragmentar el código de forma inteligente según la estructura sintáctica real, en lugar de saltos de línea arbitrarios. Estos fragmentos semánticamente coherentes se utilizan para crear un índice más eficaz para los sistemas RAG, lo que permite una recuperación de código más precisa y una mejor preservación del contexto.

Pase rápido 🚀: puedes encontrar el código completo aquí . Solo unas 50 líneas de código Python para la canalización RAG. ¡Échale un vistazo! 🤗

Si te gusta nuestro trabajo, dale una estrella a CocoIndex en Github para apoyarnos. Muchas gracias con un cálido abrazo de coco 🥥🤗.

Prerrequisitos

Si no tiene Postgres instalado, consulte la guía de instalación . CocoIndex utiliza Postgres para gestionar el índice de datos. Lo tenemos previsto para ser compatible con otras bases de datos, incluidas las que están en desarrollo. Si le interesan otras bases de datos, infórmenos creando un problema en GitHub o Discord .

Definir el flujo de cocoIndex

Definamos el flujo cocoIndex para leer desde una base de código e indexarlo para RAG.

El diagrama de flujo anterior ilustra cómo procesaremos nuestra base de código:

1. Leer archivos de código del sistema de archivos local
2. Extraer extensiones de archivo
3. Dividir el código en fragmentos semánticos usando Tree-sitter
4. Generar incrustaciones para cada fragmento
5. Almacenar en una base de datos vectorial para su recuperación

Implementemos este flujo paso a paso.

1. Agregue el código base como fuente .

 @cocoindex.flow_def(name="CodeEmbedding") def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope): """ Define an example flow that embeds files into a vector database. """ data_scope["files"] = flow_builder.add_source( cocoindex.sources.LocalFile(path="../..", included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"], excluded_patterns=[".*", "target", "**/node_modules"])) code_embeddings = data_scope.add_collector()

En este ejemplo, indexaremos el código base de cocoindex desde el directorio raíz. Puede cambiar la ruta del código base que desea indexar. Indexaremos todos los archivos con las extensiones .py , .rs , .toml , .md y .mdx , y omitiremos los directorios que empiecen por ., target (en la raíz) y node_modules (en cualquier directorio).

flow_builder.add_source creará una tabla con los siguientes subcampos, consulte la documentación aquí.

• filename (clave, tipo: str ): el nombre del archivo, p. ej. dir1/file1.md
• content (tipo: str si binary es False , de lo contrario bytes ): el contenido del archivo

2. Procesa cada archivo y recopila la información .

2.1 Extraer la extensión de unnombre de archivo

Primero, definamos una función para extraer la extensión de un nombre de archivo al procesar cada archivo. Puede encontrar la documentación de la función personalizada aquí .

 @cocoindex.op.function() def extract_extension(filename: str) -> str: """Extract the extension of a filename.""" return os.path.splitext(filename)[1]

Luego vamos a procesar cada archivo y recopilar la información.

 # ... with data_scope["files"].row() as file: file["extension"] = file["filename"].transform(extract_extension)

Aquí extraemos la extensión del nombre del archivo y la almacenamos en el campo extension . Por ejemplo, si el nombre del archivo es spec.rs , el campo extension será .rs .

2.2 Dividir el archivo en fragmentos

A continuación, dividiremos el archivo en fragmentos. Usamos la función SplitRecursively para dividirlo. Puedes encontrar la documentación de la función aquí .

CocoIndex ofrece compatibilidad integrada con Tree-sitter, lo que permite pasar el idioma al parámetro language . Para ver todos los nombres y extensiones de idiomas compatibles, consulte la documentación aquí . Se admiten todos los lenguajes principales, como Python, Rust, JavaScript, TypeScript, Java, C++, etc. Si no se especifica o el idioma especificado no es compatible, se tratará como texto sin formato.

 with data_scope["files"].row() as file: # ... file["chunks"] = file["content"].transform( cocoindex.functions.SplitRecursively(), language=file["extension"], chunk_size=1000, chunk_overlap=300)

2.3 Incrustar los fragmentos

Usaremos la función SentenceTransformerEmbed para incrustar los fragmentos. Puedes encontrar la documentación de la función aquí . Hay 12 000 modelos compatibles con 🤗 Hugging Face . Puedes elegir tu modelo favorito.

 def code_to_embedding(text: cocoindex.DataSlice) -> cocoindex.DataSlice: """ Embed the text using a SentenceTransformer model. """ return text.transform( cocoindex.functions.SentenceTransformerEmbed( model="sentence-transformers/all-MiniLM-L6-v2"))

Luego, para cada fragmento, lo integraremos usando la función code_to_embedding y recopilaremos las integraciones en el recopilador code_embeddings .

Extraemos esta función code_to_embedding en lugar de llamar directamente a transform(cocoindex.functions.SentenceTransformerEmbed(...)) en su lugar.

Esto se debe a que queremos que este se comparta entre la creación del flujo de indexación y la definición del controlador de consultas. Alternativamente, para simplificarlo, también podemos evitar esta función adicional y ejecutar las acciones directamente en el lugar; no es complicado copiar y pegar un poco; lo hicimos para el proyecto de inicio rápido .

 with data_scope["files"].row() as file: # ... with file["chunks"].row() as chunk: chunk["embedding"] = chunk["text"].call(code_to_embedding) code_embeddings.collect(filename=file["filename"], location=chunk["location"], code=chunk["text"], embedding=chunk["embedding"])

2.4 Recopilar las incrustaciones

Por último, exportemos las incrustaciones a una tabla.

 code_embeddings.export( "code_embeddings", cocoindex.storages.Postgres(), primary_key_fields=["filename", "location"], vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])

3. Configure el controlador de consultas para su índice

Usaremos SimpleSemanticsQueryHandler para consultar el índice. Tenga en cuenta que debemos pasar la función code_to_embedding al parámetro query_transform_flow . Esto se debe a que el controlador de consultas usará el mismo modelo de incrustación que el usado en el flujo.

 query_handler = cocoindex.query.SimpleSemanticsQueryHandler( name="SemanticsSearch", flow=code_embedding_flow, target_name="code_embeddings", query_transform_flow=code_to_embedding, default_similarity_metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)

Defina una función principal para ejecutar el controlador de consultas.

 @cocoindex.main_fn() def _run(): # Run queries in a loop to demonstrate the query capabilities. while True: try: query = input("Enter search query (or Enter to quit): ") if query == '': break results, _ = query_handler.search(query, 10) print("\nSearch results:") for result in results: print(f"[{result.score:.3f}] {result.data['filename']}") print(f" {result.data['code']}") print("---") print() except KeyboardInterrupt: break if __name__ == "__main__": load_dotenv(override=True) _run()

El decorador @cocoindex.main_fn() inicializa la biblioteca con la configuración cargada desde las variables de entorno. Consulte la documentación sobre la inicialización para obtener más detalles.

Ejecutar la configuración y actualización del índice

🎉¡Ya está todo listo!

Ejecute los siguientes comandos para configurar y actualizar el índice.

 python main.py cocoindex setup python main.py cocoindex update

Verá el estado de actualizaciones del índice en la terminal.

Probar la consulta

En este punto, puede iniciar el servidor cocoindex y desarrollar su tiempo de ejecución RAG contra los datos.

Para probar su índice, hay dos opciones:

Opción 1: Ejecutar el servidor de índice en la terminal

 python main.py

Cuando vea el mensaje, puede ingresar su consulta de búsqueda, por ejemplo: spec.

 Enter search query (or Enter to quit): spec

Puedes encontrar los resultados de la búsqueda en la terminal.

Resultados devueltos: cada entrada contiene la puntuación (similitud del coseno), el nombre del archivo y el fragmento de código que coincide. En cocoindex, utilizamos cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY para medir la similitud entre la consulta y los datos indexados. También puede cambiar a otras métricas y probarlas rápidamente.

Para obtener más información sobre similitud consina, consulte Wiki .

Opción 2: Ejecute CocoInsight para comprender su flujo de datos y su índice de datos

CocoInsight es una herramienta que le ayuda a comprender su flujo de datos y su índice. Se conecta a su servidor local de CocoInsight sin retención de datos.

CocoInsight ya está en acceso anticipado (gratis). ¡Nos encontraste! Un videotutorial rápido de 3 minutos sobre CocoInsight: Míralo en YouTube .

Ejecute el servidor CocoIndex

 python main.py cocoindex server -c https://cocoindex.io

Una vez que el servidor esté en funcionamiento, abra CocoInsight en su navegador. Podrá conectarse a su servidor local de CocoIndex y explorar su flujo de datos e índice.

En el lado derecho, puedes ver el flujo de datos que definimos.

En el lado izquierdo, puede ver el índice de datos en la vista previa de datos.

Puede hacer clic en cualquier fila para ver los detalles de esa entrada de datos, incluido el contenido completo de los fragmentos de código y sus incrustaciones.

Comunidad

¡Nos encanta saber de la comunidad! Puedes encontrarnos en Github y Discord .

Si te gusta esta publicación y nuestro trabajo, apoya a CocoIndex en Github con una estrella ⭐. ¡Gracias con un cálido abrazo de coco! 🥥🤗.