Spaces:

toasted1920
/

GitRead

Sleeping

App Files Files Community

GitRead / back_end /main.py

toasted1920

Claude Opus 5.5

Add bring-your-own-key LLMs, ONNX code embeddings and security hardening

5e9ddc1 5 days ago

Raw History Blame Contribute Delete

18.5 kB

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# main.py
import asyncio
import json
import logging
import threading
import time
from contextlib import asynccontextmanager
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from back_end.agent.graph import build_workflow
from back_end.agent.tools import get_code_search_tools
from back_end.config import (
    ALLOW_SERVER_KEY_FALLBACK,
    ALLOWED_ORIGINS,
    EMBEDDING_TIME_BUDGET_SECONDS,
    ENABLE_API_DOCS,
    MAX_CHAT_MESSAGE_CHARS,
    MAX_CHUNKS_TO_EMBED,
    MAX_CONCURRENT_INDEXING,
    MAX_FILES_TO_CREATE_VECTOR_DB,
    RATE_LIMIT_CHAT_PER_MIN,
    RATE_LIMIT_GLOBAL_SESSION_PER_MIN,
    RATE_LIMIT_INIT_PER_MIN,
    RATE_LIMIT_IP_CHAT_PER_MIN,
    RATE_LIMIT_IP_INIT_PER_MIN,
    RATE_LIMIT_SESSION_PER_MIN,
    SESSION_IDLE_TTL_SECONDS,
)
from back_end.core.downloader import RepoError, delete_dir, download_github_repo, parse_github_url
from back_end.core.embeddings import build_vector_db, get_embedding_model
from back_end.core.loader import count_valid_supported_files, load_repository_as_documents
from back_end.core.splitter import custom_splitter
from back_end.llm import PROVIDERS, LLMConfigError, friendly_llm_error, resolve_llm_config
from back_end.security import RateLimiter, SecurityHeadersMiddleware, client_ip
from back_end.sessions import Session, SessionCapacityError, store
from langchain_core.messages import HumanMessage

logger = logging.getLogger("gitread")
# HTTP client libraries log every request at INFO; keep them quiet
for noisy in ("httpx", "httpcore", "urllib3", "google_genai", "openai", "anthropic"):
    logging.getLogger(noisy).setLevel(logging.WARNING)

MAX_HISTORY_MESSAGES = 40


# ─────────────────────────── lifecycle ───────────────────────────

async def _periodic_cleanup():
    while True:
        await asyncio.sleep(300)
        try:
            await asyncio.to_thread(store.cleanup_expired)
        except Exception:
            logger.exception("Session cleanup failed")


@asynccontextmanager
async def lifespan(app: FastAPI):
    store.wipe_storage_root()
    if os.getenv("PRELOAD_EMBEDDING_MODEL", "true").lower() == "true":
        # Warm the model in the background so the first visitor doesn't wait for it
        threading.Thread(target=get_embedding_model, daemon=True).start()
    task = asyncio.create_task(_periodic_cleanup())
    yield
    task.cancel()


app = FastAPI(
    lifespan=lifespan,
    debug=False,
    docs_url="/docs" if ENABLE_API_DOCS else None,
    redoc_url=None,
    openapi_url="/openapi.json" if ENABLE_API_DOCS else None,
)

# The frontend is served from this same origin, so CORS is only needed if you
# explicitly list other origins in ALLOWED_ORIGINS. Never "*" with credentials.
if ALLOWED_ORIGINS:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=ALLOWED_ORIGINS,
        allow_credentials=False,
        allow_methods=["GET", "POST"],
        allow_headers=["Authorization", "Content-Type"],
    )

# Added last = outermost, so even CORS / error responses get the headers
app.add_middleware(SecurityHeadersMiddleware)


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    # FastAPI's default 422 echoes the submitted values back (which can include
    # an API key). Only report which fields were wrong.
    fields = sorted({".".join(str(p) for p in e.get("loc", [])[1:]) or "body" for e in exc.errors()})
    return JSONResponse({"detail": f"Invalid request: {', '.join(fields)}"}, status_code=422)


