Spaces:
Paused
Paused
File size: 2,170 Bytes
5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 84412f7 5d8d9b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import os, faiss
from pathlib import Path
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import numpy as np
class CodebaseIndex:
def __init__(self, project_path: str):
self.project_path = Path(project_path)
self.model = SentenceTransformer('all-MiniLM-L6-v2')
self.index = None
self.file_list = []
self.chunks = []
def build_index(self):
exts = {'.py', '.js', '.ts', '.tsx', '.jsx', '.go', '.rs', '.java', '.cpp', '.c', '.h'}
for filepath in self.project_path.rglob('*'):
if filepath.suffix in exts and filepath.is_file():
self.file_list.append(str(filepath.relative_to(self.project_path)))
content = filepath.read_text()
chunks = self._split_into_chunks(content)
for chunk in chunks:
self.chunks.append((str(filepath.relative_to(self.project_path)), chunk))
texts = [chunk[1] for chunk in self.chunks]
embeddings = self.model.encode(texts, show_progress_bar=False)
dim = embeddings.shape[1]
self.index = faiss.IndexFlatL2(dim)
self.index.add(np.array(embeddings).astype('float32'))
def search(self, query: str, top_k=5) -> List[Dict]:
if not self.index:
return []
q_emb = self.model.encode([query]).astype('float32')
D, I = self.index.search(q_emb, top_k)
results = []
for i, idx in enumerate(I[0]):
if idx != -1:
fpath, chunk = self.chunks[idx]
results.append({"file": fpath, "snippet": chunk, "score": float(D[0][i])})
return results
def _split_into_chunks(self, text, max_chars=1000):
paragraphs = text.split('\n\n')
chunks = []
current = ""
for para in paragraphs:
if len(current) + len(para) < max_chars:
current += para + "\n\n"
else:
if current:
chunks.append(current.strip())
current = para + "\n\n"
if current:
chunks.append(current.strip())
return chunks |