build your ai coding assistant
v0.0.2
Em 2023, a popularidade da IA generativa levará cada vez mais organizações a introduzir a codificação assistida por IA. O que é um pouco diferente do GitHub Copilot lançado em 2021 é que a conclusão do código é apenas um dos muitos cenários. Um grande número de empresas está explorando cenários como a geração de código completo e revisão de código com base em requisitos, e também está introduzindo IA generativa para melhorar a eficiência do desenvolvimento.
Nesse contexto, nós (comunidade de código aberto da Thoughtworks) também abrimos o código-fonte de uma série de ferramentas auxiliares de IA para ajudar mais organizações a construir seus próprios assistentes de codificação assistidos por IA:
Porque quando projetamos o AutoDev, vários modelos de código aberto estavam em constante evolução. Nesse contexto, suas etapas são:
Portanto, este tutorial também está centrado nessas três etapas. Além disso, com base em nossa experiência, o exemplo de pilha de tecnologia para este tutorial:
Como nossa experiência em IA é relativamente limitada, inevitavelmente haverá alguns erros. Portanto, também esperamos trabalhar com mais desenvolvedores para construir este projeto de código aberto.
Combinado com a parte de inteligência artificial do relatório “Ecossistema de Desenvolvedor” da JetBrains 2023, podemos resumir alguns cenários gerais que refletem as áreas onde a IA generativa pode desempenhar um papel no processo de desenvolvimento. Aqui estão alguns dos principais cenários:
Quando construímos o AutoDev, também descobrimos cenários como criação de SQL DDL, geração de requisitos, TDD, etc. então. Fornecemos a capacidade de personalizar cenários para que os desenvolvedores possam personalizar seus próprios recursos de IA. Para obter detalhes, consulte: https://ide.unitmesh.cc/customize.
Na codificação diária, existem vários cenários diferentes com diferentes requisitos de velocidade de resposta da IA (apenas como exemplo):
| cena | Velocidade de resposta | Gerar requisitos de qualidade | Tamanho esperado | ilustrar |
|---|---|---|---|---|
| conclusão de código | rápido | meio | 1~6B | A conclusão do código é o cenário mais comum na codificação diária e a velocidade de resposta é crucial. |
| Geração de documentos | meio | meio | 1 | A geração de documentação requer um entendimento completo da estrutura do código, e velocidade e qualidade são igualmente importantes. |
| revisão de código | rápido | meio | 1 | As revisões de código exigem aconselhamento de alta qualidade, mas também precisam ser o mais responsivas possível. |
| Geração de teste unitário | rápido | meio | 6B ~ | Os testes unitários geram menos contexto, e a capacidade de resposta e a qualidade da IA são igualmente importantes. |
| refatoração de código | meio | alto | 32B~ | A refatoração de código pode exigir mais compreensão contextual e os tempos de resposta podem ser moderadamente mais lentos. |
| geração de demanda | meio | alto | 32B~ | A geração de demanda é um cenário relativamente complexo e a velocidade de resposta pode ser moderadamente desacelerada para garantir a precisão. |
| Pesquisa e interpretação de código em linguagem natural | Médio-Baixo | alto | 32B~ | A pesquisa e interpretação de código em linguagem natural são cenários relativamente complexos e a velocidade de resposta pode ser moderadamente desacelerada para garantir a precisão. |
PS: O 32B aqui é expresso apenas como uma ordem de grandeza, pois o efeito será melhor com um modelo maior.
Portanto, resumimo-lo como: um modelo grande, um médio, um micro e três modelos, fornecendo codificação abrangente assistida por IA:
A conclusão de código de IA pode combinar ferramentas IDE para analisar o contexto do código e as regras da linguagem de programação, e a IA gerará ou sugerirá automaticamente trechos de código. Em ferramentas de autocompletar código semelhantes ao GitHub Copilot, elas geralmente são divididas em três modos de subdivisão:
Conclusão inline (inline)
Semelhante ao modo FIM (preencher no meio), o conteúdo concluído está na linha atual. Por exemplo: BlotPost blogpost = new , a conclusão é: BlogPost(); para alcançar: BlogPost blogpost = new BlogPost();
Podemos usar o Deepseek Coder como exemplo para ver o efeito neste cenário:
< ｜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｜ >Aqui, precisamos combinar o código antes e depois do cursor.
Conclusão no bloco (InBlock)
Obtido por meio de aprendizagem contextual (Aprendizagem In-Context), o conteúdo de conclusão está no bloco de função atual. Por exemplo, o código original é:
fun createBlog ( blogDto : CreateBlogDto ): BlogPost {
}O código completo é:
val blogPost = BlogPost (
title = blogDto.title,
content = blogDto.content,
author = blogDto.author
)
return blogRepository.save(blogPost)Pós-Bloco
Alcançado por meio de aprendizagem contextual (Aprendizagem In-Context), conclusão após o bloco de função atual, como: conclusão de uma nova função após o bloco de função atual. Por exemplo, o código original é:
fun createBlog ( blogDto : CreateBlogDto ): BlogPost {
// ...
}O código completo é:
fun updateBlog ( id : Long , blogDto : CreateBlogDto ): BlogPost {
// ...
}
fun deleteBlog ( id : Long ) {
// ...
}Quando construímos a função de conclusão de IA correspondente, também precisamos considerar aplicá-la ao conjunto de dados padrão correspondente para melhorar a qualidade da conclusão e fornecer uma melhor experiência ao usuário.
Alguns recursos relacionados para escrever este artigo:
As explicações de código são projetadas para ajudar os desenvolvedores a gerenciar e compreender grandes bases de código de maneira mais eficaz. Esses assistentes podem responder perguntas sobre a base de código, fornecer documentação, pesquisar código, identificar fontes de erros, reduzir a duplicação de código, etc., melhorando assim a eficiência do desenvolvimento, reduzindo as taxas de erros e reduzindo a carga de trabalho dos desenvolvedores.
Neste cenário, dependendo da qualidade da geração que esperamos, geralmente é composto por dois modelos: um grande e um micro ou um médio e um micro. O modelo maior apresenta melhores resultados em termos de qualidade de geração. Combinada com nossa experiência de design na ferramenta Fábrica de Chocolate, normalmente essa função pode ser dividida em várias etapas:
Por ser uma aplicação RAG, ela é dividida em duas partes: indexação e consulta.
Na etapa de indexação, precisamos indexar a base de código, que envolve segmentação de texto, vetorização, indexação de banco de dados e outras tecnologias. Um dos elementos mais desafiadores é a divisão. As regras de divisão às quais nos referimos são: https://docs.sweep.dev/blogs/chunking-2m-files. Agora mesmo:
Em diferentes cenários, também podemos dividir de diferentes maneiras. Por exemplo, na Chocolate Factory, dividimos através do AST para garantir a qualidade do contexto gerado.
Na fase de consulta, precisamos combinar algumas de nossas tecnologias de busca tradicionais, como busca por vetorização, busca por caminho, etc., para garantir a qualidade da busca. Ao mesmo tempo, no cenário chinês, também precisamos considerar a questão da conversão para o chinês, como a conversão do inglês para o chinês para garantir a qualidade da pesquisa.
Para assistência diária, também podemos alcançá-lo por meio de IA generativa, como criação automática de SQL DDL, criação automática de casos de teste, criação automática de requisitos, etc. Isso só pode ser alcançado personalizando palavras de prompt e combinando conhecimento de domínio específico, por isso não entrarei em detalhes aqui.
Além do modelo, o contexto também é um fator importante que afeta as capacidades de assistência da IA. Quando construímos o AutoDev, também descobrimos dois modos de contexto diferentes:
Uma comparação simples é a seguinte:
| contexto relevante | contexto semelhante | |
|---|---|---|
| Tecnologia de pesquisa | análise de código estático | Pesquisa de similaridade |
| informações de estrutura de dados | AST, CFG | Pedaço semelhante |
| Capacidades multiplataforma | Depende do IDE ou analisador independente | Não depende de plataformas específicas |
| qualidade contextual | extremamente alto | alto |
| Gerar resultados | extremamente alto | alto |
| custo de construção | Depende do idioma e da plataforma | Baixo |
Quando o suporte para IDE é limitado, o contexto trará desempenho de custo mais alto.
GitHub Copilot adota um padrão de arquitetura de contexto semelhante e sua arquitetura detalhada é dividida em camadas da seguinte forma:
Nos materiais de pesquisa do projeto "público" Copilot-Explorer, você pode ver como o Prompt é construído. A seguir está a solicitação imediata enviada para:
{
"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
}
]
}em:
prefix usada para construir o prompt é construída a partir de promptElements, que inclui: BeforeCursor , AfterCursor , SimilarFile , ImportedFile , LanguageMarker , PathMarker , RetrievalSnippet e outros tipos. Pelos nomes de vários PromptElementKind , também podemos ver seu verdadeiro significado.suffix utilizada para construir o prompt é determinada pela parte onde o cursor está localizado. De acordo com o limite superior de tokens (2048), quantas posições faltam para serem calculadas. O cálculo do token aqui é o cálculo do token LLM real no Copilot, é calculado por Cushman002. O comprimento do token dos caracteres chineses é diferente, como: { context: "console.log('你好，世界')", lineCount: 1, tokenLength: 30 } , onde o comprimento do conteúdo no contexto é 20, mas tokenLength é 30, o comprimento dos caracteres chineses é 5 (incluindo ， ) e o token ocupado por um único caractere é 3.Aqui está um exemplo mais detalhado de um contexto de aplicativo 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 {
//...No contexto de computação, o GitHub Copilot usa o coeficiente Jaccard (Jaccard Similarity). Esta parte da implementação é implementada no Agent. Para uma lógica mais detalhada, consulte: Depois de passar mais de meio mês, finalmente fiz a engenharia reversa do Github Copilot.
Recursos relacionados:
Conforme mencionado acima, o código relevante depende da análise estática do código , principalmente com a ajuda de informações estruturais do código, como AST, CFG, DDG, etc. Em diferentes cenários e plataformas, podemos combinar diferentes ferramentas de análise de código estático. A seguir estão algumas ferramentas comuns de análise de código estático:
No cenário de conclusão, por meio da análise estática do código, podemos obter o contexto atual, como: função atual, classe atual, arquivo atual, etc. A seguir está um exemplo do contexto do AutoDev para gerar testes de unidade:
// 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 ) { Neste exemplo, o contexto da função createBlog é analisado para obter as classes de entrada e saída da função: CreateBlogRequest , informações BlogPost e informações de classe BlogService, que são fornecidas ao modelo como contexto (fornecidas em comentários). Neste ponto, o modelo gera construtores mais precisos, bem como casos de teste mais precisos.
Como o contexto relevante depende da análise estática de código de diferentes linguagens e APIs de diferentes IDEs, também precisamos nos adaptar a diferentes linguagens e diferentes IDEs. Em termos de custo de construção, é mais caro em relação a contextos semelhantes.
IDEs e editores são as principais ferramentas para desenvolvedores e seus custos de design e aprendizado são relativamente altos. Primeiro, podemos usar o modelo oficial para gerar:
Depois, adicione funcionalidade por cima (não é muito simples), claro que não. A seguir estão alguns recursos do plug-in IDEA para referência:
Claro, é mais apropriado consultar o plug-in AutoDev.
Você pode usar diretamente o modelo oficial para gerar o plug-in correspondente: https://github.com/JetBrains/intellij-platform-plugin-template
Para implementação do plug-in IDEA, ele é implementado principalmente por meio de Action e Listener, que só precisam ser registrados em plugin.xml . Para obter detalhes, consulte a documentação oficial: IntelliJ Platform Plugin SDK
Como não consideramos o problema de compatibilidade com a versão IDE do AutoDev no estágio inicial, para sermos compatíveis com a versão antiga do IDE posteriormente, precisamos realizar o processamento de compatibilidade no plug-in. Portanto, conforme descrito no documento oficial: Build Number Ranges, podemos ver que diferentes versões têm requisitos diferentes para o JDK.
| Número da filial | Versão da plataforma IntelliJ |
|---|---|
| 233 | 2023.3 |
| 232 | 2023.2 |
| 231 | 2023.1 |
| 223 | 2022.3 |
| 222 | 2022.2 NOTA Java 17 agora é necessário (postagem no blog) |
| 221 | 2022.1 |
| 213 | 2021.3 |
| 212 | 2021.2 |
| 211 | 2021.1 |
| 203 | 2020.3 NOTA Java 11 agora é necessário (postagem no blog) |
E configure-o em gradle.properties :
pluginSinceBuild = 223
pluginUntilBuild = 233.*A configuração subsequente de compatibilidade é problemática, então você pode consultar o design do AutoDev.
Em termos de preenchimento automático de código, os fabricantes nacionais referem-se principalmente à implementação do GitHub Copilot, e a lógica não é complicada.
Acionar usando teclas de atalho
Ele monitora principalmente a entrada do usuário em Ação e, em seguida:
| Função | tecla de atalho | ilustrar |
|---|---|---|
| requestCompletions | Alt + / | Obtenha o contexto atual e, em seguida, obtenha os resultados de conclusão por meio do modelo |
| aplicarInlays | TAB | Exibir resultados de conclusão no IDE |
| descartarInlays | ESC | Cancelar conclusão |
| cicloNextInlays | Alt + ] | Mudar para o próximo resultado de conclusão |
| cicloPrevInlays | Alt + [ | Mudar para o resultado de conclusão anterior |
Use método de disparo automático
Ele monitora principalmente a entrada do usuário por meio de EditorFactoryListener e, em seguida, aciona diferentes resultados de conclusão com base em diferentes entradas. O código principal é o seguinte:
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)
}
}
}Então, de acordo com diferentes entradas, diferentes resultados de conclusão são acionados e a estrutura é processada.
Código de conclusão de renderização
Posteriormente, precisamos implementar um Inlay Render, que herda de EditorCustomElementRenderer .
Combinado com os recursos de interface do IDE, precisamos adicionar a Ação correspondente, o Grupo correspondente e o Ícone correspondente. A seguir está um exemplo de Ação:
<add-to-group group-id="ShowIntentionsGroup" relative-to-action="ShowIntentionActions" anchor="after"/>
A seguir estão alguns ActionGroups do AutoDev:
| ID do grupo | Usos de IA | Descrição |
|---|---|---|
| MostrarIntençõesGrupo | Refatoração de código, interpretação de código, geração de código, teste de código | Usado para exibir dicas no contexto do código e acessado por meio Alt + Enter e ⌥ + Enter no macOS. |
| ConsoleEditorPopupMenu | corrigir erros | O menu exibido no console, como o console da estrutura de execução do programa. |
| Vcs.MessageActionGroup | Geração de informações de código | Menu para escrever mensagens de commit no VCS. |
| Vcs.Log.ContextMenu | Revisão de código, interpretação de código, geração de código | Menu para visualização de logs no VCS, funções disponíveis: inspeção de código por IA, geração de logs de lançamento. |
| EditorPopupMenu | Todos são aceitáveis | Menu do botão direito, você também pode adicionar o ActionGroup correspondente |
Ao escrever ShowIntentionsGroup, podemos nos referir à implementação do AutoDev para construir o Grupo correspondente:
< 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 >Devido à estratégia de plataforma do Intellij, a diferença entre rodar em um IDE Java (Intellij IDEA) e outros IDEs como Python IDE (Pycharm) torna-se ainda maior. Precisamos fornecer compatibilidade com base em produtos multiplataforma. Para introdução detalhada, consulte: Compatibilidade de plug-ins com produtos da plataforma IntelliJ.
Primeiro, a arquitetura do plug-in é ainda mais modularizada, ou seja, diferentes módulos são fornecidos para diferentes linguagens. A seguir está a arquitetura modular do 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 Em plugin/plugin.xml , precisamos adicionar as depends e extensions correspondentes.
< 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 > Em java/plugin.xml , precisamos adicionar as depends e extensions correspondentes.
< 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, o Intellij carregará automaticamente o módulo correspondente para obter suporte multilíngue. Dependendo dos diferentes idiomas que esperamos oferecer suporte, precisamos plugin.xml correspondente, como:
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.xmlPor fim, basta implementar as funções correspondentes em diferentes módulos de linguagem.
Para simplificar este processo, usamos Unit Eval para mostrar como construir dois contextos semelhantes.
Através da análise estática de código, podemos obter a função atual, a classe atual, o arquivo atual, etc. Em seguida, combine a similaridade do caminho para encontrar o contexto mais 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 o código acima, podemos usar as informações de importação do código como parte do código relevante. Em seguida, encontre o código relevante por meio do relacionamento de herança do código. Finalmente, o contexto mais próximo é encontrado através da similaridade de caminhos.
Pesquise primeiro e depois encontre o código relacionado por meio de similaridade de código. A lógica central é mostrada:
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 obter detalhes, consulte: SimilarChunker
PENDÊNCIA
TreeSitter é uma estrutura para gerar analisadores personalizados eficientes, desenvolvida pelo GitHub. Ele usa um analisador LR(1), o que significa que pode analisar qualquer idioma em tempo O(n) em vez de tempo O(n²). Ele também usa uma técnica chamada "reutilização de árvore de sintaxe" que permite atualizar árvores de sintaxe sem analisar novamente o arquivo inteiro.
Como o TreeSitter já oferece suporte multilíngue, você pode usar Node.js, Rust e outras linguagens para construir plug-ins correspondentes. Veja: TreeSitter para detalhes.
Dependendo de nossas intenções, existem diferentes maneiras de usar o TreeSitter:
Símbolo de análise
No mecanismo de busca de código em linguagem natural Bloop, usamos TreeSitter para analisar símbolos para obter melhor qualidade de pesquisa.
; ; methods
(method_declaration
name: (identifier) @hoist.definition.method)Em seguida, decida como exibi-lo com base no 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 do pedaço
A seguir está como o TreeSitter é usado para melhorar o Code Chunker do LlamaIndex limpando os CSTs do 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
