This change introduces status updates to indicate the start and successful completion of the report generation. These progress messages improve user feedback during the asynchronous operation.
48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
from fastapi import FastAPI, HTTPException, Request, Depends
|
|
from pydantic import BaseModel
|
|
from phoenix_technologies import ReportGenerator
|
|
from fastapi.responses import StreamingResponse
|
|
import os
|
|
import asyncio
|
|
|
|
# FastAPI app instance
|
|
app = FastAPI()
|
|
|
|
# Define a request body structure using Pydantic
|
|
class ReportRequest(BaseModel):
|
|
query: str
|
|
report_type: str
|
|
|
|
# Define a dependency to validate the API Key
|
|
def verify_api_key(request: Request):
|
|
# Define the API key from the environment variables
|
|
expected_api_key = os.getenv("API_KEY", None)
|
|
if not expected_api_key:
|
|
raise HTTPException(status_code=500, detail="API key is not configured on the server.")
|
|
|
|
# Get the API key from the request headers
|
|
provided_api_key = request.headers.get("X-API-KEY", None)
|
|
|
|
# Check if the API key is correct
|
|
if not provided_api_key or provided_api_key != expected_api_key:
|
|
raise HTTPException(status_code=403, detail="Invalid or missing API key.")
|
|
|
|
@app.post("/get_report", dependencies=[Depends(verify_api_key)])
|
|
async def get_report_endpoint(request: ReportRequest):
|
|
"""
|
|
Expose the `get_report` function as a POST API endpoint, with a streaming response.
|
|
"""
|
|
|
|
async def generate_report():
|
|
try:
|
|
# Call the asynchronous get_report function
|
|
yield "Report generation started...\n"
|
|
generator = ReportGenerator(request.query, request.report_type)
|
|
custom_logs_handler = await generator.generate_report()
|
|
yield "Report generation completed successfully!\n"
|
|
except Exception as e:
|
|
yield f"Error: {str(e)}"
|
|
|
|
# Return streaming response
|
|
return StreamingResponse(generate_report(), media_type="text/plain")
|
|
|