Vil du have, at AI rent faktisk forstår din kode? Dette værktøj siger, at det kan hjælpe

I denne blog vil vi vise dig, hvordan du indekserer kodebase til RAG med CocoIndex. CocoIndex er et værktøj til at hjælpe dig med at indeksere og forespørge på dine data. Det er designet til at blive brugt som en ramme til at bygge din egen datapipeline. CocoIndex giver indbygget support til kodebase-chunking med indbygget Tree-sitter-understøttelse.

Træsitter

Tree-sitter er et parser-generatorværktøj og et inkrementelt parsingbibliotek, det er tilgængeligt i Rust 🦀 - GitHub . CocoIndex har indbygget Rust-integration med Tree-sitter for effektivt at parse kode og udtrække syntakstræer til forskellige programmeringssprog.

Codebase chunking er processen med at nedbryde en kodebase i mindre, semantisk meningsfulde bidder. CocoIndex udnytter Tree-sitters evner til intelligent at dele kode baseret på den faktiske syntaksstruktur frem for vilkårlige linjeskift. Disse semantisk sammenhængende bidder bruges derefter til at bygge et mere effektivt indeks for RAG-systemer, hvilket muliggør mere præcis kodehentning og bedre kontekstbevarelse.

Fast pass 🚀 - du kan finde den fulde kode her . Kun ~ 50 linjer Python-kode til RAG-rørledningen, tjek det ud 🤗!

Giv venligst CocoIndex på Github en stjerne for at støtte os, hvis du kan lide vores arbejde. Tusind tak med et varmt kokosnøddekram 🥥🤗.

Forudsætninger

Hvis du ikke har Postgres installeret, se venligst installationsvejledningen . CocoIndex bruger Postgres til at administrere dataindekset, vi har det på vores køreplan for at understøtte andre databaser, inklusive de igangværende. Hvis du er interesseret i andre databaser, så lad os det vide ved at oprette et GitHub-problem eller Discord .

Definer cocoIndex Flow

Lad os definere cocoIndex-flowet til at læse fra en kodebase og indeksere det til RAG.

Flowdiagrammet ovenfor illustrerer, hvordan vi behandler vores kodebase:

1. Læs kodefiler fra det lokale filsystem
2. Udpak filtypenavne
3. Opdel kode i semantiske bidder ved hjælp af Tree-sitter
4. Generer indlejringer for hver del
5. Gem i en vektordatabase til hentning

Lad os implementere dette flow trin for trin.

1. Tilføj kodebasen som en kilde .

 @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()

I dette eksempel skal vi indeksere cocoindex-kodebasen fra rodmappen. Du kan ændre stien til den kodebase, du vil indeksere. Vi vil indeksere alle filerne med filtypenavnene .py , .rs , .toml , .md , .mdx , og springe mapper over startende med ., target (i roden) og node_modules (under enhver mappe).

flow_builder.add_source vil oprette en tabel med følgende underfelter, se dokumentation her.

• filename (nøgle, type: str ): filnavnet på filen, f.eks. dir1/file1.md
• content (type: str hvis binary er False , ellers bytes ): indholdet af filen

2. Behandl hver fil og indsaml oplysningerne .

2.1 Udpakfiltypenavnet

Lad os først definere en funktion til at udtrække filtypenavnet, mens hver fil behandles. Du kan finde dokumentationen for brugerdefineret funktion her .

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

Derefter skal vi behandle hver fil og indsamle oplysningerne.

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

Her udtrækker vi filtypenavnet og gemmer det i extension . for eksempel, hvis filnavnet er spec.rs , vil extension være .rs .

2.2 Del filen i bidder

Dernæst skal vi opdele filen i bidder. Vi bruger funktionen SplitRecursively til at opdele filen i bidder. Du kan finde dokumentationen til funktionen her .

CocoIndex giver indbygget support til Tree-sitter, så du kan videregive sproget til language . For at se alle understøttede sprognavne og udvidelser, se dokumentationen her . Alle de større sprog er understøttet, f.eks. Python, Rust, JavaScript, TypeScript, Java, C++ osv. Hvis det er uspecificeret, eller det angivne sprog ikke understøttes, vil det blive behandlet som almindelig tekst.

 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 Integrer bidderne