@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    # Log server-side, never send stack traces / internals to the client
    logger.exception("Unhandled error on %s", request.url.path)
    return JSONResponse({"detail": "Internal server error."}, status_code=500)


session_ip_limiter = RateLimiter(RATE_LIMIT_SESSION_PER_MIN)
session_global_limiter = RateLimiter(RATE_LIMIT_GLOBAL_SESSION_PER_MIN)
init_session_limiter = RateLimiter(RATE_LIMIT_INIT_PER_MIN)
init_ip_limiter = RateLimiter(RATE_LIMIT_IP_INIT_PER_MIN)
chat_session_limiter = RateLimiter(RATE_LIMIT_CHAT_PER_MIN)
chat_ip_limiter = RateLimiter(RATE_LIMIT_IP_CHAT_PER_MIN)
indexing_slots = asyncio.Semaphore(MAX_CONCURRENT_INDEXING)


# ─────────────────────────── auth ───────────────────────────

def require_session(request: Request) -> Session:
    auth = request.headers.get("authorization", "")
    scheme, _, credentials = auth.partition(" ")
    session = store.get(credentials.strip()) if scheme.lower() == "bearer" else None
    if session is None:
        raise HTTPException(
            status_code=401,
            detail="Session expired. Please reload the page.",
            headers={"WWW-Authenticate": "Bearer"},
        )
    return session


# ─────────────────────────── request models ───────────────────────────

class RepoRequest(BaseModel):
    url: str = Field(..., min_length=19, max_length=200)


class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=MAX_CHAT_MESSAGE_CHARS)
    provider: str = Field("google", max_length=20)
    model: str | None = Field(None, max_length=100)
    api_key: str | None = Field(None, max_length=300)


def sse(payload: dict) -> str:
    return f"data: {json.dumps(payload)}\n\n"


async def iterate_in_thread(make_iterator, stop: threading.Event):
    """Run a blocking iterator in its own thread so it can't freeze the event loop
    (and with it every other visitor's requests)."""
    loop = asyncio.get_running_loop()
    queue: asyncio.Queue = asyncio.Queue()
    done = object()

    def worker():
        try:
            for item in make_iterator():
                loop.call_soon_threadsafe(queue.put_nowait, item)
                if stop.is_set():  # client went away: stop spending their tokens
                    break
        except BaseException as e:  # forwarded to the async side
            loop.call_soon_threadsafe(queue.put_nowait, e)
        finally:
            loop.call_soon_threadsafe(queue.put_nowait, done)

    threading.Thread(target=worker, daemon=True).start()
    while True:
        item = await queue.get()
        if item is done:
            return
        if isinstance(item, BaseException):
            raise item
        yield item


class _EmbeddingTooSlow(Exception):
    def __init__(self, projected_seconds: float):
        self.projected_seconds = projected_seconds


class _EmbeddingCancelled(Exception):
    pass


async def embed_with_progress(splits: list, outcome: dict, format_msg):
    """Embed in a worker thread while streaming progress (which also keeps the
    connection alive through the proxy). Sets outcome["vector_db"] on success, or
    outcome["warning"] if embedding was skipped (too slow for this hardware, or failed)."""
    state = {"t0": None, "fraction": 0.0}
    cancelled = threading.Event()

    def on_progress(done_chars: int, total_chars: int):
        if cancelled.is_set():
            raise _EmbeddingCancelled()
        now = time.monotonic()
        if state["t0"] is None:  # first call happens before the first batch
            state["t0"] = now
            return
        fraction = done_chars / total_chars
        state["fraction"] = fraction
        elapsed = now - state["t0"]
        if elapsed >= 15 or fraction >= 0.25:
            projected = elapsed / max(fraction, 1e-6)
            if projected > EMBEDDING_TIME_BUDGET_SECONDS:
                raise _EmbeddingTooSlow(projected)

    task = asyncio.ensure_future(asyncio.to_thread(build_vector_db, splits, on_progress))
    task.add_done_callback(lambda t: t.cancelled() or t.exception())  # never "exception never retrieved"
    try:
        while not task.done():
            await asyncio.wait({task}, timeout=5)
            if not task.done():
                fraction, t0 = state["fraction"], state["t0"]
                eta = ""
                if t0 is not None and fraction > 0.02:
                    elapsed = time.monotonic() - t0
                    eta = f" (~{max(1, round((elapsed / fraction - elapsed) / 60))} min left)"
                yield format_msg("PROGRESS", f"Embedding {len(splits)} chunks: {int(fraction * 100)}%{eta}")
        outcome["vector_db"] = task.result()
    except _EmbeddingTooSlow as e:
        minutes = round(e.projected_seconds / 60)
        outcome["warning"] = f"Repo is large for this server (~{minutes} min to embed). Using keyword + file search instead."
    except Exception:
        logger.exception("Embedding failed")
        outcome["warning"] = "Semantic search is unavailable right now. Using keyword + file search instead."
    finally:
        cancelled.set()  # if the visitor disconnected, stop the worker at the next batch


