shovo896/codedocumentation

1import hashlib2import os3import re4import shutil5import sys6import tempfile7from pathlib import Path8 9from dotenv import load_dotenv10from langchain_core.documents import Document11from langchain_openai import OpenAIEmbeddings12from langchain_pinecone import PineconeVectorStore13from langchain_text_splitters import RecursiveCharacterTextSplitter14from pinecone import Pinecone, ServerlessSpec15 16load_dotenv()17 18 19def get_required_env(name: str) -> str:20    value = os.getenv(name, "").strip().strip('"').strip("'")21    if not value:22        raise RuntimeError(f"Missing required environment variable: {name}")23    return value24 25 26PINECONE_API_KEY = get_required_env("PINECONE_API_KEY")27INDEX_NAME = "code-doc-search-openai"28EMBED_MODEL = "text-embedding-3-small"29EMBED_DIM = 153630MAX_FILE_SIZE_KB = 150031CHUNK_SIZE = 100032CHUNK_OVERLAP = 20033 34ALLOWED_EXTENSIONS = {35    ".py",36    ".md",37    ".rst",38    ".txt",39    ".js",40    ".jsx",41    ".ts",42    ".tsx",43    ".json",44    ".html",45    ".css",46    ".scss",47    ".sass",48    ".less",49    ".vue",50    ".svelte",51    ".java",52    ".go",53    ".rs",54    ".php",55    ".rb",56    ".cs",57    ".cpp",58    ".c",59    ".h",60    ".yml",61    ".yaml",62    ".toml",63}64SKIP_DIRS = {"node_modules", "__pycache__", ".git", "dist", "build", "venv", ".venv"}65SKIP_FILES = {66    "package-lock.json",67    "yarn.lock",68    "pnpm-lock.yaml",69    "bun.lockb",70    "poetry.lock",71    "pipfile.lock",72}73 74 75def normalize_repo_url(repo_url: str) -> str:76    """Return a normalized public GitHub repository URL."""77    repo_url = repo_url.strip().removesuffix("/")78    repo_url = repo_url.removesuffix(".git")79 80    match = re.fullmatch(r"https://github\.com/([A-Za-z0-9_.-]+)/([A-Za-z0-9_.-]+)", repo_url)81    if not match:82        raise ValueError("Please provide a public GitHub repository URL like https://github.com/owner/repo")83 84    owner, repo = match.groups()85    return f"https://github.com/{owner}/{repo}"86 87 88def namespace_for_repo(repo_url: str, branch: str) -> str:89    """Create a stable Pinecone namespace for a repository and branch."""90    normalized = normalize_repo_url(repo_url)91    branch = branch.strip()92    digest = hashlib.sha1(f"{normalized}:{branch}".encode("utf-8")).hexdigest()[:16]93    return f"repo-{digest}"94 95 96def should_load(file_path: str) -> bool:97    """Load only small, relevant source and documentation files."""98    p = Path(file_path)99 100    for part in p.parts:101        if part in SKIP_DIRS:102            return False103 104    if p.suffix.lower() not in ALLOWED_EXTENSIONS:105        return False106 107    name = p.name.lower()108    if name in SKIP_FILES:109        return False110 111    if "test" in name or "spec" in name:112        return False113 114    return True115 116 117def clone_and_load(repo_url: str, branch: str = ""):118    """Clone the GitHub repository and load relevant files as documents."""119    tmp_dir = tempfile.mkdtemp()120    print(f"Temp directory: {tmp_dir}")121 122    try:123        from git import Blob, Repo124 125        repo = Repo.clone_from(repo_url, tmp_dir)126        requested_branch = (branch or "").strip()127        if requested_branch:128            repo.git.checkout(requested_branch)129 130        try:131            resolved_branch = repo.active_branch.name132        except TypeError:133            resolved_branch = repo.git.rev_parse("--abbrev-ref", "HEAD")134 135        docs = []136        for item in repo.tree().traverse():137            if not isinstance(item, Blob):138                continue139 140            file_path = os.path.join(tmp_dir, item.path)141            if repo.ignored([file_path]):142                continue143 144            if not should_load(file_path):145                continue146 147            try:148                with open(file_path, "rb") as f:149                    content = f.read()150                    text_content = content.decode("utf-8")151            except UnicodeDecodeError:152                continue153 154            file_type = os.path.splitext(item.name)[1]155            docs.append(156                Document(157                    page_content=text_content,158                    metadata={159                        "source": item.path,160                        "file_path": item.path,161                        "file_name": item.name,162                        "file_type": file_type,163                    },164                )165            )166 167        filtered = []168        skipped = 0169        for doc in docs:170            size_kb = len(doc.page_content.encode("utf-8")) / 1024171            if size_kb > MAX_FILE_SIZE_KB:172                skipped += 1173                continue174 175            source = doc.metadata.get("source", "")176            try:177                source_path = Path(source).resolve().relative_to(Path(tmp_dir).resolve())178                doc.metadata["source"] = source_path.as_posix()179            except (OSError, ValueError):180                doc.metadata["source"] = str(source).replace("\\", "/")181 182            filtered.append(doc)183 184        print(f"Loaded files: {len(filtered)} | Skipped large files: {skipped}")185        print(f"Resolved branch: {resolved_branch}")186        return filtered, resolved_branch187 188    except Exception as e:189        raise RuntimeError(190            f"Clone failed: {e}. Check that the repository is public and the branch exists."191        ) from e192 193    finally:194        shutil.rmtree(tmp_dir, ignore_errors=True)195 196 197def chunk_documents(docs):198    """Split loaded documents into searchable chunks."""199    splitter = RecursiveCharacterTextSplitter(200        chunk_size=CHUNK_SIZE,201        chunk_overlap=CHUNK_OVERLAP,202        separators=["\nclass ", "\ndef ", "\n\n", "\n", " ", ""],203    )204    chunks = splitter.split_documents(docs)205    print(f"Total chunks: {len(chunks)}")206    return chunks207 208 209def setup_pinecone():210    """Create the Pinecone index if it does not already exist."""211    pc = Pinecone(api_key=PINECONE_API_KEY)212    existing = [i.name for i in pc.list_indexes()]213 214    if INDEX_NAME not in existing:215        print(f"Creating Pinecone index '{INDEX_NAME}'...")216        pc.create_index(217            name=INDEX_NAME,218            dimension=EMBED_DIM,219            metric="cosine",220            spec=ServerlessSpec(cloud="aws", region="us-east-1"),221        )222        print("Index created.")223    else:224        print(f"Index '{INDEX_NAME}' already exists.")225 226    return pc227 228 229def clear_namespace(pc: Pinecone, namespace: str) -> None:230    """Clear existing vectors for a repo namespace before re-ingesting it."""231    try:232        index = pc.Index(INDEX_NAME)233        index.delete(delete_all=True, namespace=namespace)234        print(f"Cleared namespace '{namespace}'.")235    except Exception as e:236        if "Namespace not found" in str(e):237            print(f"No existing namespace found for '{namespace}'.")238        else:239            print(f"Namespace clear skipped: {e}")240 241 242def store_in_pinecone(chunks, namespace: str):243    """Embed chunks with OpenAI and store them in Pinecone."""244    print("Loading OpenAI embedding model...")245    embeddings = OpenAIEmbeddings(246        model=EMBED_MODEL,247        dimensions=EMBED_DIM,248        api_key=get_required_env("OPENAI_API_KEY"),249    )250 251    pc = setup_pinecone()252    clear_namespace(pc, namespace)253 254    print(f"Storing {len(chunks)} chunks in Pinecone namespace '{namespace}'.")255    PineconeVectorStore.from_documents(256        chunks,257        embeddings,258        index_name=INDEX_NAME,259        namespace=namespace,260    )261    print("Ingestion complete. Data has been stored in Pinecone.")262 263 264def ingest_repository(repo_url: str, branch: str = ""):265    """Ingest any public GitHub repository and return the namespace metadata."""266    repo_url = normalize_repo_url(repo_url)267    requested_branch = (branch or "").strip()268 269    docs, resolved_branch = clone_and_load(repo_url, requested_branch)270    if not docs:271        raise ValueError("No files were loaded. Check the repository URL or branch name.")272 273    for doc in docs:274        doc.metadata["repo_url"] = repo_url275        doc.metadata["branch"] = resolved_branch276 277    chunks = chunk_documents(docs)278    namespace = namespace_for_repo(repo_url, resolved_branch)279    store_in_pinecone(chunks, namespace)280 281    return {282        "repo_url": repo_url,283        "branch": resolved_branch,284        "namespace": namespace,285        "files": len(docs),286        "chunks": len(chunks),287    }288 289 290def main():291    print("=" * 50)292    print("   Code Documentation Search - Ingestion")293    print("=" * 50)294 295    if len(sys.argv) > 1:296        repo_url = sys.argv[1].strip()297    else:298        repo_url = input("\nGitHub repo URL:\n(e.g. https://github.com/tiangolo/fastapi): ").strip()299 300    if len(sys.argv) > 2:301        branch = sys.argv[2].strip()302    else:303        branch = input("Branch name (press Enter to use the default branch): ").strip()304 305    print(f"\nStarting ingestion for: {repo_url} [{branch}]\n")306 307    try:308        result = ingest_repository(repo_url, branch)309    except Exception as e:310        print(f"Error: {e}")311        sys.exit(1)312 313    print(314        "Ingested {files} files into {chunks} chunks for {repo_url} [{branch}].".format(315            **result316        )317    )318    print(f"Namespace: {result['namespace']}")319 320 321if __name__ == "__main__":322    main()323