back_end/main.py

# main.py
from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import asyncio
import os
from pathlib import Path
from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, AIMessage

# Import your existing backend logic
from core.downloader import download_github_repo, delete_dir
from core.loader import count_valid_supported_files, load_repository_as_documents
from core.splitter import custom_splitter
from core.embeddings import build_vector_db
from agent.graph import build_workflow
from config import MAX_FILES_TO_CREATE_VECTOR_DB
import json

load_dotenv()

app = FastAPI()

# 1. CORS Middleware (Crucial for frontend connection)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 2. Global State Variables
GLOBAL_APP = None
GLOBAL_CHAT_HISTORY = []

# 3. Request Models
class RepoRequest(BaseModel):
    url: str

class ChatRequest(BaseModel):
    message: str


# --- ENDPOINTS ---
@app.post("/init-repo")
async def initialize_repository(request: RepoRequest):
    """Downloads repo, builds vector DB, and streams status updates as JSON."""
    
    async def process_stream():
        global GLOBAL_APP, GLOBAL_CHAT_HISTORY
        
        current_dir = Path(os.getcwd())
        temp_git_repo_storage = current_dir / "temp_git_repo_storage"

        # Helper function that JUST formats the string
        def format_msg(status: str, task_name: str):
            payload = {
                "status": status.upper(), # Forces uppercase to prevent JS errors
                "task": task_name,
            }
            return f"data: {json.dumps(payload)}\n\n"

        try:
            yield format_msg("START", "Preparing to load repository")
            await asyncio.sleep(0.01) # Flush to UI

            delete_dir(temp_git_repo_storage)
            temp_git_repo_storage.mkdir(parents=True, exist_ok=True)

            yield format_msg("SUCCESS", "Ready to load repository")
            
            yield format_msg("START", f"Loading repository from {request.url}...")
            await asyncio.sleep(0.01)
            
            extracted_path = await asyncio.to_thread(download_github_repo, request.url, temp_git_repo_storage)
            
            yield format_msg("SUCCESS", "Repo successfully loaded.")
            await asyncio.sleep(0.01)

        except Exception as e:
            yield format_msg("ERROR", f"Failed to download repository: {e}")
            return

        yield format_msg("START", "Counting supported files...")
        await asyncio.sleep(0.01)
        
        file_count = await asyncio.to_thread(count_valid_supported_files, temp_git_repo_storage)
        
        if file_count > MAX_FILES_TO_CREATE_VECTOR_DB:
            yield format_msg("WARNING", f"Repo is large ({file_count} files). Building workflow without Vector DB...")
            await asyncio.sleep(0.01)
            GLOBAL_APP = await asyncio.to_thread(build_workflow, temp_git_repo_storage, False)
        else:
            yield format_msg("SUCCESS", f"Found {file_count} files to process")
            
            yield format_msg("START", "Loading repository files as documents...")
            await asyncio.sleep(0.01)
            doc = await asyncio.to_thread(load_repository_as_documents, temp_git_repo_storage)
            yield format_msg("SUCCESS", "Files loaded")

            yield format_msg("START", "Preparing files to analyse")
            await asyncio.sleep(0.01)
            all_splits = await asyncio.to_thread(custom_splitter, doc, current_dir)
            yield format_msg("SUCCESS", "Done preparing")

            yield format_msg("START", "Analysing files (This may take 1 to 10 minutes)...")
            await asyncio.sleep(0.01)
            vector_db = await asyncio.to_thread(build_vector_db, all_splits)
            yield format_msg("SUCCESS", "Done Analysing")

            yield format_msg("START", "Loading model")
            await asyncio.sleep(0.01)
            GLOBAL_APP = await asyncio.to_thread(build_workflow, temp_git_repo_storage, True, all_splits, vector_db)
            yield format_msg("SUCCESS", "Model Loaded")

        GLOBAL_CHAT_HISTORY = []
        
        # Final success message
        yield format_msg("FINISHED", "System ready. Switching to chat.")
        await asyncio.sleep(0.01)

    return StreamingResponse(process_stream(), media_type="text/event-stream")


import json
import asyncio
from fastapi.responses import StreamingResponse

@app.post("/chat")
async def chat_stream(request: ChatRequest):
    """Streams the LangGraph response back to the frontend."""
    
    if not GLOBAL_APP:
        return {"error": "System not initialized. Please load a repo first."}

    async def generate_response():
        global GLOBAL_CHAT_HISTORY
        
        user_input = request.message
        GLOBAL_CHAT_HISTORY.append(HumanMessage(content=user_input))
        
        config = {"recursion_limit": 100}
        final_ai_message = None
        
        # Helper function for JSON SSE (Just like init-repo!)
        def format_chat_chunk(msg_type: str, text: str):
            payload = {
                "type": msg_type, # "thinking", "tool", "message", or "end"
                "text": text
            }
            return f"data: {json.dumps(payload)}\n\n"

        # Stream the response from LangGraph
        for event in GLOBAL_APP.stream({"messages": GLOBAL_CHAT_HISTORY}, stream_mode="values", config=config):
            message = event["messages"][-1]
            message.pretty_print()
            
            # 1. Check if the AI is calling a tool
            if message.type == "ai" and getattr(message, "tool_calls", None):
                for tool in message.tool_calls:
                    yield format_chat_chunk("tool", f"Browsing Codebase..")
                    await asyncio.sleep(0.01)

            # 2. Process the actual AI content
            if message.type == "ai":
                
                # SCENARIO A: Model is Thinking (Content is a list of dicts)
                if isinstance(message.content, list):
                    for block in message.content:
                        if isinstance(block, dict) and "thinking" in block:
                            clean_think = block["thinking"]
                            yield format_chat_chunk("thinking", clean_think)
                            await asyncio.sleep(0.01)
                            
                # SCENARIO B: Final Output (Content is a plain string)
                elif isinstance(message.content, str) and message.content.strip():
                    clean_text = message.content
                    yield format_chat_chunk("message", clean_text)
                    await asyncio.sleep(0.01)
                    final_ai_message = message # Save this to append to history later

        # Append final message to maintain context for the next turn
        if final_ai_message:
            GLOBAL_CHAT_HISTORY.append(final_ai_message)
            
        # Tell the frontend the stream is completely done
        yield format_chat_chunk("end", "[END]")
        await asyncio.sleep(0.01)

    return StreamingResponse(generate_response(), media_type="text/event-stream")

# 1. Find the path to the front_end folder
current_dir = os.path.dirname(os.path.abspath(__file__))
frontend_dir = os.path.join(current_dir, "..", "front_end")

# 2. Tell FastAPI to serve your index.html when people visit the main URL
app.mount("/", StaticFiles(directory=frontend_dir, html=True), name="frontend")