build your ai coding assistant
v0.0.2
En 2023, la popularidad de la IA generativa llevará a que cada vez más organizaciones introduzcan la codificación asistida por IA. Lo que es ligeramente diferente del GitHub Copilot lanzado en 2021 es que la finalización del código es solo uno de muchos escenarios. Una gran cantidad de empresas están explorando escenarios como generar código completo y revisión de código según los requisitos, y también están introduciendo IA generativa para mejorar la eficiencia del desarrollo.
En este contexto, nosotros (la comunidad de código abierto de Thoughtworks) también hemos abierto una serie de herramientas auxiliares de IA para ayudar a más organizaciones a crear sus propios asistentes de codificación asistidos por IA:
Porque cuando diseñamos AutoDev, varios modelos de código abierto evolucionaban constantemente. En este contexto, sus pasos son:
Por lo tanto, este tutorial también se centra en estos tres pasos. Además, según nuestra experiencia, la pila de tecnología de muestra para este tutorial:
Dado que nuestra experiencia en IA es relativamente limitada, inevitablemente habrá algunos errores. Por lo tanto, también esperamos trabajar con más desarrolladores para construir este proyecto de código abierto.
Combinado con la parte de inteligencia artificial del informe "Ecosistema de desarrolladores" de JetBrains 2023, podemos resumir algunos escenarios generales que reflejan las áreas donde la IA generativa puede desempeñar un papel en el proceso de desarrollo. Éstos son algunos de los principales escenarios:
Cuando creamos AutoDev, también descubrimos escenarios como la creación de SQL DDL, la generación de requisitos, TDD, etc. entonces. Brindamos la capacidad de personalizar escenarios para que los desarrolladores puedan personalizar sus propias capacidades de IA. Para obtener más detalles, consulte: https://ide.unitmesh.cc/customize.
En la codificación diaria, existen varios escenarios diferentes con diferentes requisitos de velocidad de respuesta de la IA (solo como ejemplo):
| escena | Velocidad de respuesta | Generar requisitos de calidad. | Tamaño esperado | ilustrar |
|---|---|---|---|---|
| finalización del código | rápido | medio | 1~6B | La finalización del código es el escenario más común en la codificación diaria y la velocidad de respuesta es crucial. |
| Generación de documentos | medio | medio | 1 | La generación de documentación requiere una comprensión completa de la estructura del código, y la velocidad y la calidad son igualmente importantes. |
| revisión de código | rápido | medio | 1 | Las revisiones de código requieren asesoramiento de alta calidad, pero también deben ser lo más receptivas posible. |
| Generación de pruebas unitarias | rápido | medio | 6B ~ | Las pruebas unitarias generan menos contexto y la capacidad de respuesta y la calidad de la IA son igualmente importantes. |
| refactorización de código | medio | alto | 32B ~ | La refactorización de código puede requerir una mayor comprensión contextual y los tiempos de respuesta pueden ralentizarse moderadamente. |
| generación de demanda | medio | alto | 32B ~ | La generación de demanda es un escenario relativamente complejo y la velocidad de respuesta se puede reducir moderadamente para garantizar la precisión. |
| Búsqueda e interpretación de códigos en lenguaje natural. | Medio-bajo | alto | 32B ~ | La búsqueda e interpretación de códigos en lenguaje natural son escenarios relativamente complejos y la velocidad de respuesta se puede reducir moderadamente para garantizar la precisión. |
PD: El 32B aquí solo se expresa como un orden de magnitud, porque el efecto será mejor con un modelo más grande.
Por lo tanto, lo resumimos como: uno grande, uno mediano, un micro y tres modelos, que proporcionan codificación completa asistida por IA:
La finalización de código con IA puede combinar herramientas IDE para analizar el contexto del código y las reglas del lenguaje de programación, y la IA generará o sugerirá automáticamente fragmentos de código. En herramientas de finalización de código similares a GitHub Copilot, generalmente se dividen en tres modos de subdivisión:
Finalización en línea (en línea)
Similar al modo FIM (rellenar en el medio), el contenido completado está en la línea actual. Por ejemplo: BlotPost blogpost = new , la finalización es: BlogPost(); para lograr: BlogPost blogpost = new BlogPost();
Podemos usar Deepseek Coder como ejemplo para ver el efecto en este escenario:
< ｜fim▁begin｜ > def quick_sort(arr):
if len(arr) < = 1:
return arr
pivot = arr[0]
left = []
right = []
< ｜fim▁hole｜ >
if arr[i] < pivot:
left.append(arr[i])
else:
right.append(arr[i])
return quick_sort(left) + [pivot] + quick_sort(right) < ｜fim▁end｜ >Aquí, necesitamos combinar el código antes y después del cursor.
Finalización en bloque (InBlock)
Logrado a través del aprendizaje contextual (aprendizaje en contexto), el contenido completado está en el bloque de funciones actual. Por ejemplo, el código original es:
fun createBlog ( blogDto : CreateBlogDto ): BlogPost {
}El código completo es:
val blogPost = BlogPost (
title = blogDto.title,
content = blogDto.content,
author = blogDto.author
)
return blogRepository.save(blogPost)Después del bloque
Se logra mediante el aprendizaje contextual (aprendizaje en contexto), finalización después del bloque de funciones actual, como por ejemplo: finalización de una nueva función después del bloque de funciones actual. Por ejemplo, el código original es:
fun createBlog ( blogDto : CreateBlogDto ): BlogPost {
// ...
}El código completo es:
fun updateBlog ( id : Long , blogDto : CreateBlogDto ): BlogPost {
// ...
}
fun deleteBlog ( id : Long ) {
// ...
}Cuando creamos la función de finalización de IA correspondiente, también debemos considerar aplicarla al conjunto de datos de patrón correspondiente para mejorar la calidad de finalización y brindar una mejor experiencia de usuario.
Algunos recursos relacionados para escribir este artículo:
Las explicaciones del código están diseñadas para ayudar a los desarrolladores a administrar y comprender bases de código grandes de manera más efectiva. Estos asistentes pueden responder preguntas sobre la base del código, proporcionar documentación, buscar código, identificar fuentes de errores, reducir la duplicación de código, etc., mejorando así la eficiencia del desarrollo, reduciendo las tasas de error y reduciendo la carga de trabajo de los desarrolladores.
En este escenario, dependiendo de la calidad de generación que esperamos, suele estar compuesto por dos modelos: uno grande y otro micro o uno mediano y uno micro. El modelo más grande tiene mejores resultados en cuanto a calidad de generación. Combinada con nuestra experiencia en diseño en la herramienta Chocolate Factory, normalmente esta función se puede dividir en varios pasos:
Como aplicación RAG, se divide en dos partes: indexación y consulta.
En la etapa de indexación, necesitamos indexar la base del código, lo que implica segmentación de texto, vectorización, indexación de bases de datos y otras tecnologías. Uno de los elementos más desafiantes es la división. Las reglas de división a las que nos referimos son: https://docs.sweep.dev/blogs/chunking-2m-files. Ahora mismo:
En diferentes escenarios, también podemos dividir de diferentes maneras. Por ejemplo, en Chocolate Factory, dividimos a través de AST para garantizar la calidad del contexto generado.
En la etapa de consulta, debemos combinar algunas de nuestras tecnologías de búsqueda tradicionales, como la búsqueda por vectorización, la búsqueda de ruta, etc., para garantizar la calidad de la búsqueda. Al mismo tiempo, en el escenario chino, también debemos considerar la cuestión de la conversión al chino, como la conversión del inglés al chino para garantizar la calidad de la búsqueda.
Para la asistencia diaria, también podemos lograrlo a través de IA generativa, como crear SQL DDL automáticamente, crear casos de prueba automáticamente, crear requisitos automáticamente, etc. Esto sólo se puede lograr personalizando las palabras clave y combinando conocimientos de dominios específicos, por lo que no entraré en detalles aquí.
Además del modelo, el contexto también es un factor importante que afecta las capacidades de asistencia de la IA. Cuando creamos AutoDev, también descubrimos dos modos de contexto diferentes:
Una comparación simple es la siguiente:
| contexto relevante | contexto similar | |
|---|---|---|
| Tecnología de búsqueda | análisis de código estático | Búsqueda de similitud |
| información de la estructura de datos | AST, CFG | Trozo similar |
| Capacidades multiplataforma | Depende del IDE o analizador independiente | No depende de plataformas específicas |
| calidad contextual | extremadamente alto | alto |
| Generar resultados | extremadamente alto | alto |
| costo de construcción | Depende del idioma y la plataforma. | Bajo |
Cuando el soporte para IDE es limitado, el rendimiento relacionado con el contexto traerá un mayor costo .
GitHub Copilot adopta un patrón arquitectónico de contexto similar y su arquitectura detallada se divide en capas de la siguiente manera:
En los materiales de investigación del proyecto "público" Copilot-Explorer, puede ver cómo se construye Prompt. La siguiente es la solicitud enviada a:
{
"prefix" : " # Path: codeviz \ app.py n #.... " ,
"suffix" : " if __name__ == '__main__': rn app.run(debug=True) " ,
"isFimEnabled" : true ,
"promptElementRanges" : [
{
"kind" : " PathMarker " ,
"start" : 0 ,
"end" : 23
},
{
"kind" : " SimilarFile " ,
"start" : 23 ,
"end" : 2219
},
{
"kind" : " BeforeCursor " ,
"start" : 2219 ,
"end" : 3142
}
]
}en:
prefix utilizada para crear el mensaje se crea a partir de PromptElements, que incluye: BeforeCursor , AfterCursor , SimilarFile , ImportedFile , LanguageMarker , PathMarker , RetrievalSnippet y otros tipos. A partir de los nombres de varios PromptElementKind , también podemos ver su verdadero significado.suffix utilizada para construir el mensaje está determinada por la parte donde se encuentra el cursor y cuántas posiciones quedan por calcular en función del límite superior de tokens (2048). El cálculo del token aquí es el cálculo del token LLM real en Copilot, lo calcula Cushman002. La longitud del token de los caracteres chinos es diferente, como: { context: "console.log('你好，世界')", lineCount: 1, tokenLength: 30 } , donde la longitud del contenido en contexto es 20, pero tokenLength es 30, la longitud de los caracteres chinos es 5 (incluido ， ) y el token ocupado por un solo carácter es 3.A continuación se muestra un ejemplo más detallado de un contexto de aplicación Java:
// Path: src/main/cc/unitmesh/demo/infrastructure/repositories/ProductRepository.java
// Compare this snippet from src/main/cc/unitmesh/demo/domain/product/Product.java:
// ....
// Compare this snippet from src/main/cc/unitmesh/demo/application/ProductService.java:
// ...
// @Component
// public class ProductService {
// //...
// }
//
package cc . unitmesh . demo . repositories ;
// ...
@ Component
public class ProductRepository {
//...En el contexto informático, GitHub Copilot utiliza el coeficiente Jaccard (Similitud de Jaccard). Esta parte de la implementación se implementa en Agent. Para obtener una lógica más detallada, consulte: Después de pasar más de medio mes, finalmente realicé ingeniería inversa en Github Copilot.
Recursos relacionados:
Como se mencionó anteriormente, el código relevante se basa en el análisis de código estático , principalmente con la ayuda de la información estructural del código, como AST, CFG, DDG, etc. En diferentes escenarios y plataformas, podemos combinar diferentes herramientas de análisis de código estático. Las siguientes son algunas herramientas de análisis de código estático comunes:
En el escenario de finalización, mediante el análisis de código estático, podemos obtener el contexto actual, como: función actual, clase actual, archivo actual, etc. El siguiente es un ejemplo del contexto de AutoDev para generar pruebas unitarias:
// here are related classes:
// 'filePath: /Users/phodal/IdeaProjects/untitled/src/main/java/cc/unitmesh/untitled/demo/service/BlogService.java
// class BlogService {
// blogRepository
// + public BlogPost createBlog(BlogPost blogDto)
// + public BlogPost getBlogById(Long id)
// + public BlogPost updateBlog(Long id, BlogPost blogDto)
// + public void deleteBlog(Long id)
// }
// 'filePath: /Users/phodal/IdeaProjects/untitled/src/main/java/cc/unitmesh/untitled/demo/dto/CreateBlogRequest.java
// class CreateBlogRequest ...
// 'filePath: /Users/phodal/IdeaProjects/untitled/src/main/java/cc/unitmesh/untitled/demo/entity/BlogPost.java
// class BlogPost {...
@ ApiOperation ( value = "Create a new blog" )
@ PostMapping ( "/" )
public BlogPost createBlog ( @ RequestBody CreateBlogRequest request ) { En este ejemplo, se analiza el contexto de createBlog para obtener las clases de entrada y salida de la función: CreateBlogRequest , información BlogPost e información de clase BlogService, que se proporcionan al modelo como contexto (proporcionado en los comentarios). En este punto, el modelo genera constructores más precisos, así como casos de prueba más precisos.
Dado que el contexto relevante se basa en el análisis de código estático de diferentes idiomas y API de diferentes IDE, también debemos adaptarnos a diferentes idiomas y diferentes IDE. En términos de costo de construcción, es más caro en relación con contextos similares.
Los IDE y los editores son las principales herramientas para los desarrolladores y sus costos de diseño y aprendizaje son relativamente altos. Primero, podemos usar la plantilla oficial para generar:
Luego, agregue funcionalidad encima (¿no es muy simple?), por supuesto que no. Los siguientes son algunos recursos del complemento IDEA que puede consultar:
Por supuesto, es más apropiado referirse al complemento AutoDev.
Puede utilizar directamente la plantilla oficial para generar el complemento correspondiente: https://github.com/JetBrains/intellij-platform-plugin-template
Para la implementación del complemento IDEA, se implementa principalmente a través de Action y Listener, que solo deben registrarse en plugin.xml . Para obtener más información, consulte la documentación oficial: IntelliJ Platform Plugin SDK
Dado que no consideramos el problema de compatibilidad con las versiones IDE de AutoDev en la etapa inicial, para ser compatibles con versiones anteriores de IDE en el futuro, debemos realizar un procesamiento de compatibilidad en el complemento. Por lo tanto, como se describe en el documento oficial: Rangos de números de compilación, podemos ver que diferentes versiones tienen diferentes requisitos para el JDK. Los siguientes son los requisitos para diferentes versiones:
| Número de sucursal | Versión de la plataforma IntelliJ |
|---|---|
| 233 | 2023.3 |
| 232 | 2023.2 |
| 231 | 2023.1 |
| 223 | 2022.3 |
| 222 | 2022.2 NOTA Ahora se requiere Java 17 (publicación de blog) |
| 221 | 2022.1 |
| 213 | 2021.3 |
| 212 | 2021.2 |
| 211 | 2021.1 |
| 203 | 2020.3 NOTA Ahora se requiere Java 11 (publicación de blog) |
Y configúrelo en gradle.properties :
pluginSinceBuild = 223
pluginUntilBuild = 233.*La configuración posterior de compatibilidad es problemática, por lo que puede consultar el diseño de AutoDev.
En términos de finalización automática de código, los fabricantes nacionales se refieren principalmente a la implementación de GitHub Copilot, y la lógica no es complicada.
Activar mediante teclas de método abreviado
Principalmente monitorea la entrada del usuario en Acción y luego:
| Función | tecla de acceso directo | ilustrar |
|---|---|---|
| solicitudFinalizaciones | Alt + / | Obtenga el contexto actual y luego obtenga los resultados de finalización a través del modelo. |
| aplicarIncrustaciones | TAB | Mostrar resultados de finalización en el IDE |
| disponerIncrustaciones | ESC | Cancelar finalización |
| cicloSiguienteIncrustaciones | Alt + ] | Cambiar al siguiente resultado de finalización |
| cicloAnteriorIncrustaciones | Alt + [ | Cambiar al resultado de finalización anterior |
Utilice el método de activación automática
Principalmente monitorea la entrada del usuario a través de EditorFactoryListener y luego activa diferentes resultados de finalización en función de diferentes entradas. El código central es el siguiente:
class AutoDevEditorListener : EditorFactoryListener {
override fun editorCreated ( event : EditorFactoryEvent ) {
// ...
editor.document.addDocumentListener( AutoDevDocumentListener (editor), editorDisposable)
editor.caretModel.addCaretListener( AutoDevCaretListener (editor), editorDisposable)
// ...
}
class AutoDevCaretListener ( val editor : Editor ) : CaretListener {
override fun caretPositionChanged ( event : CaretEvent ) {
// ...
val wasTypeOver = TypeOverHandler .getPendingTypeOverAndReset(editor)
// ...
llmInlayManager.disposeInlays(editor, InlayDisposeContext . CaretChange )
}
}
class AutoDevDocumentListener ( val editor : Editor ) : BulkAwareDocumentListener {
override fun documentChangedNonBulk ( event : DocumentEvent ) {
// ...
val llmInlayManager = LLMInlayManager .getInstance()
llmInlayManager
.editorModified(editor, changeOffset)
}
}
}Luego, según diferentes entradas, se activan diferentes resultados de finalización y se procesa la estructura.
Código de finalización de renderizado
Posteriormente, necesitamos implementar un Inlay Render, que hereda de EditorCustomElementRenderer .
Combinado con las capacidades de interfaz del IDE, necesitamos agregar la Acción correspondiente, el Grupo correspondiente y el Icono correspondiente. El siguiente es un ejemplo de acción:
<add-to-group group-id="ShowIntentionsGroup" relative-to-action="ShowIntentionActions" anchor="after"/>
Los siguientes son algunos grupos de acción de AutoDev:
| ID de grupo | usos de la IA | Descripción |
|---|---|---|
| Mostrar grupo de intenciones | Refactorización de código, interpretación de código, generación de código, prueba de código | Se utiliza para mostrar sugerencias en el contexto del código y se accede a ellas mediante Alt + Enter y ⌥ + Enter en macOS. |
| ConsolaEditorPopupMenu | corregir errores | El menú que se muestra en la consola, como la consola de la estructura de ejecución del programa. |
| Vcs.MessageActionGroup | Generación de información de código. | Menú para escribir mensajes de confirmación en VCS. |
| Vcs.Log.ContextMenu | Revisión de código, interpretación de código, generación de código. | Menú para ver registros en VCS, funciones disponibles: inspección de código por IA, generación de registros de lanzamiento. |
| EditorMenú emergente | todos son aceptables | Menú contextual, también puede agregar el grupo de acción correspondiente |
Al escribir ShowIntentionsGroup, podemos referirnos a la implementación de AutoDev para construir el grupo correspondiente:
< group id = " AutoDevIntentionsActionGroup " class = " cc.unitmesh.devti.intentions.IntentionsActionGroup "
icon = " cc.unitmesh.devti.AutoDevIcons.AI_COPILOT " searchable = " false " >
< add-to-group group-id = " ShowIntentionsGroup " relative-to-action = " ShowIntentionActions " anchor = " after " />
</ group >Debido a la estrategia de plataforma de Intellij, la diferencia entre ejecutar un IDE de Java (Intellij IDEA) y otros IDE como Python IDE (Pycharm) se vuelve aún mayor. Necesitamos proporcionar compatibilidad basada en productos multiplataforma. Para una introducción detallada, consulte: Compatibilidad de complementos con productos de la plataforma IntelliJ.
Primero, la arquitectura del complemento se modulariza aún más, es decir, se proporcionan diferentes módulos para diferentes idiomas. La siguiente es la arquitectura modular de AutoDev:
java/ # Java 语言插件
src/main/java/cc/unitmesh/autodev/ # Java 语言入口
src/main/resources/META-INF/plugin.xml
plugin/ # 多平台入口
src/main/resources/META-INF/plugin.xml
src/ # 即核心模块
main/resource/META-INF/core.plugin.xml En plugin/plugin.xml , debemos agregar las depends y extensions correspondientes. El siguiente es un ejemplo:
< idea-plugin package = " cc.unitmesh " xmlns : xi = " http://www.w3.org/2001/XInclude " allow-bundled-update = " true " >
< xi : include href = " /META-INF/core.xml " xpointer = " xpointer(/idea-plugin/*) " />
< content >
< module name = " cc.unitmesh.java " />
<!-- 其它模块 -->
</ content >
</ idea-plugin > En java/plugin.xml , debemos agregar las depends y extensions correspondientes. El siguiente es un ejemplo:
< idea-plugin package = " cc.unitmesh.java " >
<!-- suppress PluginXmlValidity -->
< dependencies >
< plugin id = " com.intellij.modules.java " />
< plugin id = " org.jetbrains.plugins.gradle " />
</ dependencies >
</ idea-plugin > Posteriormente, Intellij cargará automáticamente el módulo correspondiente para lograr soporte en varios idiomas. Dependiendo de los diferentes idiomas que esperamos admitir, necesitamos plugin.xml correspondiente, como por ejemplo:
cc.unitmesh.javascript.xml
cc.unitmesh.rust.xml
cc.unitmesh.python.xml
cc.unitmesh.kotlin.xml
cc.unitmesh.java.xml
cc.unitmesh.go.xml
cc.unitmesh.cpp.xmlFinalmente, simplemente implemente las funciones correspondientes en diferentes módulos de idioma.
Para simplificar este proceso, utilizamos Unit Eval para mostrar cómo construir dos contextos similares.
Mediante el análisis de código estático, podemos obtener la función actual, la clase actual, el archivo actual, etc. Combinado con la similitud de rutas, encuentre el contexto más relevante.
private fun findRelatedCode ( container : CodeContainer ): List < CodeDataStruct > {
// 1. collects all similar data structure by imports if exists in a file tree
val byImports = container. Imports
.mapNotNull {
context.fileTree[it. Source ]?.container?. DataStructures
}
.flatten()
// 2. collects by inheritance tree for some node in the same package
val byInheritance = container. DataStructures
.map {
(it. Implements + it. Extend ).mapNotNull { i ->
context.fileTree[i]?.container?. DataStructures
}.flatten()
}
.flatten()
val related = (byImports + byInheritance).distinctBy { it. NodeName }
// 3. convert all similar data structure to uml
return related
}
class RelatedCodeStrategyBuilder ( private val context : JobContext ) : CodeStrategyBuilder {
override fun build (): List < TypedIns > {
// ...
val findRelatedCodeDs = findRelatedCode(container)
val relatedCodePath = findRelatedCodeDs.map { it. FilePath }
val jaccardSimilarity = SimilarChunker .pathLevelJaccardSimilarity(relatedCodePath, currentPath)
val relatedCode = jaccardSimilarity.mapIndexed { index, d ->
findRelatedCodeDs[index] to d
}.sortedByDescending {
it.second
}.take( 3 ).map {
it.first
}
// ...
}
}Para el código anterior, podemos usar la información de Importaciones del código como parte del código relevante. Luego encuentre el código relevante a través de la relación de herencia del código. Finalmente, el contexto más cercano se encuentra a través de la similitud de rutas.
Busque primero y luego encuentre el código relacionado a través de la similitud de código. Se muestra la lógica central:
fun pathLevelJaccardSimilarity ( chunks : List < String >, text : String ): List < Double > {
// ...
}
fun tokenize ( chunk : String ): List < String > {
return chunk.split( Regex ( " [^a-zA-Z0-9] " )).filter { it.isNotBlank() }
}
fun similarityScore ( set1 : Set < String >, set2 : Set < String >): Double {
// ...
}Para obtener más información, consulte: SimilarChunker
HACER
TreeSitter es un marco para generar analizadores personalizados eficientes, desarrollado por GitHub. Utiliza un analizador LR(1), lo que significa que puede analizar cualquier idioma en tiempo O(n) en lugar de tiempo O(n²). También utiliza una técnica llamada "reutilización de árboles de sintaxis" que le permite actualizar árboles de sintaxis sin tener que analizar todo el archivo.
Dado que TreeSitter ya proporciona soporte en varios idiomas, puede utilizar Node.js, Rust y otros lenguajes para crear los complementos correspondientes. Consulte: TreeSitter para obtener más detalles.
Dependiendo de nuestras intenciones, existen diferentes formas de utilizar TreeSitter:
Símbolo de análisis
En el código del motor de búsqueda de lenguaje natural Bloop, utilizamos TreeSitter para analizar símbolos y lograr una mejor calidad de búsqueda.
; ; methods
(method_declaration
name: (identifier) @hoist.definition.method)Luego, decide cómo mostrarlo según el tipo:
pub static JAVA : TSLanguageConfig = TSLanguageConfig {
language_ids : & [ "Java" ] ,
file_extensions : & [ "java" ] ,
grammar : tree_sitter_java :: language ,
scope_query : MemoizedQuery :: new ( include_str ! ( "./scopes.scm" ) ) ,
hoverable_query : MemoizedQuery :: new (
r#"
[(identifier)
(type_identifier)] @hoverable
"# ,
) ,
namespaces : & [ & [
// variables
"local" ,
// functions
"method" ,
// namespacing, modules
"package" ,
"module" ,
// types
"class" ,
"enum" ,
"enumConstant" ,
"record" ,
"interface" ,
"typedef" ,
// misc.
"label" ,
] ] ,
} ;código fragmentado
A continuación se explica cómo se utiliza TreeSitter para mejorar el fragmentador de código de LlamaIndex mediante la limpieza de los CST de Tree-Sitter:
from tree_sitter import Tree
def chunker (
tree : Tree ,
source_code : bytes ,
MAX_CHARS = 512 * 3 ,
coalesce = 50 # Any chunk less than 50 characters long gets coalesced with the next chunk
) -> list [ Span ]:
# 1. Recursively form chunks based on the last post (https://docs.sweep.dev/blogs/chunking-2m-files)
def chunk_node ( node : Node ) -> list [ Span ]:
chunks : list [ Span ] = []
current_chunk : Span = Span ( node . start_byte , node . start_byte )
node_children = node . children
for child in node_children :
if child . end_byte - child . start_byte > MAX_CHARS :
chunks . append ( current_chunk )
current_chunk = Span ( child . end_byte , child . end_byte )
chunks . extend ( chunk_node ( child ))
elif child . end_byte - child . start_byte + len ( current_chunk ) > MAX_CHARS :
chunks . append ( current_chunk )
current_chunk = Span ( child . start_byte , child . end_byte )
else :
current_chunk += Span ( child . start_byte , child . end_byte )
chunks . append ( current_chunk )
return chunks
chunks = chunk_node ( tree . root_node )
# 2. Filling in the gaps
for prev , curr in zip ( chunks [: - 1 ], chunks [ 1 :]):
prev . end = curr . start
curr . start = tree . root_node . end_byte
# 3. Combining small chunks with bigger ones
new_chunks = []
current_chunk = Span ( 0 , 0 )
for chunk in chunks :
current_chunk += chunk
if non_whitespace_len ( current_chunk . extract ( source_code )) > coalesce