# ─────────────────────────── endpoints ───────────────────────────

@app.get("/api/providers")
async def list_providers():
    return {
        "providers": {k: {"label": v["label"], "default_model": v["default_model"]} for k, v in PROVIDERS.items()},
        "server_key_available": ALLOW_SERVER_KEY_FALLBACK and bool(os.getenv("GOOGLE_API_KEY")),
    }


@app.post("/api/session")
async def create_session(request: Request):
    session_ip_limiter.check(client_ip(request))
    session_global_limiter.check("global")
    try:
        session = await asyncio.to_thread(store.create)
    except SessionCapacityError as e:
        raise HTTPException(status_code=503, detail=str(e))
    return {"token": f"{session.session_id}.{session.token}", "expires_in": SESSION_IDLE_TTL_SECONDS}


@app.post("/api/init-repo")
async def initialize_repository(body: RepoRequest, request: Request, session: Session = Depends(require_session)):
    """Downloads repo, builds this session's vector DB, and streams status updates."""
    init_ip_limiter.check(client_ip(request))
    init_session_limiter.check(session.session_id)
    try:
        parse_github_url(body.url)
    except RepoError as e:
        raise HTTPException(status_code=400, detail=str(e))
    if session.busy.locked():
        raise HTTPException(status_code=409, detail="Already working on a request for this session.")
    try:
        await asyncio.to_thread(store.claim_repo_slot, session)
    except SessionCapacityError as e:
        raise HTTPException(status_code=503, detail=str(e))

    async def process_stream():
        def format_msg(status: str, task_name: str):
            session.touch()  # long indexing runs must not count as idle time
            return sse({"status": status.upper(), "task": task_name})

        async with session.busy:
            session.reset_repo()
            storage = session.storage_dir
            ok = False

            if indexing_slots.locked():
                yield format_msg("START", "Waiting for a free slot (another repo is being analysed)...")
            async with indexing_slots:
                try:
                    yield format_msg("START", "Preparing to load repository")
                    await asyncio.to_thread(delete_dir, storage)
                    storage.mkdir(parents=True, exist_ok=True)
                    yield format_msg("SUCCESS", "Ready to load repository")

                    yield format_msg("START", "Downloading repository from GitHub...")
                    await asyncio.to_thread(download_github_repo, body.url, storage)
                    yield format_msg("SUCCESS", "Repo successfully loaded.")

                    yield format_msg("START", "Counting supported files...")
                    file_count = await asyncio.to_thread(count_valid_supported_files, storage)

                    vector_db, all_splits = None, None
                    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...")
                    else:
                        yield format_msg("SUCCESS", f"Found {file_count} files to process")

                        yield format_msg("START", "Loading repository files as documents...")
                        docs = await asyncio.to_thread(load_repository_as_documents, storage)
                        yield format_msg("SUCCESS", "Files loaded")

                        yield format_msg("START", "Preparing files to analyse")
                        all_splits = await asyncio.to_thread(custom_splitter, docs, storage)
                        yield format_msg("SUCCESS", "Done preparing")

                        if len(all_splits) > MAX_CHUNKS_TO_EMBED:
                            yield format_msg("WARNING", f"Repo is large ({len(all_splits)} chunks). Using keyword + file search without Vector DB...")
                        elif all_splits:
                            yield format_msg("START", f"Embedding {len(all_splits)} chunks...")
                            outcome: dict = {}
                            async for msg in embed_with_progress(all_splits, outcome, format_msg):
                                yield msg
                            vector_db = outcome.get("vector_db")
                            if vector_db is not None:
                                yield format_msg("SUCCESS", "Done Analysing")
                            else:
                                yield format_msg("WARNING", outcome["warning"])

                    session.tools = await asyncio.to_thread(
                        get_code_search_tools, storage, vector_db is not None, all_splits, vector_db
                    )
                    session.vector_db = vector_db
                    session.chat_history = []
                    ok = True
                except RepoError as e:
                    # Our own validation messages (bad URL, too large, ...) are safe to show
                    yield format_msg("ERROR", f"Failed to load repository: {e}")
                except Exception:
                    logger.exception("Indexing failed")
                    yield format_msg("ERROR", "Failed to analyse this repository. Check the URL and that the repo is public.")
                finally:
                    if not ok:
                        await asyncio.to_thread(store.release_repo, session)

        if ok:
            yield format_msg("FINISHED", "System ready. Switching to chat.")

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


