Update app/main.py

app/main.py

@@ -8,7 +8,7 @@ from typing import List, Optional, Dict
 from concurrent.futures import ThreadPoolExecutor
 from dotenv import load_dotenv
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, status
 from pydantic import BaseModel
 import uvicorn
 
@@ -16,7 +16,6 @@ import uvicorn
 load_dotenv()
 
 # Import internal services
-# Ensure these modules exist in your /app directory
 from app.predictor import classifier, guide_generator, reviewer
 from app.core.model_loader import llm_engine
 
@@ -28,7 +27,6 @@ logger = logging.getLogger(__name__)
 app = FastAPI(title="GitGud AI Service")
 
 # Global embedding cache
-# Structure: { "repo_name": { "file_path": [embedding_vector] } }
 REPO_CACHE: Dict[str, Dict[str, List[float]]] = {}
 
 # 3. Data Models
@@ -56,23 +54,31 @@ class ChatRequest(BaseModel):
 
 # 4. Core Logic Helpers
 def calculate_repo_health(total_vulns: int, avg_maint: float) -> int:
-    """
-    Standardized health logic to keep Dashboard and Review metrics in sync.
-    Uses a weighted scale of maintainability vs security risk.
-    """
-    # Base starts from the maintainability average (0.0 to 10.0 scale mapped to 100)
     base_score = avg_maint * 10
-    # Penalty: Subtract 8 points per vulnerability
     penalty = total_vulns * 8
-    
     final_score = base_score - penalty
     return int(max(10, min(100, final_score)))
 
+def parse_tree_to_list(raw_tree: str):
+    """Parses text tree into JSON objects for Compose LazyColumn"""
+    nodes = []
+    for line in raw_tree.strip().split('\n'):
+        # Detect depth based on tree characters
+        level = line.count('|') + (line.count(' ') // 2)
+        # Clean the name
+        name = re.sub(r'[|└├─]', '', line).strip()
+        if name:
+            nodes.append({
+                "name": name,
+                "type": "file" if '.' in name else "folder",
+                "level": level
+            })
+    return nodes
+
 # 5. Endpoints
 
 @app.get("/")
 def health_check():
-    """Checks server status, GPU availability, and cached data."""
     return {
         "status": "online",
         "model": "microsoft/codebert-base",
@@ -82,15 +88,12 @@ def health_check():
 
 @app.get("/usage")
 def get_usage():
-    """Returns AI Service usage statistics from the LLM engine."""
     return llm_engine.get_usage_stats()
 
 @app.post("/classify")
 async def classify_file(request: FileRequest):
-    """Classifies file into architectural layers and caches embeddings."""
     try:
         result = classifier.predict(request.fileName, request.content)
-        
         if request.repoName:
             if request.repoName not in REPO_CACHE:
                 REPO_CACHE[request.repoName] = {}
@@ -108,7 +111,6 @@ async def classify_file(request: FileRequest):
 
 @app.post("/review-batch-code")
 async def review_batch_code(request: BatchReviewRequest):
-    """Detailed review results for the Review Section UI."""
     try:
         loop = asyncio.get_event_loop()
         with ThreadPoolExecutor() as executor:
@@ -117,14 +119,14 @@ async def review_batch_code(request: BatchReviewRequest):
             )
         return {"results": reviews}
     except Exception as e:
+        if "429" in str(e):
+            raise HTTPException(status_code=429, detail="AI Quota Exceeded")
         traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/repo-dashboard-stats")
 async def get_dashboard_stats(request: BatchReviewRequest):
-    """Aggregated stats using the exact same logic as batch review."""
     try:
-        # Run heavy AI review in thread pool to keep FastAPI responsive
         loop = asyncio.get_event_loop()
         with ThreadPoolExecutor() as executor:
             raw_reviews = await loop.run_in_executor(
@@ -137,15 +139,11 @@ async def get_dashboard_stats(request: BatchReviewRequest):
         api_regex = re.compile(r'(?:get|post|put|delete|patch)\([\'"]\/(.*?)[\'"]', re.IGNORECASE)
 
         for i, review in enumerate(raw_reviews):
-            # Sync Vulnerability count
             vulns = review.get("vulnerabilities", [])
             total_vulns += len(vulns)
-            
-            # Sync Maintainability
             m_score = review.get("metrics", {}).get("maintainability", 8.0)
             maint_scores.append(m_score)
             
-            # Extract APIs (No cap, show all discovered)
             content = request.files[i].content
             if content:
                 matches = api_regex.findall(content)
@@ -154,8 +152,6 @@ async def get_dashboard_stats(request: BatchReviewRequest):
 
         num_files = len(maint_scores)
         avg_maint = (sum(maint_scores) / num_files) if num_files > 0 else 0
-        
-        # Calculate health using shared logic
         health_score = calculate_repo_health(total_vulns, avg_maint)
 
         return {
@@ -168,13 +164,13 @@ async def get_dashboard_stats(request: BatchReviewRequest):
             "average_maintainability": round(avg_maint, 1)
         }
     except Exception as e:
+        if "429" in str(e):
+            raise HTTPException(status_code=429, detail="Quota exceeded")
         logger.error(f"Dashboard stats failed: {e}")
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail="Failed to sync dashboard metrics")
 
 @app.post("/analyze-file")
 async def analyze_file(request: FileRequest):
-    """Deep analysis: Summary, Tags, and Layer Classification."""
     try:
         result = classifier.predict(request.fileName, request.content)
         summary = classifier.generate_file_summary(request.content, request.fileName)
@@ -188,12 +184,12 @@ async def analyze_file(request: FileRequest):
             "embedding": result["embedding"],
         }
     except Exception as e:
-        traceback.print_exc()
+        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):
-    """Search code using natural language and vector similarity."""
     try:
         embeddings = request.embeddings
         if not embeddings and request.repoName and request.repoName in REPO_CACHE:
@@ -205,12 +201,10 @@ async def semantic_search(request: SearchRequest):
         results = classifier.semantic_search(request.query, embeddings)
         return {"results": results}
     except Exception as e:
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    """RAG-based chat using provided file context."""
     start_time = time.time()
     try:
         context_str = ""
@@ -224,21 +218,35 @@ async def chat(request: ChatRequest):
         USER QUESTION: {request.query}
         """
         response = llm_engine.generate_text(prompt)
-        
-        logger.info(f"Chat generated in {time.time() - start_time:.2f}s")
-        return {"response": response}
+        return {"response": response, "status": "success"}
     except Exception as e:
-        traceback.print_exc()
+        if "429" in str(e):
+            # Return a structured error so Compose can show a nice UI
+            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):
-    """Generates markdown documentation for the repo."""
+    """
+    FIXED: Now parses the messy tree text into JSON for your Compose UI.
+    """
     try:
         markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
-        return {"markdown": markdown}
+        
+        # Extract the messy tree part and clean it
+        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:
-        traceback.print_exc()
+        if "429" in str(e):
+            raise HTTPException(status_code=429, detail="AI Quota Exceeded")
         raise HTTPException(status_code=500, detail=str(e))
 
 # 6. Application Entry Point