Spaces:
Sleeping
Sleeping
File size: 4,938 Bytes
35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 e3b17f1 35c7218 | 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 | """
FastAPI endpoints for the RAG system with streaming support.
"""
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional, List, AsyncGenerator
import json
import asyncio
from src.app_hf import answer_question, add_urls, build_rag_chain, load_documents_from_crawler_cache
from src.crawler import crawl_and_persist, async_crawl_and_persist
app = FastAPI(
title="Documentation RAG API",
description="Retrieval-Augmented Generation API for technical documentation",
version="1.0.0"
)
class QueryRequest(BaseModel):
question: str
urls: Optional[List[str]] = None
doc_dir: Optional[str] = "./my_docs"
class CrawlRequest(BaseModel):
base_url: str
max_depth: Optional[int] = 3
max_pages: Optional[int] = 100
class QueryResponse(BaseModel):
question: str
answer: str
source_urls: Optional[List[str]] = None
@app.get("/health")
async def health_check():
"""Health check endpoint."""
return {"status": "ok", "service": "RAG API"}
@app.post("/query")
async def query(request: QueryRequest) -> QueryResponse:
"""
Query the RAG system.
Returns a complete answer in one response.
"""
try:
if request.urls:
add_urls(request.urls)
answer = answer_question(request.question, request.doc_dir, request.urls)
return QueryResponse(
question=request.question,
answer=answer,
source_urls=request.urls
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query/stream")
async def query_stream(request: QueryRequest):
"""
Query the RAG system with streaming response.
Streams tokens from the LLM as they're generated.
"""
async def generate():
try:
if request.urls:
add_urls(request.urls)
# Get the RAG chain
rag_chain = build_rag_chain(request.doc_dir, request.urls)
# Stream response
yield json.dumps({
"type": "start",
"question": request.question
}).encode() + b"\n"
# Invoke with streaming
response = rag_chain.invoke(request.question)
answer_text = response.content if hasattr(response, "content") else str(response)
# Stream the answer in chunks
chunk_size = 10
for i in range(0, len(answer_text), chunk_size):
chunk = answer_text[i:i + chunk_size]
yield json.dumps({
"type": "token",
"content": chunk
}).encode() + b"\n"
await asyncio.sleep(0.01) # Simulate streaming delay
yield json.dumps({
"type": "end",
"answer": answer_text
}).encode() + b"\n"
except Exception as e:
yield json.dumps({
"type": "error",
"error": str(e)
}).encode() + b"\n"
return StreamingResponse(generate(), media_type="application/x-ndjson")
@app.post("/crawl/prepare")
async def prepare_crawl(request: CrawlRequest):
"""
Endpoint to crawl a website and prepare documents for indexing.
This is an async endpoint that returns crawl job status.
"""
try:
documents = await async_crawl_and_persist(request.base_url, output_path="./crawler_docs.json", max_pages=request.max_pages)
return {
"status": "success",
"documents_crawled": len(documents),
"failed_urls": 0,
"base_url": request.base_url,
"message": f"Successfully crawled and persisted {len(documents)} pages from {request.base_url}"
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/index/from-crawl")
async def index_from_crawl(request: CrawlRequest):
"""
Crawl a website and automatically index its content.
"""
try:
cached_docs = load_documents_from_crawler_cache()
if not cached_docs:
raise HTTPException(
status_code=404,
detail="No cached crawler documents found. Run /crawl/prepare first."
)
build_rag_chain(request.doc_dir, urls=[])
return {
"status": "success",
"documents_indexed": len(cached_docs),
"base_url": request.base_url,
"message": f"Successfully indexed {len(cached_docs)} cached crawler pages"
}
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|