Voleu que la IA entengui realment el vostre codi? Aquesta eina diu que pot ajudar

En aquest bloc, us mostrarem com indexar la base de codi per a RAG amb CocoIndex. CocoIndex és una eina per ajudar-vos a indexar i consultar les vostres dades. Està dissenyat per utilitzar-lo com a marc per crear el vostre propi canal de dades. CocoIndex proporciona suport integrat per a la fragmentació de la base de codi, amb suport natiu de Tree-sitter.

Arbredor

Tree-sitter és una eina generadora d'analitzadors i una biblioteca d'anàlisi incremental, està disponible a Rust 🦀 - GitHub . CocoIndex té integrada la integració Rust amb Tree-sitter per analitzar codi i extreure arbres de sintaxi de manera eficient per a diversos llenguatges de programació.

La fragmentació de la base de codi és el procés de desglossar una base de codi en fragments més petits i significatius semànticament. CocoIndex aprofita les capacitats de Tree-sitter per fragmentar el codi de manera intel·ligent basant-se en l'estructura de sintaxi real en lloc de salts de línia arbitraris. Aquests fragments semànticament coherents s'utilitzen després per construir un índex més eficaç per als sistemes RAG, permetent una recuperació de codi més precisa i una millor preservació del context.

Passatge ràpid 🚀: podeu trobar el codi complet aquí . Només ~ 50 línies de codi Python per al pipeline RAG, comproveu-ho 🤗!

Si us plau, doneu una estrella a CocoIndex a Github per donar-nos suport si us agrada el nostre treball. Moltes gràcies amb una càlida abraçada de coco 🥥🤗.

Requisits previs

Si no teniu Postgres instal·lat, consulteu la guia d'instal·lació . CocoIndex utilitza Postgres per gestionar l'índex de dades, el tenim al nostre full de ruta per donar suport a altres bases de dades, incloses les en curs. Si esteu interessats en altres bases de dades, feu-nos-ho saber creant un problema de GitHub o Discord .

Definiu el flux cocoIndex

Definim el flux cocoIndex per llegir des d'una base de codi i indexem-lo per a RAG.

El diagrama de flux anterior il·lustra com processarem la nostra base de codi:

1. Llegiu fitxers de codi del sistema de fitxers local
2. Extreu extensions de fitxer
3. Dividiu el codi en fragments semàntics mitjançant Tree-sitter
4. Genereu incrustacions per a cada tros
5. Emmagatzemar en una base de dades vectorial per a la seva recuperació

Implementem aquest flux pas a pas.

1. Afegiu la base de codi com a font .

 @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 aquest exemple, indexarem la base de codis cocoindex des del directori arrel. Podeu canviar el camí a la base de codi que voleu indexar. Indexarem tots els fitxers amb les extensions de .py , .rs , .toml , .md , .mdx , i saltarem directoris que comencen per ., target (a l'arrel) i node_modules (a qualsevol directori).

flow_builder.add_source crearà una taula amb els subcamps següents, vegeu la documentació aquí.

• filename (clau, tipus: str ): el nom de fitxer del fitxer, p. ex. dir1/file1.md
• content (tipus: str si binary és False , en cas contrari bytes ): el contingut del fitxer

2. Processa cada fitxer i recull la informació .

2.1 Extreu l'extensió d'unnom de fitxer

Primer anem a definir una funció per extreure l'extensió d'un nom de fitxer mentre processem cada fitxer. Podeu trobar la documentació de la funció personalitzada aquí .

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

A continuació, processarem cada fitxer i recollirem la informació.

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

Aquí extreu l'extensió del nom del fitxer i l'emmagatzemem al camp extension . per exemple, si el nom del fitxer és spec.rs , el camp extension serà .rs .

2.2 Dividiu el fitxer en trossos

A continuació, dividirem el fitxer en trossos. Utilitzem la funció SplitRecursively per dividir el fitxer en trossos. Podeu trobar la documentació de la funció aquí .

CocoIndex proporciona suport integrat per a Tree-sitter, de manera que podeu passar l'idioma al paràmetre language . Per veure tots els noms i extensions d'idiomes admesos, consulteu la documentació aquí . S'admeten tots els llenguatges principals, per exemple, Python, Rust, JavaScript, TypeScript, Java, C++, etc. Si no s'especifica o el llenguatge especificat no és compatible, es tractarà com a text sense format.

 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 Insereix els fragments

Utilitzarem la funció SentenceTransformerEmbed per incrustar els fragments. Podeu trobar la documentació de la funció aquí . Hi ha 12k models compatibles amb 🤗 Hugging Face . Només pots triar el teu model preferit.

 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"))

A continuació, per a cada tros, l'incorporarem mitjançant la funció code_to_embedding . i recull les incrustacions al col·lector code_embeddings .

Extraïm aquesta funció code_to_embedding en lloc de cridar directament transform(cocoindex.functions.SentenceTransformerEmbed(...)) al seu lloc.

Això es deu al fet que volem que aquest es comparteixi entre la construcció del flux d'indexació i la definició del controlador de consultes. Alternativament, per fer-ho més senzill. També està bé evitar aquesta funció addicional i fer les coses directament al seu lloc; no és gran cosa copiar i enganxar una mica, ho vam fer per al projecte d'inici ràpid .

 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 Recolliu les incrustacions

Finalment, exportem les incrustacions a una taula.

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

3. Configura el gestor de consultes per al teu índex

Utilitzarem el SimpleSemanticsQueryHandler per consultar l'índex. Tingueu en compte que hem de passar la funció code_to_embedding al paràmetre query_transform_flow . Això es deu al fet que el gestor de consultes utilitzarà el mateix model d'incrustació que l'utilitzat al flux.

 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)

