Spaces:
Paused
Paused
| 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 |