@app.post("/api/chat")
async def chat_stream(body: ChatRequest, request: Request, session: Session = Depends(require_session)):
    """Streams the LangGraph response back to the frontend using the caller's own LLM key."""
    chat_ip_limiter.check(client_ip(request))
    chat_session_limiter.check(session.session_id)

    if not session.tools:
        raise HTTPException(status_code=400, detail="System not initialized. Please load a repo first.")
    try:
        llm_cfg = resolve_llm_config(body.provider, body.api_key, body.model)
    except LLMConfigError as e:
        raise HTTPException(status_code=400, detail=str(e))
    if session.busy.locked():
        raise HTTPException(status_code=409, detail="Please wait for the current answer to finish.")

    async def generate_response():
        def format_chat_chunk(msg_type: str, text: str):
            return sse({"type": msg_type, "text": text})

        async with session.busy:
            history = session.chat_history[-MAX_HISTORY_MESSAGES:] + [HumanMessage(content=body.message)]
            final_ai_message = None
            stop = threading.Event()
            try:
                workflow = build_workflow(session.tools, llm_cfg)
                events = iterate_in_thread(
                    lambda: workflow.stream({"messages": history}, stream_mode="values", config={"recursion_limit": 100}),
                    stop,
                )
                async for event in events:
                    session.touch()
                    message = event["messages"][-1]

                    if message.type == "ai" and getattr(message, "tool_calls", None):
                        for _ in message.tool_calls:
                            yield format_chat_chunk("tool", "Browsing Codebase..")

                    if message.type == "ai":
                        if isinstance(message.content, list):
                            for block in message.content:
                                # Some models return thinking blocks with empty text (hidden reasoning)
                                if isinstance(block, dict) and block.get("thinking"):
                                    yield format_chat_chunk("thinking", block["thinking"])
                        elif isinstance(message.content, str) and message.content.strip():
                            yield format_chat_chunk("message", message.content)
                            final_ai_message = message
            except Exception as e:
                logger.warning("Chat failed: %s", type(e).__name__)
                yield format_chat_chunk("error", friendly_llm_error(e))
            else:
                session.chat_history = history + ([final_ai_message] if final_ai_message else [])
            finally:
                stop.set()

        yield format_chat_chunk("end", "[END]")

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


