Spaces:
Build error
Build error
File size: 4,151 Bytes
2a72045 | 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 | from fastapi import FastAPI, UploadFile, File, Body, HTTPException, Request
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import os
import json
import sys
from dotenv import load_dotenv
from typing import Optional, List
import logging
from model_manager import model_manager
from ocr_engine import ocr_engine
from database import db_manager
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
stream=sys.stdout
)
logger = logging.getLogger(__name__)
load_dotenv()
app = FastAPI(title="AI Platform API")
# Configure CORS
app.add_middleware(
CORSMiddleware,
allow_origins=[
"https://frontend-one-gamma-14.vercel.app",
"http://localhost:3000", # For local development
"http://localhost:8000"
],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/")
async def root():
return {
"name": "Alpha Core AI API",
"version": "1.0.0",
"status": "online",
"endpoints": {
"health": "/health",
"chat": "/chat",
"upload": "/upload-image",
"cleanup": "/cleanup"
}
}
@app.get("/health")
async def health_check():
return {"status": "healthy", "version": "1.0.0"}
class ChatRequest(BaseModel):
message: str
model: str = "tinyllama"
user_id: str = "default_user"
context: Optional[List[dict]] = None
temperature: Optional[float] = 0.7
top_p: Optional[float] = 0.95
max_tokens: Optional[int] = 2048
repeat_penalty: Optional[float] = 1.1
@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
try:
logger.info(f"Chat request: model={request.model}, user={request.user_id}")
def stream_response():
full_response = ""
try:
# Pass context and settings to model manager for memory
params = {
"temperature": request.temperature,
"top_p": request.top_p,
"max_tokens": request.max_tokens,
"repeat_penalty": request.repeat_penalty
}
for token in model_manager.generate_stream(request.model, request.message, request.context, **params):
full_response += token
yield f"data: {json.dumps({'token': token})}\n\n"
logger.info(f"Response generated: {len(full_response)} tokens")
# Final output and DB storage
db_manager.store_message(request.user_id, request.message, "user", request.model)
db_manager.store_message(request.user_id, full_response, "assistant", request.model)
yield "data: [DONE]\n\n"
except Exception as e:
logger.error(f"Stream error: {str(e)}", exc_info=True)
yield f"data: {json.dumps({'error': str(e)})}\n\n"
return StreamingResponse(stream_response(), media_type="text/event-stream")
except Exception as e:
logger.error(f"Chat endpoint error: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.post("/upload-image")
async def upload_image(file: UploadFile = File(...)):
if not file.content_type.startswith("image/"):
raise HTTPException(status_code=400, detail="File must be an image")
try:
content = await file.read()
extracted_text = ocr_engine.extract_text(content)
return {"text": extracted_text}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/cleanup")
async def cleanup_chats():
try:
db_manager.cleanup_old_messages()
return {"message": "Cleanup successful"}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
port = int(os.getenv("PORT", 8000))
uvicorn.run(app, host="0.0.0.0", port=port)
|