Update app/main.py

app/main.py

@@ -1,212 +1,233 @@
-import json
-import logging
+import os
 import re
-import hashlib
-import sqlite3
-import threading
-from typing import List
+import logging
+import traceback
+import time
+import asyncio
+from typing import List, Optional, Dict
+from concurrent.futures import ThreadPoolExecutor
+from dotenv import load_dotenv
+from fastapi import FastAPI, HTTPException, status
+from pydantic import BaseModel
+import uvicorn
+
+load_dotenv()
+
+from app.predictor import classifier, guide_generator, reviewer
 from app.core.model_loader import llm_engine
 
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-DB_PATH = "review_cache.db"
-
-class AIReviewerService:
-    def __init__(self):
-        self.batch_size = 3  # Reduced from 10 to speed up individual prompts
-        self.db_lock = threading.Lock()
-        self._init_db()
-
-    def _init_db(self):
-        """Initializes SQLite cache table with proper timeout."""
-        with self.db_lock:
-            with sqlite3.connect(DB_PATH, timeout=30.0) as conn:
-                cursor = conn.cursor()
-                cursor.execute("""
-                    CREATE TABLE IF NOT EXISTS review_cache (
-                        file_hash TEXT PRIMARY KEY,
-                        result_json TEXT
-                    )
-                """)
-                conn.commit()
-
-    def _generate_file_hash(self, file_obj) -> str:
-        content = file_obj.content or ""
-        raw_payload = f"{file_obj.fileName}:{content}"
-        return hashlib.sha256(raw_payload.encode("utf-8")).hexdigest()
-
-    def _get_cached_result(self, file_hash: str) -> dict:
-        with self.db_lock:
-            with sqlite3.connect(DB_PATH, timeout=30.0) as conn:
-                cursor = conn.cursor()
-                cursor.execute("SELECT result_json FROM review_cache WHERE file_hash = ?", (file_hash,))
-                row = cursor.fetchone()
-                if row:
-                    try:
-                        return json.loads(row[0])
-                    except json.JSONDecodeError:
-                        return None
-        return None
-
-    def _save_cached_results(self, hash_result_pairs: List[tuple]):
-        if not hash_result_pairs:
-            return
-        with self.db_lock:
-            with sqlite3.connect(DB_PATH, timeout=30.0) as conn:
-                cursor = conn.cursor()
-                cursor.executemany(
-                    "INSERT OR REPLACE INTO review_cache (file_hash, result_json) VALUES (?, ?)",
-                    [(h, json.dumps(res)) for h, res in hash_result_pairs]
-                )
-                conn.commit()
-
-    def review_batch_code(self, files: list) -> list:
-        if not files:
-            return []
-
-        results = [None] * len(files)
-        uncached_files = []
-        uncached_indices = []
-
-        # 1. Check SQLite cache
-        for idx, f in enumerate(files):
-            file_hash = self._generate_file_hash(f)
-            cached_res = self._get_cached_result(file_hash)
-            if cached_res:
-                results[idx] = cached_res
-            else:
-                uncached_files.append(f)
-                uncached_indices.append(idx)
-
-        if not uncached_files:
-            logger.info(f"⚡ Served all {len(files)} files instantly from SQLite cache.")
-            return results
-
-        # 2. Sequential batch processing (prevents LLM thread lock contention)
-        batches = [
-            uncached_files[i : i + self.batch_size]
-            for i in range(0, len(uncached_files), self.batch_size)
-        ]
-        
-        new_results = []
-        for batch in batches:
-            try:
-                batch_results = self._process_batch(batch)
-                new_results.extend(batch_results)
-            except Exception as e:
-                logger.error(f"Batch processing error: {e}")
-                for f_item in batch:
-                    new_results.append({
-                        "fileName": f_item.fileName,
-                        "vulnerabilities": [],
-                        "improvement_suggestions": [],
-                        "metrics": {"complexity": 0, "maintainability": 0},
-                        "error": "analysis_failed",
-                    })
-
-        new_results_map = {item.get("fileName"): item for item in new_results}
-        to_cache = []
-
-        # 3. Assemble final response and update cache
-        for idx, f in zip(uncached_indices, uncached_files):
-            res = new_results_map.get(
-                f.fileName,
-                {
-                    "fileName": f.fileName,
-                    "vulnerabilities": [],
-                    "improvement_suggestions": [],
-                    "metrics": {"complexity": 0, "maintainability": 0},
-                    "error": "analysis_failed",
-                },
-            )
-            if "error" not in res:
-                file_hash = self._generate_file_hash(f)
-                to_cache.append((file_hash, res))
-            results[idx] = res
-
-        self._save_cached_results(to_cache)
-        return results
 
