wikihop-server / api_server.py
stillerman's picture
stillerman HF Staff
fastapi server
76887e4
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os
from typing import List, Dict, Any, Optional
from db.wiki_db_sqlite import WikiDBSqlite
app = FastAPI(title="Wiki API Server")
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global database connection
db = None
@app.on_event("startup")
async def startup_db_client():
global db
db_path = os.environ.get("WIKI_DB_PATH")
if not db_path:
raise ValueError("WIKI_DB_PATH environment variable not set")
db = WikiDBSqlite(db_path)
@app.get("/article_count")
async def get_article_count() -> Dict[str, int]:
"""Get the number of articles in the database"""
return {"count": db.get_article_count()}
@app.get("/article_titles")
async def get_article_titles() -> Dict[str, List[str]]:
"""Get all article titles"""
return {"titles": db.get_all_article_titles()}
@app.get("/article")
async def get_article(title: str = Query(..., description="Article title")) -> Dict[str, Any]:
"""Get article data by title"""
article = db.get_article(title)
if not article:
raise HTTPException(status_code=404, detail=f"Article '{title}' not found")
return article
@app.get("/article_exists")
async def article_exists(title: str = Query(..., description="Article title")) -> Dict[str, bool]:
"""Check if an article exists"""
return {"exists": db.article_exists(title)}
@app.get("/article_text")
async def get_article_text(title: str = Query(..., description="Article title")) -> Dict[str, str]:
"""Get the text of an article"""
if not db.article_exists(title):
raise HTTPException(status_code=404, detail=f"Article '{title}' not found")
return {"text": db.get_article_text(title)}
@app.get("/article_links")
async def get_article_links(title: str = Query(..., description="Article title")) -> Dict[str, List[str]]:
"""Get the links of an article"""
if not db.article_exists(title):
raise HTTPException(status_code=404, detail=f"Article '{title}' not found")
return {"links": db.get_article_links(title)}
if __name__ == "__main__":
uvicorn.run("api_server:app", host="0.0.0.0", port=8000, reload=True)