main.py (2512B)
1 import os 2 import pickle 3 from uuid import uuid4 as uuid 4 5 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer 6 from fastapi import FastAPI, HTTPException, UploadFile 7 from fastapi.responses import HTMLResponse 8 9 address = os.environ.get("KAFKA_ADDRESS", "localhost:9092") 10 11 12 app = FastAPI() 13 14 15 @app.get("/", response_class=HTMLResponse) 16 async def index(): 17 return HTMLResponse( 18 content=""" 19 <html> 20 <form enctype="multipart/form-data" method="post"> 21 <input type="file" name="file" accept="audio/*" /> 22 <button type="submit">Submit</button> 23 </form> 24 </html> 25 """ 26 ) 27 28 29 @app.post("/") 30 async def summarize(file: UploadFile): 31 job_id = str.encode(str(uuid())) 32 33 consumer = None 34 producer = None 35 36 try: 37 consumer = AIOKafkaConsumer( 38 "transcription.output", 39 bootstrap_servers=address, 40 consumer_timeout_ms=300000, 41 ) 42 await consumer.start() 43 44 producer = AIOKafkaProducer( 45 max_request_size=26214400, bootstrap_servers=address 46 ) 47 await producer.start() 48 49 print("transcribing file") 50 await producer.send_and_wait( 51 "transcription.input", 52 key=job_id, 53 value=pickle.dumps( 54 { 55 "name": file.filename, 56 "data": await file.read(), 57 } 58 ), 59 ) 60 61 text = "" 62 async for message in consumer: 63 if message.key == job_id: 64 text = message.value.decode() 65 break 66 67 print("summarizing text") 68 await consumer.stop() 69 consumer = AIOKafkaConsumer( 70 "summarization.output", 71 bootstrap_servers=address, 72 consumer_timeout_ms=300000, 73 ) 74 await consumer.start() 75 76 await producer.send_and_wait( 77 "summarization.input", key=job_id, value=str.encode(text) 78 ) 79 80 summary = "" 81 async for message in consumer: 82 if message.key == job_id: 83 summary = message.value.decode() 84 break 85 86 print("request completed") 87 return { 88 "text": text, 89 "summary": summary, 90 } 91 92 except Exception as e: 93 print(f"Error processing request: {e}") 94 raise HTTPException(status_code=500, detail=str(e)) 95 96 finally: 97 if consumer: 98 await consumer.stop() 99 if producer: 100 await producer.stop()