-    def _process_batch(self, batch: list) -> list:
-        combined_code = ""
-        file_names = []
-        for f in batch:
-            raw_content = f.content or ""
-            minified_content = self._minify_code(raw_content[:2000], f.fileName)  # Reduced preview limit to 2000 chars
-            combined_code += f"\n--- FILE: {f.fileName} ---\n{minified_content}\n"
-            file_names.append(f.fileName)
-
-        prompt = f"""Analyze {len(batch)} files. Be ultra-concise.
-{combined_code}
-Return ONLY valid JSON array:
-[
-  {{
-    "fileName": "name",
-    "vulnerabilities": [{{"type": "type", "line": 1, "description": "short", "suggestion": "fix"}}],
-    "improvement_suggestions": ["tip"],
-    "metrics": {{"complexity": 3, "maintainability": 8}}
-  }}
-]"""
-        batch_results = []
-        try:
-            # Capped token length to 1024 max per batch turn
-            response_text = llm_engine.generate(prompt, max_tokens=1024)
-            parsed = self._parse_json(response_text)
-            processed_map = {item.get("fileName"): item for item in parsed if isinstance(item, dict)}
-
-            for fn in file_names:
-                if fn in processed_map:
-                    res = processed_map[fn]
-                    res.setdefault("vulnerabilities", [])
-                    res.setdefault("improvement_suggestions", [])
-                    res.setdefault("metrics", {"complexity": 5, "maintainability": 5})
-                    batch_results.append(res)
-                else:
-                    batch_results.append({
-                        "fileName": fn,
-                        "vulnerabilities": [],
-                        "improvement_suggestions": [],
-                        "metrics": {"complexity": 1, "maintainability": 10},
-                    })
-        except Exception as e:
-            logger.error(f"Batch processing error: {e}")
-            for fn in file_names:
-                batch_results.append({
-                    "fileName": fn,
-                    "vulnerabilities": [],
-                    "improvement_suggestions": [],
-                    "metrics": {"complexity": 0, "maintainability": 0},
-                    "error": str(e),
-                })
-        return batch_results
-
-    def _minify_code(self, code: str, filename: str) -> str:
-        lines = code.split("\n")
-        cleaned_lines = []
-        is_python = filename.endswith(".py")
-        is_js_style = filename.endswith((".js", ".ts", ".jsx", ".tsx"))
-
-        for line in lines:
-            stripped = line.strip()
-            if not stripped:
-                continue
-            if is_python and stripped.startswith("#"):
-                continue
-            if is_js_style and stripped.startswith(("//", "/*", "*")):
-                continue
-            cleaned_lines.append(stripped)
-
-        return "\n".join(cleaned_lines[:50])  # Cap at 50 lines max per file
-
-    def _parse_json(self, text: str):
-        text = text.strip()
-        if not text:
-            return []
-        text = re.sub(r"^[^\[]*\[", "[", text)
-        text = re.sub(r"\][^\]]*$", "]", text)
-        text = text.replace("```json", "").replace("```", "").strip()
-        try:
-            return json.loads(text)
-        except json.JSONDecodeError:
-            return []
-
-service = AIReviewerService()
+app = FastAPI(title="GitGud AI Service")
+
+REPO_CACHE: Dict[str, Dict[str, List[float]]] = {}
+executor = ThreadPoolExecutor(max_workers=8)
+
+class FileRequest(BaseModel):
+    fileName: str
+    content: Optional[str] = None
+    repoName: Optional[str] = None
+
+class BatchReviewRequest(BaseModel):
+    files: List[FileRequest]
+
+class GuideRequest(BaseModel):
+    repoName: str
+    filePaths: List[str]
+
+class SearchRequest(BaseModel):
+    query: str
+    embeddings: Optional[Dict[str, List[float]]] = None
+    repoName: Optional[str] = None
+
+class ChatRequest(BaseModel):
+    query: str
+    context: List[Dict[str, str]]
+    repoName: str
+
+def calculate_repo_health(total_vulns: int, avg_maint: float) -> int:
+    base_score = avg_maint * 10
+    penalty = total_vulns * 8
+    return int(max(10, min(100, base_score - penalty)))
+
+def sync_review_worker(file_list: List[FileRequest]):
+    logger.info(f"--- [DEBUG] Processing {len(file_list)} files for code review ---")
+    try:
+        results = reviewer.service.review_batch_code(file_list)
+        return results
+    except Exception as e:
+        logger.error(f"--- [ERROR] Exception during review processing: {e} ---", exc_info=True)
+        raise e
+
+def parse_tree_to_list(raw_tree: str):
+    nodes = []
+    for line in raw_tree.strip().split('\n'):
+        level = line.count('|') + (line.count(' ') // 2)
+        name = re.sub(r'[|└├─]', '', line).strip()
+        if name:
+            nodes.append({
+                "name": name,
+                "type": "file" if '.' in name else "folder",
+                "level": level
+            })
+    return nodes
+
+@app.get("/")
+def health_check():
+    return {
+        "status": "online",
+        "model": "microsoft/codebert-base",
+        "device": getattr(classifier, "device", "cpu"),
+        "cached_repos": list(REPO_CACHE.keys()),
+    }
+
+@app.get("/usage")
+def get_usage():
+    return llm_engine.get_usage_stats()
+
+@app.post("/classify")
+async def classify_file(request: FileRequest):
+    try:
+        result = classifier.predict(request.fileName, request.content)
+        if request.repoName:
+            if request.repoName not in REPO_CACHE:
+                REPO_CACHE[request.repoName] = {}
+            REPO_CACHE[request.repoName][request.fileName] = result["embedding"]
+        return {
+            "fileName": request.fileName,
+            "layer": result["label"],
+            "confidence": result["confidence"],
+            "embedding": result["embedding"]
+        }
+    except Exception as e:
+        logger.error(f"Classify failed: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/review-batch-code")
+async def review_batch_code(request: BatchReviewRequest):
+    try:
+        loop = asyncio.get_running_loop()
+        raw_results = await loop.run_in_executor(executor, sync_review_worker, request.files)
+        final_reviews = raw_results if isinstance(raw_results, list) else [raw_results]
+        return {"results": final_reviews}
+    except Exception as e:
+        logger.error(f"Batch review critical failure: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail="Internal processing error")
+
+@app.post("/repo-dashboard-stats")
+async def get_dashboard_stats(request: BatchReviewRequest):
+    try:
+        loop = asyncio.get_running_loop()
+        raw_reviews = await loop.run_in_executor(executor, sync_review_worker, request.files)
+        if not isinstance(raw_reviews, list):
+            raw_reviews = [raw_reviews]
+
+        total_vulns = 0
+        maint_scores = []
+        found_apis = set()
+        api_regex = re.compile(r'(?:get|post|put|delete|patch)\([\'"]\/(.*?)[\'"]', re.IGNORECASE)
+
+        for i, current_review in enumerate(raw_reviews):
+            vulns = current_review.get("vulnerabilities", [])
+            total_vulns += len(vulns)
+            m_score = current_review.get("metrics", {}).get("maintainability", 8.0)
+            maint_scores.append(m_score)
+
+            content = request.files[i].content if i < len(request.files) else None
+            if content:
+                matches = api_regex.findall(content)
+                for match in matches:
+                    found_apis.add(f"/{match}")
+
+        num_files = len(maint_scores)
+        avg_maint = (sum(maint_scores) / num_files) if num_files > 0 else 0
+        health_score = calculate_repo_health(total_vulns, avg_maint)
+
+        return {
+            "repo_health": health_score,
+            "health_label": "Excellent" if health_score > 85 else "Good" if health_score > 60 else "Critical",
+            "security_issues": total_vulns,
+            "performance_ratio": f"{int(avg_maint * 10)}%",
+            "exposed_apis": list(found_apis),
+            "total_files_processed": num_files,
+            "average_maintainability": round(avg_maint, 1)
+        }
+    except Exception as e:
+        logger.error(f"Stats failed: {e}")
+        raise HTTPException(status_code=500, detail="Failed to aggregate metrics")
+
+@app.post("/analyze-file")
+async def analyze_file(request: FileRequest):
+    try:
+        result = classifier.predict(request.fileName, request.content)
+        summary = classifier.generate_file_summary(request.content, request.fileName)
+        tags = classifier.extract_tags(request.content, request.fileName)
+        return {
+            "fileName": request.fileName,
+            "layer": result["label"],
+            "summary": summary,
+            "tags": tags,
+            "embedding": result["embedding"],
+        }
+    except Exception as e:
+        if "429" in str(e):
+            raise HTTPException(status_code=429, detail="Limit Reached")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/semantic-search")
+async def semantic_search(request: SearchRequest):
+    try:
+        embeddings = request.embeddings
+        if not embeddings and request.repoName and request.repoName in REPO_CACHE:
+            embeddings = REPO_CACHE[request.repoName]
+
+        if not embeddings:
+            return {"results": []}
+
+        results = classifier.semantic_search(request.query, embeddings)
+        return {"results": results}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    try:
+        context_str = ""
+        for item in request.context:
+            context_str += f"--- FILE: {item['fileName']} ---\n{item['content']}\n\n"
+
+        prompt = f"""You are "GitGud AI", an expert software architect.
+Repository: "{request.repoName}"
+CONTEXT: {context_str if request.context else "(NO CODE PROVIDED)"}
+USER QUESTION: {request.query}"""
+
+        response = llm_engine.generate_text(prompt)
+        return {"response": response, "status": "success"}
+    except Exception as e:
+        if "429" in str(e):
+            return {"response": "⚠️ Daily limit reached. Try again in a bit!", "status": "quota_error"}
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate-guide")
+async def generate_guide(request: GuideRequest):
+    try:
+        markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
+        tree_match = re.search(r"Project Structure\n\n(.*?)(?=\n\n|$)", markdown, re.S)
+        structured_tree = []
+        if tree_match:
+            structured_tree = parse_tree_to_list(tree_match.group(1))
+
+        return {
+            "markdown": markdown,
+            "structured_tree": structured_tree,
+            "project_name": request.repoName
+        }
+    except Exception as e:
+        if "429" in str(e):
+            raise HTTPException(status_code=429, detail="AI Quota Exceeded")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)