# ── Serve frontend — must be LAST ──────────────────────────────────────────────
FRONTEND_DIR = Path(__file__).resolve().parent.parent / "front_end"
app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")

	import sys
	import os
	sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
	# main.py
	import asyncio
	import json
	import logging
	import threading
	import time
	from contextlib import asynccontextmanager
	from pathlib import Path

	from dotenv import load_dotenv

	load_dotenv()

	from fastapi import Depends, FastAPI, HTTPException, Request
	from fastapi.exceptions import RequestValidationError
	from fastapi.middleware.cors import CORSMiddleware
	from fastapi.responses import JSONResponse, StreamingResponse
	from fastapi.staticfiles import StaticFiles
	from pydantic import BaseModel, Field

	from back_end.agent.graph import build_workflow
	from back_end.agent.tools import get_code_search_tools
	from back_end.config import (
	ALLOW_SERVER_KEY_FALLBACK,
	ALLOWED_ORIGINS,
	EMBEDDING_TIME_BUDGET_SECONDS,
	ENABLE_API_DOCS,
	MAX_CHAT_MESSAGE_CHARS,
	MAX_CHUNKS_TO_EMBED,
	MAX_CONCURRENT_INDEXING,
	MAX_FILES_TO_CREATE_VECTOR_DB,
	RATE_LIMIT_CHAT_PER_MIN,
	RATE_LIMIT_GLOBAL_SESSION_PER_MIN,
	RATE_LIMIT_INIT_PER_MIN,
	RATE_LIMIT_IP_CHAT_PER_MIN,
	RATE_LIMIT_IP_INIT_PER_MIN,
	RATE_LIMIT_SESSION_PER_MIN,
	SESSION_IDLE_TTL_SECONDS,
	)
	from back_end.core.downloader import RepoError, delete_dir, download_github_repo, parse_github_url
	from back_end.core.embeddings import build_vector_db, get_embedding_model
	from back_end.core.loader import count_valid_supported_files, load_repository_as_documents
	from back_end.core.splitter import custom_splitter
	from back_end.llm import PROVIDERS, LLMConfigError, friendly_llm_error, resolve_llm_config
	from back_end.security import RateLimiter, SecurityHeadersMiddleware, client_ip
	from back_end.sessions import Session, SessionCapacityError, store
	from langchain_core.messages import HumanMessage

	logger = logging.getLogger("gitread")
	# HTTP client libraries log every request at INFO; keep them quiet
	for noisy in ("httpx", "httpcore", "urllib3", "google_genai", "openai", "anthropic"):
	logging.getLogger(noisy).setLevel(logging.WARNING)

	MAX_HISTORY_MESSAGES = 40


	# ─────────────────────────── lifecycle ───────────────────────────

	async def _periodic_cleanup():
	while True:
	await asyncio.sleep(300)
	try:
	await asyncio.to_thread(store.cleanup_expired)
	except Exception:
	logger.exception("Session cleanup failed")


	@asynccontextmanager
	async def lifespan(app: FastAPI):
	store.wipe_storage_root()
	if os.getenv("PRELOAD_EMBEDDING_MODEL", "true").lower() == "true":
	# Warm the model in the background so the first visitor doesn't wait for it
	threading.Thread(target=get_embedding_model, daemon=True).start()
	task = asyncio.create_task(_periodic_cleanup())
	yield
	task.cancel()


	app = FastAPI(
	lifespan=lifespan,
	debug=False,
	docs_url="/docs" if ENABLE_API_DOCS else None,
	redoc_url=None,
	openapi_url="/openapi.json" if ENABLE_API_DOCS else None,
	)

	# The frontend is served from this same origin, so CORS is only needed if you
	# explicitly list other origins in ALLOWED_ORIGINS. Never "*" with credentials.
	if ALLOWED_ORIGINS:
	app.add_middleware(
	CORSMiddleware,
	allow_origins=ALLOWED_ORIGINS,
	allow_credentials=False,
	allow_methods=["GET", "POST"],
	allow_headers=["Authorization", "Content-Type"],
	)

	# Added last = outermost, so even CORS / error responses get the headers
	app.add_middleware(SecurityHeadersMiddleware)


	@app.exception_handler(RequestValidationError)
	async def validation_exception_handler(request: Request, exc: RequestValidationError):
	# FastAPI's default 422 echoes the submitted values back (which can include
	# an API key). Only report which fields were wrong.
	fields = sorted({".".join(str(p) for p in e.get("loc", [])[1:]) or "body" for e in exc.errors()})
	return JSONResponse({"detail": f"Invalid request: {', '.join(fields)}"}, status_code=422)


	@app.exception_handler(Exception)
	async def unhandled_exception_handler(request: Request, exc: Exception):
	# Log server-side, never send stack traces / internals to the client
	logger.exception("Unhandled error on %s", request.url.path)
	return JSONResponse({"detail": "Internal server error."}, status_code=500)


	session_ip_limiter = RateLimiter(RATE_LIMIT_SESSION_PER_MIN)
	session_global_limiter = RateLimiter(RATE_LIMIT_GLOBAL_SESSION_PER_MIN)
	init_session_limiter = RateLimiter(RATE_LIMIT_INIT_PER_MIN)
	init_ip_limiter = RateLimiter(RATE_LIMIT_IP_INIT_PER_MIN)
	chat_session_limiter = RateLimiter(RATE_LIMIT_CHAT_PER_MIN)
	chat_ip_limiter = RateLimiter(RATE_LIMIT_IP_CHAT_PER_MIN)
	indexing_slots = asyncio.Semaphore(MAX_CONCURRENT_INDEXING)


	# ─────────────────────────── auth ───────────────────────────

	def require_session(request: Request) -> Session:
	auth = request.headers.get("authorization", "")
	scheme, _, credentials = auth.partition(" ")
	session = store.get(credentials.strip()) if scheme.lower() == "bearer" else None
	if session is None:
	raise HTTPException(
	status_code=401,
	detail="Session expired. Please reload the page.",
	headers={"WWW-Authenticate": "Bearer"},
	)
	return session


	# ─────────────────────────── request models ───────────────────────────

	class RepoRequest(BaseModel):
	url: str = Field(..., min_length=19, max_length=200)


	class ChatRequest(BaseModel):
	message: str = Field(..., min_length=1, max_length=MAX_CHAT_MESSAGE_CHARS)
	provider: str = Field("google", max_length=20)
	model: str \| None = Field(None, max_length=100)
	api_key: str \| None = Field(None, max_length=300)


	def sse(payload: dict) -> str:
	return f"data: {json.dumps(payload)}\n\n"


	async def iterate_in_thread(make_iterator, stop: threading.Event):
	"""Run a blocking iterator in its own thread so it can't freeze the event loop
	(and with it every other visitor's requests)."""
	loop = asyncio.get_running_loop()
	queue: asyncio.Queue = asyncio.Queue()
	done = object()

	def worker():
	try:
	for item in make_iterator():
	loop.call_soon_threadsafe(queue.put_nowait, item)
	if stop.is_set(): # client went away: stop spending their tokens
	break
	except BaseException as e: # forwarded to the async side
	loop.call_soon_threadsafe(queue.put_nowait, e)
	finally:
	loop.call_soon_threadsafe(queue.put_nowait, done)

	threading.Thread(target=worker, daemon=True).start()
	while True:
	item = await queue.get()
	if item is done:
	return
	if isinstance(item, BaseException):
	raise item
	yield item


	class _EmbeddingTooSlow(Exception):
	def __init__(self, projected_seconds: float):
	self.projected_seconds = projected_seconds


	class _EmbeddingCancelled(Exception):
	pass


	async def embed_with_progress(splits: list, outcome: dict, format_msg):
	"""Embed in a worker thread while streaming progress (which also keeps the
	connection alive through the proxy). Sets outcome["vector_db"] on success, or
	outcome["warning"] if embedding was skipped (too slow for this hardware, or failed)."""
	state = {"t0": None, "fraction": 0.0}
	cancelled = threading.Event()

	def on_progress(done_chars: int, total_chars: int):
	if cancelled.is_set():
	raise _EmbeddingCancelled()
	now = time.monotonic()
	if state["t0"] is None: # first call happens before the first batch
	state["t0"] = now
	return
	fraction = done_chars / total_chars
	state["fraction"] = fraction
	elapsed = now - state["t0"]
	if elapsed >= 15 or fraction >= 0.25:
	projected = elapsed / max(fraction, 1e-6)
	if projected > EMBEDDING_TIME_BUDGET_SECONDS:
	raise _EmbeddingTooSlow(projected)

	task = asyncio.ensure_future(asyncio.to_thread(build_vector_db, splits, on_progress))
	task.add_done_callback(lambda t: t.cancelled() or t.exception()) # never "exception never retrieved"
	try:
	while not task.done():
	await asyncio.wait({task}, timeout=5)
	if not task.done():
	fraction, t0 = state["fraction"], state["t0"]
	eta = ""
	if t0 is not None and fraction > 0.02:
	elapsed = time.monotonic() - t0
	eta = f" (~{max(1, round((elapsed / fraction - elapsed) / 60))} min left)"
	yield format_msg("PROGRESS", f"Embedding {len(splits)} chunks: {int(fraction * 100)}%{eta}")
	outcome["vector_db"] = task.result()
	except _EmbeddingTooSlow as e:
	minutes = round(e.projected_seconds / 60)
	outcome["warning"] = f"Repo is large for this server (~{minutes} min to embed). Using keyword + file search instead."
	except Exception:
	logger.exception("Embedding failed")
	outcome["warning"] = "Semantic search is unavailable right now. Using keyword + file search instead."
	finally:
	cancelled.set() # if the visitor disconnected, stop the worker at the next batch


	# ─────────────────────────── endpoints ───────────────────────────

	@app.get("/api/providers")
	async def list_providers():
	return {
	"providers": {k: {"label": v["label"], "default_model": v["default_model"]} for k, v in PROVIDERS.items()},
	"server_key_available": ALLOW_SERVER_KEY_FALLBACK and bool(os.getenv("GOOGLE_API_KEY")),
	}


	@app.post("/api/session")
	async def create_session(request: Request):
	session_ip_limiter.check(client_ip(request))
	session_global_limiter.check("global")
	try:
	session = await asyncio.to_thread(store.create)
	except SessionCapacityError as e:
	raise HTTPException(status_code=503, detail=str(e))
	return {"token": f"{session.session_id}.{session.token}", "expires_in": SESSION_IDLE_TTL_SECONDS}


	@app.post("/api/init-repo")
	async def initialize_repository(body: RepoRequest, request: Request, session: Session = Depends(require_session)):
	"""Downloads repo, builds this session's vector DB, and streams status updates."""
	init_ip_limiter.check(client_ip(request))
	init_session_limiter.check(session.session_id)
	try:
	parse_github_url(body.url)
	except RepoError as e:
	raise HTTPException(status_code=400, detail=str(e))
	if session.busy.locked():
	raise HTTPException(status_code=409, detail="Already working on a request for this session.")
	try:
	await asyncio.to_thread(store.claim_repo_slot, session)
	except SessionCapacityError as e:
	raise HTTPException(status_code=503, detail=str(e))

	async def process_stream():
	def format_msg(status: str, task_name: str):
	session.touch() # long indexing runs must not count as idle time
	return sse({"status": status.upper(), "task": task_name})

	async with session.busy:
	session.reset_repo()
	storage = session.storage_dir
	ok = False

	if indexing_slots.locked():
	yield format_msg("START", "Waiting for a free slot (another repo is being analysed)...")
	async with indexing_slots:
	try:
	yield format_msg("START", "Preparing to load repository")
	await asyncio.to_thread(delete_dir, storage)
	storage.mkdir(parents=True, exist_ok=True)
	yield format_msg("SUCCESS", "Ready to load repository")

	yield format_msg("START", "Downloading repository from GitHub...")
	await asyncio.to_thread(download_github_repo, body.url, storage)
	yield format_msg("SUCCESS", "Repo successfully loaded.")

	yield format_msg("START", "Counting supported files...")
	file_count = await asyncio.to_thread(count_valid_supported_files, storage)

	vector_db, all_splits = None, None
	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...")
	else:
	yield format_msg("SUCCESS", f"Found {file_count} files to process")

	yield format_msg("START", "Loading repository files as documents...")
	docs = await asyncio.to_thread(load_repository_as_documents, storage)
	yield format_msg("SUCCESS", "Files loaded")

	yield format_msg("START", "Preparing files to analyse")
	all_splits = await asyncio.to_thread(custom_splitter, docs, storage)
	yield format_msg("SUCCESS", "Done preparing")

	if len(all_splits) > MAX_CHUNKS_TO_EMBED:
	yield format_msg("WARNING", f"Repo is large ({len(all_splits)} chunks). Using keyword + file search without Vector DB...")
	elif all_splits:
	yield format_msg("START", f"Embedding {len(all_splits)} chunks...")
	outcome: dict = {}
	async for msg in embed_with_progress(all_splits, outcome, format_msg):
	yield msg
	vector_db = outcome.get("vector_db")
	if vector_db is not None:
	yield format_msg("SUCCESS", "Done Analysing")
	else:
	yield format_msg("WARNING", outcome["warning"])

	session.tools = await asyncio.to_thread(
	get_code_search_tools, storage, vector_db is not None, all_splits, vector_db
	)
	session.vector_db = vector_db
	session.chat_history = []
	ok = True
	except RepoError as e:
	# Our own validation messages (bad URL, too large, ...) are safe to show
	yield format_msg("ERROR", f"Failed to load repository: {e}")
	except Exception:
	logger.exception("Indexing failed")
	yield format_msg("ERROR", "Failed to analyse this repository. Check the URL and that the repo is public.")
	finally:
	if not ok:
	await asyncio.to_thread(store.release_repo, session)

	if ok:
	yield format_msg("FINISHED", "System ready. Switching to chat.")

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


	@app.post("/api/chat")
	async def chat_stream(body: ChatRequest, request: Request, session: Session = Depends(require_session)):
	"""Streams the LangGraph response back to the frontend using the caller's own LLM key."""
	chat_ip_limiter.check(client_ip(request))
	chat_session_limiter.check(session.session_id)

	if not session.tools:
	raise HTTPException(status_code=400, detail="System not initialized. Please load a repo first.")
	try:
	llm_cfg = resolve_llm_config(body.provider, body.api_key, body.model)
	except LLMConfigError as e:
	raise HTTPException(status_code=400, detail=str(e))
	if session.busy.locked():
	raise HTTPException(status_code=409, detail="Please wait for the current answer to finish.")

	async def generate_response():
	def format_chat_chunk(msg_type: str, text: str):
	return sse({"type": msg_type, "text": text})

	async with session.busy:
	history = session.chat_history[-MAX_HISTORY_MESSAGES:] + [HumanMessage(content=body.message)]
	final_ai_message = None
	stop = threading.Event()
	try:
	workflow = build_workflow(session.tools, llm_cfg)
	events = iterate_in_thread(
	lambda: workflow.stream({"messages": history}, stream_mode="values", config={"recursion_limit": 100}),
	stop,
	)
	async for event in events:
	session.touch()
	message = event["messages"][-1]

	if message.type == "ai" and getattr(message, "tool_calls", None):
	for _ in message.tool_calls:
	yield format_chat_chunk("tool", "Browsing Codebase..")

	if message.type == "ai":
	if isinstance(message.content, list):
	for block in message.content:
	# Some models return thinking blocks with empty text (hidden reasoning)
	if isinstance(block, dict) and block.get("thinking"):
	yield format_chat_chunk("thinking", block["thinking"])
	elif isinstance(message.content, str) and message.content.strip():
	yield format_chat_chunk("message", message.content)
	final_ai_message = message
	except Exception as e:
	logger.warning("Chat failed: %s", type(e).__name__)
	yield format_chat_chunk("error", friendly_llm_error(e))
	else:
	session.chat_history = history + ([final_ai_message] if final_ai_message else [])
	finally:
	stop.set()

	yield format_chat_chunk("end", "[END]")

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


	# ── Serve frontend — must be LAST ──────────────────────────────────────────────
	FRONTEND_DIR = Path(__file__).resolve().parent.parent / "front_end"
	app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")