File size: 4,522 Bytes
23bd097
43e97e3
a537e5e
ec10d0e
c0a983b
ec10d0e
c4f4d59
43e97e3
428a61d
 
 
 
6d73c15
 
8e6b116
43e97e3
c4f4d59
 
 
 
 
a537e5e
43e97e3
c4f4d59
 
 
 
 
428a61d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c4f4d59
23bd097
 
 
 
8e6b116
6d73c15
 
 
 
8e6b116
6d73c15
c0a983b
a537e5e
 
8e6b116
 
 
 
 
 
 
 
 
 
 
c0a983b
 
 
 
 
 
 
 
 
 
 
 
 
 
ec10d0e
 
 
 
 
 
 
c4f4d59
 
 
43e97e3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
129
130
131
132
133
134
135
# main.py
from fastapi import FastAPI,status,Response,Request,Depends,HTTPException,UploadFile, File
from fastapi.responses import StreamingResponse,FileResponse
from models import load_text_model,generate_text,load_audio_model,generate_audio,load_image_model, generate_image
from schemas import VoicePresets
from utils import audio_array_to_buffer,img_to_bytes
from contextlib import asynccontextmanager
from typing import AsyncIterator,Callable,Awaitable,Annotated
from uuid import uuid4
import time
from datetime import datetime, timezone
import csv
from dependencies import get_urls_content
from schemas import TextModelResponse,TextModelRequest
import shutil, uuid
from upload import save_file

models = {} 

@asynccontextmanager 
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    # models["text2image"] = load_image_model() 
    # models["text"]=load_text_model()
    yield 
    models.clear() 

app = FastAPI(lifespan=lifespan)

csv_header = [
    "Request ID", "Datetime", "Endpoint Triggered", "Client IP Address",
    "Response Time", "Status Code", "Successful"
]


@app.middleware("http") 
async def monitor_service(
    req: Request, call_next: Callable[[Request], Awaitable[Response]]
) -> Response: 
    request_id = uuid4().hex 
    request_datetime = datetime.now(timezone.utc).isoformat()
    start_time = time.perf_counter()
    response: Response = await call_next(req)
    response_time = round(time.perf_counter() - start_time, 4) 
    response.headers["X-Response-Time"] = str(response_time)
    response.headers["X-API-Request-ID"] = request_id 
    with open("usage.csv", "a", newline="") as file:
        writer = csv.writer(file)
        if file.tell() == 0:
            writer.writerow(csv_header)
        writer.writerow( 
            [
                request_id,
                request_datetime,
                req.url,
                req.client.host,
                response_time,
                response.status_code,
                response.status_code < 400,
            ]
        )
    return response




# app = FastAPI() 
@app.get("/")
def root_controller():
    return {"status": "healthy"}

@app.post("/generate/text") 
async def serve_language_model_controller(request: Request,
    body: TextModelRequest ,
    urls_content: str = Depends(get_urls_content)) -> TextModelResponse: 
    prompt = body.prompt + " " + urls_content
    output =  generate_text(models["text"], prompt, body.temperature)
    return TextModelResponse(content=output, ip=request.client.host)

@app.get("/logs")
def get_logs():
    # return FileResponse("usage.csv", media_type='text/csv', filename="usage.csv")
    temp_file = f"temp_{uuid.uuid4().hex}.csv"
    shutil.copyfile("usage.csv", temp_file)

    # Return file and ensure FastAPI deletes it after sending
    return FileResponse(
        temp_file,
        media_type="text/csv",
        filename="logs.csv",
        headers={"Content-Disposition": "attachment; filename=logs.csv"}
    )

@app.get(
    "/generate/audio",
    responses={status.HTTP_200_OK: {"content": {"audio/wav": {}}}},
    response_class=StreamingResponse,
) 
def serve_text_to_audio_model_controller(
    prompt: str,
    preset: VoicePresets = "v2/en_speaker_1",
):
    processor, model = load_audio_model()
    output, sample_rate = generate_audio(processor, model, prompt, preset)
    return StreamingResponse(
        audio_array_to_buffer(output, sample_rate), media_type="audio/wav"
    )


@app.get("/generate/image",
         responses={status.HTTP_200_OK: {"content": {"image/png": {}}}}, 
         response_class=Response) 
def serve_text_to_image_model_controller(prompt: str):
    # pipe = load_image_model()
    # output = generate_image(pipe, prompt) 
    output = generate_image(models["text2image"], prompt)
    return Response(content=img_to_bytes(output), media_type="image/png") 

@app.post("/upload")
async def file_upload_controller(
    file: Annotated[UploadFile, File(description="Uploaded PDF documents")]
):
    if file.content_type != "application/pdf":
        raise HTTPException(
            detail=f"Only uploading PDF documents are supported",
            status_code=status.HTTP_400_BAD_REQUEST,
        )
    try:
        await save_file(file)
    except Exception as e:
        raise HTTPException(
            detail=f"An error occurred while saving file - Error: {e}",
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    return {"filename": file.filename, "message": "File uploaded successfully"}