File size: 2,006 Bytes
28d8100 |
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 |
import os
import json
import asyncio
from typing import AsyncGenerator
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi import FastAPI, HTTPException
from websearch import QueryRequest, PerplexityClient, parse_perplexity_response
# load .env file
from dotenv import load_dotenv
load_dotenv()
# Initialize FastAPI app
app = FastAPI()
# Add CORS middleware to allow frontend connections
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Adjust this in production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize Perplexity client
perplexity_client = PerplexityClient(
api_key=os.environ["PERPLEXITY_AUTH_TOKEN"]
)
async def generate_stream(query: str) -> AsyncGenerator[str, None]:
"""
Async generator to stream JSON response
Args:
query (str): User query
Yields:
str: JSON-encoded chunks of response
"""
try:
# Fetch response from Perplexity
response = await perplexity_client.generate_response(query)
# Parse the response
parsed_response = parse_perplexity_response(response)
# Stream the parsed response as JSON chunks
yield json.dumps(parsed_response)
except Exception as e:
yield json.dumps({"error": str(e)})
@app.post("/websearch")
async def handle_query(request: QueryRequest):
"""
Endpoint to handle user queries and stream responses
Args:
request (QueryRequest): Query request model
Returns:
StreamingResponse: Streaming JSON response
"""
return StreamingResponse(
generate_stream(request.query),
media_type="application/json"
)
# Optional: Health check endpoint
@app.get("/health")
async def health_check():
return {"status": "healthy"}
# Run the app with:
# uvicorn main:app --reload
# Make sure to set PERPLEXITY_API_KEY environment variable |