Update app.py

app.py

@@ -6,9 +6,13 @@ import jwt
 import requests
 import tempfile
 import shutil
+import ast
+import faiss
+import numpy as np
 from fastapi import FastAPI, Request, HTTPException, status
 from dotenv import load_dotenv
 from git import Repo
+from sentence_transformers import SentenceTransformer
 
 # --- Configuration ---
 load_dotenv()
@@ -16,65 +20,133 @@ GITHUB_WEBHOOK_SECRET = os.getenv("GITHUB_WEBHOOK_SECRET")
 GITHUB_APP_ID = os.getenv("GITHUB_APP_ID")
 GITHUB_PRIVATE_KEY = os.getenv("GITHUB_PRIVATE_KEY")
 
-# In-memory storage for our repository data (for this simple version)
+# In-memory storage for our repository data
 repo_data_store = {}
 
-# --- GitHub App Authentication Logic (Unchanged) ---
+# --- NEW: Code Processing and Vectorization Class ---
+
+class CodeProcessor:
+    def __init__(self, model_name='all-MiniLM-L6-v2'):
+        print("Initializing SentenceTransformer model...")
+        self.model = SentenceTransformer(model_name)
+        print("Model initialized.")
+
+    def parse_python_file(self, file_path):
+        """Parses a Python file to extract functions and their docstrings using AST."""
+        with open(file_path, 'r', encoding='utf-8') as source:
+            try:
+                tree = ast.parse(source.read())
+            except SyntaxError:
+                return [] # Ignore files with syntax errors
+        
+        functions = []
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef):
+                # We extract the function's name and its full source code
+                # The source code provides more context than just the docstring
+                function_source = ast.get_source_segment(tree, node)
+                functions.append({
+                    "name": node.name,
+                    "source": function_source,
+                    "file_path": file_path
+                })
+        return functions
+
+    def process_repo(self, repo_path):
+        """Walks a repository, parses Python files, and creates a FAISS index."""
+        all_functions = []
+        print(f"Walking through Python files in {repo_path}...")
+        for root, _, files in os.walk(repo_path):
+            for file in files:
+                if file.endswith('.py'):
+                    file_path = os.path.join(root, file)
+                    functions = self.parse_python_file(file_path)
+                    all_functions.extend(functions)
+        
+        if not all_functions:
+            print("No Python functions found in the repository.")
+            return None, None
+
+        print(f"Found {len(all_functions)} functions. Generating embeddings...")
+        # We create embeddings for the source code of each function
+        function_sources = [f["source"] for f in all_functions]
+        embeddings = self.model.encode(function_sources, show_progress_bar=False)
+        
+        # Create a FAISS index
+        embedding_dim = embeddings.shape[1]
+        index = faiss.IndexFlatL2(embedding_dim)
+        index.add(np.array(embeddings, dtype=np.float32))
+        
+        print(f"FAISS index created successfully with {index.ntotal} vectors.")
+        
+        # We return the index and the metadata (the list of function details)
+        return index, all_functions
+
+# --- Initialize the processor globally ---
+# This ensures the model is loaded only once when the app starts.
+code_processor = CodeProcessor()
+
+
+# --- GitHub App Authentication & Repo Management (Largely Unchanged) ---
 
 def create_jwt(app_id, private_key):
+    # (This function is unchanged)
     now = int(time.time())
-    payload = {
-        "iat": now,
-        "exp": now + (10 * 60),
-        "iss": app_id,
-    }
+    payload = {"iat": now, "exp": now + (10 * 60), "iss": app_id}
     return jwt.encode(payload, private_key, algorithm="RS256")
 
 def get_installation_access_token(installation_id, app_id, private_key):
+    # (This function is unchanged)
     app_jwt = create_jwt(app_id, private_key)
-    headers = {
-        "Authorization": f"Bearer {app_jwt}",
-        "Accept": "application/vnd.github.v3+json",
-    }
+    headers = {"Authorization": f"Bearer {app_jwt}", "Accept": "application/vnd.github.v3+json"}
     url = f"https://api.github.com/app/installations/{installation_id}/access_tokens"
     response = requests.post(url, headers=headers)
     response.raise_for_status()
     return response.json()["token"]
 
-# --- Repository Management Logic (Unchanged) ---
-
 def process_repository(repo_url, token, repo_full_name):
+    """Clones a repo and hands it off to the CodeProcessor."""
     temp_dir = tempfile.mkdtemp()
-    print(f"Cloning {repo_full_name} into {temp_dir}")
+    print(f"Cloning {repo_full_name} into {temp_dir}...")
     
     try:
         clone_url = repo_url.replace("https://", f"https://x-access-token:{token}@")
         Repo.clone_from(clone_url, temp_dir)
         
