main.py

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