Vi vil bruge SentenceTransformerEmbed -funktionen til at indlejre bidderne. Du kan finde dokumentationen til funktionen her . Der er 12k modeller understøttet af 🤗 Hugging Face . Du kan bare vælge din yndlingsmodel.

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

Så for hver chunk vil vi indlejre den ved hjælp af code_to_embedding funktionen. og saml indlejringerne til code_embeddings samleren.

Vi udtrækker denne code_to_embedding-funktion i stedet for direkte at kalde transform(cocoindex.functions.SentenceTransformerEmbed(...)) på plads.

Dette er fordi vi ønsker at gøre denne delt mellem indekseringsflowbygningen og forespørgselshåndteringsdefinitionen. Alternativt for at gøre det nemmere. Det er også OK at undgå denne ekstra funktion og direkte gøre tingene på plads - ikke en big deal at kopiere indsætte en lille smule, vi gjorde dette til quickstart- projektet.

 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 Saml indlejringerne

Lad os endelig eksportere indlejringerne til en tabel.

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

3. Konfigurer forespørgselshandler til dit indeks

Vi vil bruge SimpleSemanticsQueryHandler til at forespørge indekset. Bemærk, at vi skal overføre funktionen code_to_embedding til parameteren query_transform_flow . Dette skyldes, at forespørgselshandleren vil bruge den samme indlejringsmodel som den, der blev brugt i flowet.

 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)

Definer en hovedfunktion til at køre forespørgselshåndteringen.

 @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()

@cocoindex.main_fn() dekoratoren initialiserer biblioteket med indstillinger indlæst fra miljøvariabler. Se dokumentationen for initialisering for flere detaljer.

Kør indeksopsætning og opdatering

🎉 Nu er du klar!

Kør følgende kommandoer for at opsætte og opdatere indekset.

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

Du vil se status for indeksopdateringer i terminalen

Test forespørgslen

På dette tidspunkt kan du starte cocoindex-serveren og udvikle din RAG-runtime mod dataene.

For at teste dit indeks er der to muligheder:

Mulighed 1: Kør indeksserveren i terminalen

 python main.py

Når du ser prompten, kan du indtaste din søgeforespørgsel. for eksempel: spec.

 Enter search query (or Enter to quit): spec

Du kan finde søgeresultaterne i terminalen

De returnerede resultater - hver post indeholder score (Cosinus-lighed), filnavn og kodestykket, der bliver matchet. På cocoindex bruger vi cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY til at måle ligheden mellem forespørgslen og de indekserede data. Du kan også skifte til andre målinger og hurtigt teste det.

For at lære mere om Consine Similarity, se Wiki .

Mulighed 2: Kør CocoInsight for at forstå din datapipeline og dataindeks

CocoInsight er et værktøj til at hjælpe dig med at forstå din datapipeline og dataindeks. Den opretter forbindelse til din lokale CocoIndex-server uden dataopbevaring.

CocoInsight er i Early Access nu (gratis) 😊 Du fandt os! En hurtig 3 minutters videovejledning om CocoInsight: Se på YouTube .

Kør CocoIndex -serveren

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

Når serveren kører, skal du åbne CocoInsight i din browser. Du vil være i stand til at oprette forbindelse til din lokale CocoIndex-server og udforske din datapipeline og indeks.

På højre side kan du se det dataflow, vi definerede.

I venstre side kan du se dataindekset i dataeksemplet.

Du kan klikke på en hvilken som helst række for at se detaljerne for denne dataindtastning, inklusive det fulde indhold af kodestykker og deres indlejringer.

Fællesskabet

Vi elsker at høre fra fællesskabet! Du kan finde os på Github og Discord .

Hvis du kan lide dette indlæg og vores arbejde, så støtte CocoIndex på Github med en stjerne ⭐. Tak med et varmt kokosnøddekram 🥥🤗.