vision / app.py
abdullahalioo's picture
Update app.py
d0fc55f verified
raw
history blame
2.22 kB
import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI
from pydantic import BaseModel
import httpx
# Initialize FastAPI app
app = FastAPI()
# Define request body model for the prompt
class PromptRequest(BaseModel):
prompt: str
# Initialize OpenAI client
token = os.getenv("GITHUB_TOKEN")
if not token:
raise ValueError("GITHUB_TOKEN environment variable not set")
# Use environment variables for endpoint and model, with fallbacks
endpoint = os.getenv("API_ENDPOINT", "https://api.openai.com/v1") # Fallback to OpenAI-compatible endpoint
model = os.getenv("MODEL_NAME", "gpt-4o-mini") # Default to a known model
# Initialize AsyncOpenAI with a custom HTTP client to avoid proxies issue
client = AsyncOpenAI(
base_url=endpoint,
api_key=token,
http_client=httpx.AsyncClient() # Explicitly use httpx.AsyncClient without proxies
)
# Async generator to stream chunks
async def stream_response(prompt: str):
try:
# Create streaming chat completion
stream = await client.chat.completions.create(
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}
],
temperature=1.0,
top_p=1.0,
model=model,
stream=True
)
# Yield each chunk as it arrives
async for chunk in stream:
if chunk.choices and len(chunk.choices) > 0:
content = chunk.choices[0].delta.content or ""
if content:
yield content
except Exception as err:
yield f"Error: {str(err)}"
# Endpoint to handle prompt and stream response
@app.post("/generate")
async def generate_response(request: PromptRequest):
try:
return StreamingResponse(
stream_response(request.prompt),
media_type="text/event-stream"
)
except Exception as err:
raise HTTPException(status_code=500, detail=f"Server error: {str(err)}")
# Health check endpoint for Hugging Face Spaces
@app.get("/")
async def health_check():
return {"status": "healthy"}