SharathReddy/CodexPilot

1import os2import hmac3import hashlib4import time5import jwt6import requests7import tempfile8import shutil9import ast10import faiss11import numpy as np12from fastapi import FastAPI, Request, HTTPException, status13from dotenv import load_dotenv14from git import Repo15from sentence_transformers import SentenceTransformer16 17# --- Configuration ---18load_dotenv()19GITHUB_WEBHOOK_SECRET = os.getenv("GITHUB_WEBHOOK_SECRET")20GITHUB_APP_ID = os.getenv("GITHUB_APP_ID")21GITHUB_PRIVATE_KEY = os.getenv("GITHUB_PRIVATE_KEY")22 23# In-memory storage for our repository data24repo_data_store = {}25 26# --- Code Processing and Vectorization Class ---27 28class CodeProcessor:29    def __init__(self, model_name='all-MiniLM-L6-v2'):30        print("Initializing SentenceTransformer model...")31        self.model = SentenceTransformer(model_name)32        print("Model initialized.")33 34    def parse_python_file(self, file_path):35        """Parses a Python file to extract functions and their source code using AST."""36        # --- THE FIX IS HERE ---37        # First, read the entire file content into a string variable.38        with open(file_path, 'r', encoding='utf-8') as file:39            source_code = file.read()40        41        try:42            # Parse the source code string.43            tree = ast.parse(source_code)44        except SyntaxError:45            print(f"SyntaxError in {file_path}, skipping.")46            return []47        48        functions = []49        for node in ast.walk(tree):50            if isinstance(node, ast.FunctionDef):51                # Now, pass the correct `source_code` string to get the segment.52                function_source = ast.get_source_segment(source_code, node)53                if function_source: # Ensure we got a valid segment54                    functions.append({55                        "name": node.name,56                        "source": function_source,57                        "file_path": file_path.replace(tempfile.gettempdir(), "") # Store relative path58                    })59        return functions60 61    def process_repo(self, repo_path):62        """Walks a repository, parses Python files, and creates a FAISS index."""63        all_functions = []64        print(f"Walking through Python files in {repo_path}...")65        for root, _, files in os.walk(repo_path):66            for file in files:67                if file.endswith('.py'):68                    # Create the full path relative to the repo root for parsing69                    full_path = os.path.join(root, file)70                    # Create a "clean" path for storage in metadata71                    clean_path = os.path.relpath(full_path, repo_path)72                    73                    parsed_funcs = self.parse_python_file(full_path)74                    for func in parsed_funcs:75                        func['file_path'] = clean_path # Overwrite with the clean path76                    all_functions.extend(parsed_funcs)77        78        if not all_functions:79            print("No Python functions found in the repository.")80            return None, None81 82        print(f"Found {len(all_functions)} functions. Generating embeddings...")83        function_sources = [f["source"] for f in all_functions]84        embeddings = self.model.encode(function_sources, show_progress_bar=False)85        86        embedding_dim = embeddings.shape[1]87        index = faiss.IndexFlatL2(embedding_dim)88        index.add(np.array(embeddings, dtype=np.float32))89        90        print(f"FAISS index created successfully with {index.ntotal} vectors.")91        92        return index, all_functions93 94# --- Initialize the processor globally ---95code_processor = CodeProcessor()96 97 98# --- GitHub App Authentication & Repo Management (Unchanged) ---99# ... (The rest of the file from the previous correct version remains the same) ...100def create_jwt(app_id, private_key):101    now = int(time.time())102    payload = {"iat": now, "exp": now + (10 * 60), "iss": app_id}103    return jwt.encode(payload, private_key, algorithm="RS256")104 105def get_installation_access_token(installation_id, app_id, private_key):106    app_jwt = create_jwt(app_id, private_key)107    headers = {"Authorization": f"Bearer {app_jwt}", "Accept": "application/vnd.github.v3+json"}108    url = f"https://api.github.com/app/installations/{installation_id}/access_tokens"109    response = requests.post(url, headers=headers)110    response.raise_for_status()111    return response.json()["token"]112 113def process_repository(repo_url, token, repo_full_name):114    """Clones a repo and hands it off to the CodeProcessor."""115    temp_dir = tempfile.mkdtemp()116    print(f"Cloning {repo_full_name} into {temp_dir}...")117    118    try:119        clone_url = repo_url.replace("https://", f"https://x-access-token:{token}@")120        Repo.clone_from(clone_url, temp_dir)121        122        faiss_index, metadata = code_processor.process_repo(temp_dir)123        124        if faiss_index:125            repo_data_store[repo_full_name] = {126                "status": "processed",127                "faiss_index": faiss_index,128                "metadata": metadata129            }130            print(f"Successfully processed and indexed {repo_full_name}")131        else:132            print(f"No processable data found for {repo_full_name}")133 134    except Exception as e:135        print(f"Failed to process repository {repo_full_name}: {e}")136    finally:137        print(f"Cleaning up temporary directory: {temp_dir}")138        shutil.rmtree(temp_dir)139 140# --- FastAPI App (Webhook handler is unchanged) ---141app = FastAPI()142 143async def verify_signature(request: Request):144    if not GITHUB_WEBHOOK_SECRET: raise HTTPException(status_code=500, detail="Webhook secret not configured.")145    signature_header = request.headers.get("X-Hub-Signature-256")146    if not signature_header: raise HTTPException(status_code=400, detail="X-Hub-Signature-256 header is missing.")147    body = await request.body()148    sha_name, signature = signature_header.split("=")149    mac = hmac.new(GITHUB_WEBHOOK_SECRET.encode("utf-8"), msg=body, digestmod=hashlib.sha256)150    if not hmac.compare_digest(mac.hexdigest(), signature): raise HTTPException(status_code=400, detail="Invalid signature.")151 152@app.get("/")153def read_root():154    return {"message": "Docu-Pilot server is alive!"}155 156@app.post("/api/github/webhook")157async def github_webhook(request: Request):158    await verify_signature(request)159    event_type = request.headers.get("X-GitHub-Event")160    payload = await request.json()161    print(f"Received event: {event_type} with action: {payload.get('action')}")162    installation_id = payload.get("installation", {}).get("id")163    if not installation_id: return {"status": "ok", "message": "Event does not pertain to an installation."}164    repos_to_process = []165    if event_type == "installation" and payload.get("action") == "created":166        repos_to_process = payload.get("repositories", [])167    elif event_type == "installation_repositories" and payload.get("action") == "added":168        repos_to_process = payload.get("repositories_added", [])169    if repos_to_process:170        try:171            token = get_installation_access_token(installation_id, GITHUB_APP_ID, GITHUB_PRIVATE_KEY)172            for repo in repos_to_process:173                repo_full_name = repo["full_name"]174                repo_url = f"https://github.com/{repo_full_name}"175                process_repository(repo_url, token, repo_full_name)176        except Exception as e:177            print(f"Error during repository processing: {e}")178    elif event_type == "push":179        repo_name = payload.get("repository", {}).get("full_name")180        print(f"Received a push event on repo {repo_name}")181    return {"status": "ok"}