Spaces:
Sleeping
Sleeping
File size: 7,754 Bytes
a642c75 528fd91 a642c75 461dcfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | # 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")
|