commit ee882aac260fabe2e45b39a8f887dc5b0547db89
parent c18182494ca9e7f877f31cb636257ac19a0a959a
Author: Damian Myrda <damian@prime8.dev>
Date: Wed, 18 Feb 2026 14:36:41 -0600
fix critical bug and add error handling and graceful shutdown
Diffstat:
| M | main.py | | | 140 | +++++++++++++++++++++++++++++++++++++++++++++++-------------------------------- |
| M | summarize/main.py | | | 92 | ++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------------- |
| M | transcribe/main.py | | | 108 | ++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------------- |
3 files changed, 221 insertions(+), 119 deletions(-)
diff --git a/main.py b/main.py
@@ -1,21 +1,32 @@
-from fastapi import FastAPI, UploadFile
-from fastapi.responses import HTMLResponse
import os
-from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
-from uuid import uuid4 as uuid
import pickle
+import signal
+import sys
+from contextlib import asynccontextmanager
+from uuid import uuid4 as uuid
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from fastapi import FastAPI, HTTPException, UploadFile
+from fastapi.responses import HTMLResponse
-app = FastAPI()
address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+ yield
+
+
+app = FastAPI(lifespan=lifespan)
+
+
@app.get("/", response_class=HTMLResponse)
async def index():
return HTMLResponse(
content="""
<html>
<form enctype="multipart/form-data" method="post">
- <input type="file" name="file" />
+ <input type="file" name="file" accept="audio/*" />
<button type="submit">Submit</button>
</form>
</html>
@@ -27,54 +38,71 @@ async def index():
async def summarize(file: UploadFile):
job_id = str.encode(str(uuid()))
- producer = AIOKafkaProducer(
- max_request_size=26214400, bootstrap_servers=address
- )
- await producer.start()
- print("transribing file")
- consumer = AIOKafkaConsumer(
- "transcription.output", bootstrap_servers=address
- )
- await consumer.start()
- await producer.send_and_wait(
- "transcription.input",
- key=job_id,
- value=pickle.dumps(
- {
- "name": file.filename,
- "data": await file.read(),
- }
- ),
- )
- text = ""
- while True:
- message = await consumer.getone()
- if message.key == job_id:
- text = message.value.decode()
- break
- await consumer.stop()
-
- print("summarizing text")
- consumer = AIOKafkaConsumer(
- "summarization.output", bootstrap_servers=address
- )
- await consumer.start()
- await producer.send_and_wait(
- "summarization.input", key=job_id, value=str.encode(text)
- )
+ consumer = None
+ producer = None
+
+ try:
+ consumer = AIOKafkaConsumer(
+ "transcription.output",
+ bootstrap_servers=address,
+ consumer_timeout_ms=300000,
+ )
+ await consumer.start()
+
+ producer = AIOKafkaProducer(
+ max_request_size=26214400, bootstrap_servers=address
+ )
+ await producer.start()
+
+ print("transcribing file")
+ await producer.send_and_wait(
+ "transcription.input",
+ key=job_id,
+ value=pickle.dumps(
+ {
+ "name": file.filename,
+ "data": await file.read(),
+ }
+ ),
+ )
+
+ text = ""
+ async for message in consumer:
+ if message.key == job_id:
+ text = message.value.decode()
+ break
+
+ print("summarizing text")
+ await consumer.stop()
+ consumer = AIOKafkaConsumer(
+ "summarization.output",
+ bootstrap_servers=address,
+ consumer_timeout_ms=300000,
+ )
+ await consumer.start()
+
+ await producer.send_and_wait(
+ "summarization.input", key=job_id, value=str.encode(text)
+ )
+
+ summary = ""
+ async for message in consumer:
+ if message.key == job_id:
+ summary = message.value.decode()
+ break
+
+ print("request completed")
+ return {
+ "text": text,
+ "summary": summary,
+ }
+
+ except Exception as e:
+ print(f"Error processing request: {e}")
+ raise HTTPException(status_code=500, detail=str(e))
- summary = ""
- while True:
- message = await consumer.getone()
- if message.key == job_id:
- summary = message.value.decode()
- break
- await consumer.stop()
-
- await producer.stop()
-
- print("request completed")
- return {
- "text": text,
- "summary": summary,
- }
+ finally:
+ if consumer:
+ await consumer.stop()
+ if producer:
+ await producer.stop()
diff --git a/summarize/main.py b/summarize/main.py
@@ -1,40 +1,74 @@
-import os
import asyncio
-from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
+import os
+import signal
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
from openai import AsyncOpenAI
-client = AsyncOpenAI()
-address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
+openAIClient = AsyncOpenAI()
+
+kafkaAddress = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
+
+running = True
+
+
+def signal_handler(signum, frame):
+ global running
+ print(f"Received signal {signum}, shutting down...")
+ running = False
async def summarizer():
- producer = AIOKafkaProducer(bootstrap_servers=address)
- await producer.start()
- consumer = AIOKafkaConsumer(
- "summarization.input", bootstrap_servers=address
- )
- await consumer.start()
-
- while True:
- message = await consumer.getone()
- text = message.value.decode()
-
- summarization = await client.chat.completions.create(
- messages=[
- {
- "role": "user",
- "content": "summarize the following:\n" + text,
- }
- ],
- model="gpt-4-turbo",
- )
- summary = summarization.choices[0].message.content
- await producer.send_and_wait(
- "summarization.output", key=message.key, value=str.encode(summary)
+ global running
+
+ signal.signal(signal.SIGTERM, signal_handler)
+ signal.signal(signal.SIGINT, signal_handler)
+
+ producer = None
+ consumer = None
+
+ try:
+ producer = AIOKafkaProducer(bootstrap_servers=kafkaAddress)
+ await producer.start()
+
+ consumer = AIOKafkaConsumer(
+ "summarization.input", bootstrap_servers=kafkaAddress
)
+ await consumer.start()
+
+ while running:
+ try:
+ message = await asyncio.wait_for(consumer.getone(), timeout=1.0)
+ except asyncio.TimeoutError:
+ continue
+
+ text = message.value.decode()
+
+ try:
+ summarization = await openAIClient.chat.completions.create(
+ messages=[
+ {
+ "role": "user",
+ "content": "summarize the following:\n" + text,
+ }
+ ],
+ model="gpt-4-turbo",
+ )
+ summary = summarization.choices[0].message.content
+ except Exception as e:
+ print(f"Summarization failed: {e}")
+ summary = f"[ERROR: Summarization failed - {e}]"
+
+ await producer.send_and_wait(
+ "summarization.output", key=message.key, value=str.encode(summary)
+ )
- await producer.stop()
- await consumer.stop()
+ finally:
+ if consumer:
+ await consumer.stop()
+ if producer:
+ await producer.stop()
+ print("Summarizer shut down")
if __name__ == "__main__":
diff --git a/transcribe/main.py b/transcribe/main.py
@@ -1,45 +1,85 @@
-import os
import asyncio
-from openai import AsyncOpenAI
-from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
-import pickle
import io
+import os
+import pickle
+import signal
+import sys
+
+from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from openai import AsyncOpenAI
-client = AsyncOpenAI()
-address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
+CHUNK_SIZE = 26214400 - 1024
+
+openAIClient = AsyncOpenAI()
+
+kafkaAddress = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
+
+running = True
+
+
+def signal_handler(signum, frame):
+ global running
+ print(f"Received signal {signum}, shutting down...")
+ running = False
async def transcribe():
- producer = AIOKafkaProducer(
- max_request_size=26214400, bootstrap_servers=address
- )
- await producer.start()
- consumer = AIOKafkaConsumer(
- "transcription.input", bootstrap_servers=address
- )
- await consumer.start()
-
- while True:
- message = await consumer.getone()
- file = pickle.loads(message.value)
-
- text = ""
-
- CHUNK_SIZE = 26214400 - 1024
- for i in range(0, len(file["data"]), CHUNK_SIZE):
- audio_file = io.BytesIO(file["data"][i : i + CHUNK_SIZE])
- audio_file.name = file["name"]
-
- transcription = await client.audio.transcriptions.create(
- model="whisper-1", file=audio_file
- )
- text += transcription.text
- await producer.send_and_wait(
- "transcription.output", key=message.key, value=str.encode(text)
+ global running
+
+ signal.signal(signal.SIGTERM, signal_handler)
+ signal.signal(signal.SIGINT, signal_handler)
+
+ producer = None
+ consumer = None
+
+ try:
+ producer = AIOKafkaProducer(
+ max_request_size=26214400, bootstrap_servers=kafkaAddress
)
+ await producer.start()
+
+ consumer = AIOKafkaConsumer(
+ "transcription.input", bootstrap_servers=kafkaAddress
+ )
+ await consumer.start()
+
+ while running:
+ try:
+ message = await asyncio.wait_for(consumer.getone(), timeout=1.0)
+ except asyncio.TimeoutError:
+ continue
+
+ try:
+ file = pickle.loads(message.value)
+ except Exception as e:
+ print(f"Failed to deserialize message: {e}")
+ continue
+
+ text = ""
+
+ try:
+ for i in range(0, len(file["data"]), CHUNK_SIZE):
+ audio_file = io.BytesIO(file["data"][i : i + CHUNK_SIZE])
+ audio_file.name = file["name"]
+
+ transcription = await openAIClient.audio.transcriptions.create(
+ model="whisper-1", file=audio_file
+ )
+ text += transcription.text
+ except Exception as e:
+ print(f"Transcription failed: {e}")
+ text = f"[ERROR: Transcription failed - {e}]"
+
+ await producer.send_and_wait(
+ "transcription.output", key=message.key, value=str.encode(text)
+ )
- await producer.stop()
- await consumer.stop()
+ finally:
+ if consumer:
+ await consumer.stop()
+ if producer:
+ await producer.stop()
+ print("Transcriber shut down")
if __name__ == "__main__":