Definiu una funció principal per executar el controlador de consultes.

 @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() inicialitza la biblioteca amb la configuració carregada des de les variables d'entorn. Consulteu la documentació per a la inicialització per obtenir més detalls.

Executeu la configuració i l'actualització de l'índex

🎉 Ara ja esteu a punt!

Executeu les ordres següents per configurar i actualitzar l'índex.

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

Veureu l'estat de les actualitzacions de l'índex al terminal

Prova la consulta

En aquest punt, podeu iniciar el servidor cocoindex i desenvolupar el vostre temps d'execució RAG amb les dades.

Per provar el vostre índex, hi ha dues opcions:

Opció 1: executeu el servidor d'índex al terminal

 python main.py

Quan vegeu la sol·licitud, podeu introduir la vostra consulta de cerca. per exemple: spec.

 Enter search query (or Enter to quit): spec

Podeu trobar els resultats de la cerca al terminal

Els resultats retornats: cada entrada conté la puntuació (semblança del cosinus), el nom del fitxer i el fragment de codi que coincideix. A cocoindex, fem servir el cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY per mesurar la similitud entre la consulta i les dades indexades. També podeu canviar a altres mètriques i provar-les ràpidament.

Per obtenir més informació sobre Consine Similarity, vegeu Wiki .

Opció 2: executeu CocoInsight per entendre el vostre canal de dades i l'índex de dades

CocoInsight és una eina per ajudar-vos a entendre el vostre canal de dades i l'índex de dades. Es connecta al vostre servidor local de CocoIndex sense retenció de dades.

CocoInsight ja està en accés anticipat (gratuït) 😊 Ens has trobat! Un tutorial ràpid en vídeo de 3 minuts sobre CocoInsight: mira a YouTube .

Executeu el servidor CocoIndex

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

Un cop el servidor s'està executant, obriu CocoInsight al vostre navegador. Podreu connectar-vos al vostre servidor local de CocoIndex i explorar el vostre canal de dades i l'índex.

A la part dreta, podeu veure el flux de dades que hem definit.

A la part esquerra, podeu veure l'índex de dades a la vista prèvia de dades.

Podeu fer clic a qualsevol fila per veure els detalls d'aquesta entrada de dades, inclòs el contingut complet dels fragments de codi i les seves incrustacions.

Comunitat

Ens encanta escoltar de la comunitat! Ens podeu trobar a Github i Discord .

Si us agrada aquesta publicació i el nostre treball, doneu suport a CocoIndex a Github amb una estrella ⭐. Gràcies amb una càlida abraçada de coco 🥥🤗.