src/server.py

"""FastAPI server for multi-accelerator startup search."""

import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path

# Ensure project root is on sys.path when running as `python src/server.py`
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import subprocess
from fastapi import FastAPI, Query, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, PlainTextResponse, Response

from src.search import SearchEngine
from src.render import render_company_page, company_path, SITE_URL
from src.embeddings import embeddings_filename, get_provider

engine: SearchEngine | None = None


def _embedding_files_to_fetch() -> list[str]:
    """Return filenames that should be pulled from HF Hub (only when configured)."""
    if not os.environ.get("HF_DATASET_REPO"):
        return []
    provider = get_provider()
    return [
        embeddings_filename(provider, multi=True),
        embeddings_filename(provider, multi=False),
    ]


@asynccontextmanager
async def lifespan(app: FastAPI):
    global engine
    files = _embedding_files_to_fetch()
    if files:
        from src.storage import ensure_embeddings
        ensure_embeddings(files)
    engine = SearchEngine()
    print(f"Loaded {len(engine.companies)} companies with embeddings of shape {engine.embeddings.shape}")
    print(f"Embedding provider: {engine.embedder.provider}")
    yield


app = FastAPI(title="NicheFind — Startup Discovery Engine", lifespan=lifespan)

STATIC_DIR = Path(__file__).resolve().parent.parent / "static"


@app.get("/api/search")
def api_search(
    q: str = Query(..., min_length=1, description="Search query — describe your problem or idea"),
    min_score: float | None = Query(None, ge=0, le=1),
    source: str | None = Query(None),
    batch: str | None = Query(None),
    industry: str | None = Query(None),
    status: str | None = Query(None),
    tag: str | None = Query(None),
):
    score = min_score if min_score is not None else engine._default_min_score
    results = engine.search(
        q, min_score=score, source=source, batch=batch,
        industry=industry, status=status, tag=tag,
    )
    return {"query": q, "reranker": engine.reranker, "count": len(results), "results": results}


@app.get("/api/filters")
def api_filters():
    return engine.filters


@app.get("/api/company/{source}/{slug}")
def api_company(source: str, slug: str):
    company = engine.get_company(source, slug)
    if company is None:
        raise HTTPException(status_code=404, detail="Company not found")
    return {
        "company": company,
        "similar": engine.similar_companies(source, slug),
    }


app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


try:
    _GIT_REV = subprocess.check_output(
        ["git", "rev-parse", "--short", "HEAD"], stderr=subprocess.DEVNULL
    ).decode().strip()
except Exception:
    _GIT_REV = "1"


@app.get("/")
def index():
    html = (STATIC_DIR / "index.html").read_text()
    html = html.replace("/static/style.css", f"/static/style.css?v={_GIT_REV}")
    html = html.replace("/static/app.js", f"/static/app.js?v={_GIT_REV}")
    return HTMLResponse(html)


@app.get("/company/{source}/{slug}")
def company_page(source: str, slug: str):
    company = engine.get_company(source, slug)
    if company is None:
        return HTMLResponse(
            "<h1>Company not found</h1><p><a href='/'>Back to search</a></p>", status_code=404
        )
    similar = engine.similar_companies(source, slug)
    return HTMLResponse(render_company_page(company, similar, git_rev=_GIT_REV))


@app.get("/robots.txt")
def robots():
    return PlainTextResponse(f"User-agent: *\nAllow: /\nSitemap: {SITE_URL}/sitemap.xml\n")


@app.get("/sitemap.xml")
def sitemap():
    urls = [f"{SITE_URL}/"] + [SITE_URL + company_path(c) for c in engine.companies]
    body = "".join(f"<url><loc>{u}</loc></url>" for u in urls)
    xml = (
        '<?xml version="1.0" encoding="UTF-8"?>'
        '<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">'
        f"{body}</urlset>"
    )
    return Response(content=xml, media_type="application/xml")


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run("src.server:app", host="0.0.0.0", port=port)