Update app/main.py

app/main.py

@@ -1,104 +1,142 @@
-# main.py - Final Fixed Version
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import List, Optional
+import os
 import re
 import logging
+import traceback
+import time
+from typing import List, Optional, Dict
+from dotenv import load_dotenv
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import uvicorn
+
+# Load environment variables
+load_dotenv()
 
-from app.services.reviewer_service import AIReviewerService
-from app.predictor import classifier, guide_generator
+from app.predictor import classifier, guide_generator, reviewer
+# Note: AIReviewerService from the first version is typically 
+# the underlying service for the 'reviewer' object in the second.
 
 # 1. Setup Logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 2. Initialize FastAPI and Services
+# 2. Initialize FastAPI
 app = FastAPI(title="GitGud AI Service")
-reviewer_service = AIReviewerService()
+main = app # Alias for compatibility
 
-# 3. Data Models (Order matters: ReviewRequest needs FileRequest)
+# Global embedding cache
+# Structure: { "repo_name": { "file_path": [embedding_vector] } }
+REPO_CACHE: Dict[str, Dict[str, List[float]]] = {}
+
+# 3. Data Models
 class FileRequest(BaseModel):
     fileName: str
     content: Optional[str] = None
+    repoName: Optional[str] = None
 
-class ReviewRequest(BaseModel):
+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  # Path -> Embedding
+    repoName: Optional[str] = None
+
+class ChatRequest(BaseModel):
+    query: str
+    context: List[Dict[str, str]]  # List of { "fileName": str, "content": str }
+    repoName: str
+
 # 4. Endpoints
 
 @app.get("/")
 def health_check():
-    """Checks server status and GPU availability."""
+    """Checks server status, GPU availability, and cached data."""
     return {
         "status": "online",
         "model": "microsoft/codebert-base",
-        "device": classifier.device,
+        "device": getattr(classifier, "device", "cpu"),
+        "cached_repos": list(REPO_CACHE.keys()),
     }
 
+@app.get("/usage")
+def get_usage():
+    """Returns AI Service usage statistics."""
+    from app.core.model_loader import llm_engine
+    return llm_engine.get_usage_stats()
+
 @app.post("/classify")
 async def classify_file(request: FileRequest):
-    """Classifies file into architectural layers."""
+    """Classifies file into architectural layers and caches embeddings."""
     try:
         result = classifier.predict(request.fileName, request.content)
+        
+        # Cache embedding if repoName is provided
+        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, 
+            "fileName": request.fileName,
             "layer": result["label"],
             "confidence": result["confidence"],
             "embedding": result["embedding"]
         }
     except Exception as e:
         logger.error(f"Classify failed: {e}")
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/generate-guide")
-async def generate_guide(request: GuideRequest):
-    """Generates markdown guides for repositories."""
+@app.post("/review-batch-code")
+async def review_batch_code(request: BatchReviewRequest):
+    """Batch review with detailed metrics and suggestions."""
     try:
-        markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
-        return {"markdown": markdown}
-    except Exception as e:
-        logger.error(f"Guide generation failed: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        reviews = reviewer.service.review_batch_code(request.files)
+        total_files = len(reviews)
+        total_vulns = sum(len(r.get("vulnerabilities", [])) for r in reviews)
+        
+        # Calculate Average Maintainability
+        m_scores = [r.get("metrics", {}).get("maintainability", 8) for r in reviews]
+        avg_maint = sum(m_scores) / max(total_files, 1)
 
-@app.post("/review")
-async def review_code(request: ReviewRequest):
-    """Detects security and logic issues in batches of files."""
-    try:
-        # Call the batch review logic from your service
-        results = reviewer_service.review_batch_code(request.files)
-        return {"reviews": results}
+        return {
+            "totalFiles": total_files,
+            "totalVulnerabilities": total_vulns,
+            "averageMaintainability": round(avg_maint, 1),
+            "results": reviews,
+        }
     except Exception as e:
-        logger.error(f"Review endpoint failed: {e}")
+        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/repo-dashboard-stats")
-async def get_dashboard_stats(request: ReviewRequest):
+async def get_dashboard_stats(request: BatchReviewRequest):
+    """Aggregated stats for frontend dashboards including health and API sniffing."""
     try:
-        raw_reviews = reviewer_service.review_batch_code(request.files)
+        raw_reviews = reviewer.service.review_batch_code(request.files)
         
         # 1. Security Count
         total_vulns = sum(len(r.get("vulnerabilities", [])) for r in raw_reviews)
         
-        # 2. Performance Ratio (Maintainability)
-        # We use a default of 8 if the AI misses a file to avoid 0% scores
+        # 2. Performance/Maintainability Ratio
         scores = [r.get("metrics", {}).get("maintainability", 8) for r in raw_reviews]
         avg_maintainability = (sum(scores) / len(scores)) * 10 if scores else 0
         
-        # 3. API Sniffing
+        # 3. API Sniffing (Regex)
         found_apis = []
         for f in request.files:
             if f.content:
-                # Regex looks for common route decorators or methods
                 matches = re.findall(r'(?:get|post|put|delete|patch)\([\'"]\/(.*?)[\'"]', f.content.lower())
                 for match in matches:
                     found_apis.append(f"/{match}")
 
         # 4. Repo Health Calculation
-        # Every security issue drops health by 10 points
         health_score = max(10, 100 - (total_vulns * 10))
 
         return {
@@ -112,8 +150,96 @@ async def get_dashboard_stats(request: ReviewRequest):
         logger.error(f"Dashboard stats failed: {e}")
         raise HTTPException(status_code=500, detail="Failed to aggregate repository stats")
 
+@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)
+        tags = classifier.extract_tags(request.content, request.fileName)
+
+        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"],
+            "summary": summary,
+            "tags": tags,
+            "embedding": result["embedding"],
+        }
+    except Exception as e:
+        traceback.print_exc()
+        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:
+            embeddings = REPO_CACHE[request.repoName]
+
+        if not embeddings:
+            return {"results": []}
+
+        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()
+    logger.info(f"Received Chat Request for {request.repoName}")
+    
+    try:
+        from app.core.model_loader import llm_engine
+        
+        context_str = ""
+        for item in request.context:
+            context_str += f"--- FILE: {item['fileName']} ---\n{item['content']}\n\n"
+        
+        has_context = len(request.context) > 0
+        prompt = f"""
+        You are "GitGud AI", an expert software architect.
+        Repository: "{request.repoName}"
+
+        INSTRUCTIONS:
+        1. Use the provided CONTEXT to answer. 
+        2. If context is missing, state: "I am using general knowledge as I don't have specific snippets for this."
+        3. Use markdown for code.
+
+        CONTEXT:
+        {context_str if has_context else "[(NO CODE SNIPPETS PROVIDED)]"}
+
+        USER QUESTION:
+        {request.query}
+        """
+        response = llm_engine.generate_text(prompt)
+        
+        logger.info(f"Chat response generated in {time.time() - start_time:.2f}s")
+        return {"response": response}
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/generate-guide")
+async def generate_guide(request: GuideRequest):
+    """Generates markdown documentation for the repo."""
+    try:
+        markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
+        return {"markdown": markdown}
+    except Exception as e:
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
 # 5. Application Entry Point
 if __name__ == "__main__":
-    import uvicorn
-    # Port 7860 is mandatory for Hugging Face Spaces
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # Note: Using 7860 for HF Spaces compatibility, change to 8000 if preferred for local dev
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)