-        print(f"Parsing repository {repo_full_name}...")
-        # In the next step, we will add AST parsing and vectorization here.
+        # --- NEW: Use the CodeProcessor ---
+        faiss_index, metadata = code_processor.process_repo(temp_dir)
         
-        repo_data_store[repo_full_name] = {"status": "processed", "path": temp_dir}
-        print(f"Successfully processed and 'stored' data for {repo_full_name}")
+        if faiss_index:
+            # Store the live FAISS index and metadata in our in-memory store
+            repo_data_store[repo_full_name] = {
+                "status": "processed",
+                "faiss_index": faiss_index,
+                "metadata": metadata
+            }
+            print(f"Successfully processed and indexed {repo_full_name}")
+        else:
+            print(f"No processable data found for {repo_full_name}")
 
     except Exception as e:
         print(f"Failed to process repository {repo_full_name}: {e}")
+    finally:
+        # Clean up the cloned repository from disk
+        print(f"Cleaning up temporary directory: {temp_dir}")
+        shutil.rmtree(temp_dir)
+
 
-# --- FastAPI App ---
+# --- FastAPI App (Webhook handler is unchanged) ---
 app = FastAPI()
 
 async def verify_signature(request: Request):
-    # (This function remains unchanged)
-    if not GITHUB_WEBHOOK_SECRET:
-        raise HTTPException(status_code=500, detail="Webhook secret not configured.")
+    # (This function is unchanged)
+    if not GITHUB_WEBHOOK_SECRET: raise HTTPException(status_code=500, detail="Webhook secret not configured.")
     signature_header = request.headers.get("X-Hub-Signature-256")
-    if not signature_header:
-        raise HTTPException(status_code=400, detail="X-Hub-Signature-256 header is missing.")
+    if not signature_header: raise HTTPException(status_code=400, detail="X-Hub-Signature-256 header is missing.")
     body = await request.body()
     sha_name, signature = signature_header.split("=")
     mac = hmac.new(GITHUB_WEBHOOK_SECRET.encode("utf-8"), msg=body, digestmod=hashlib.sha256)
-    if not hmac.compare_digest(mac.hexdigest(), signature):
-        raise HTTPException(status_code=400, detail="Invalid signature.")
+    if not hmac.compare_digest(mac.hexdigest(), signature): raise HTTPException(status_code=400, detail="Invalid signature.")
 
 @app.get("/")
 def read_root():
@@ -82,46 +154,28 @@ def read_root():
 
 @app.post("/api/github/webhook")
 async def github_webhook(request: Request):
+    # (This webhook handler is unchanged from the last step)
     await verify_signature(request)
-    
     event_type = request.headers.get("X-GitHub-Event")
     payload = await request.json()
-
     print(f"Received event: {event_type} with action: {payload.get('action')}")
-
     installation_id = payload.get("installation", {}).get("id")
-    if not installation_id:
-        return {"status": "ok", "message": "Event does not pertain to an installation."}
-
-    # --- REVISED LOGIC ---
+    if not installation_id: return {"status": "ok", "message": "Event does not pertain to an installation."}
     repos_to_process = []
-    
     if event_type == "installation" and payload.get("action") == "created":
-        # This handles the very first time the app is installed on an account.
         repos_to_process = payload.get("repositories", [])
-        print(f"Processing 'installation.created' event for repos: {[repo['full_name'] for repo in repos_to_process]}")
-        
     elif event_type == "installation_repositories" and payload.get("action") == "added":
-        # This handles when a user adds one or more repos to an existing installation.
         repos_to_process = payload.get("repositories_added", [])
-        print(f"Processing 'installation_repositories.added' event for repos: {[repo['full_name'] for repo in repos_to_process]}")
-
     if repos_to_process:
         try:
             token = get_installation_access_token(installation_id, GITHUB_APP_ID, GITHUB_PRIVATE_KEY)
             for repo in repos_to_process:
                 repo_full_name = repo["full_name"]
-                # The 'html_url' key is not present in the 'installation_repositories' event,
-                # so we construct it manually.
                 repo_url = f"https://github.com/{repo_full_name}"
                 process_repository(repo_url, token, repo_full_name)
         except Exception as e:
             print(f"Error during repository processing: {e}")
-            # Don't raise HTTPException here to prevent GitHub from disabling the webhook
-    
-    # Handle push events (will be used in Step 4)
     elif event_type == "push":
         repo_name = payload.get("repository", {}).get("full_name")
         print(f"Received a push event on repo {repo_name}")
-
     return {"status": "ok"}