summarizer

commit 6303521f596414d8ad34dcb161e5b2e27260984a
parent f7ac254c936de7a6fb09f5ed0ffa11ba8247c631
Author: Damian Myrda <damian@prime8.dev>
Date:   Sun, 16 Jun 2024 16:51:46 -0500

added job ids

Diffstat:
M
main.py
 | 
36
+++++++++++++++++++++++++-----------
M
summarize/main.py
 | 
4
+++-
M
transcribe/main.py
 | 
5
++++-

3 files changed, 32 insertions(+), 13 deletions(-)
diff --git a/main.py b/main.py
@@ -2,6 +2,7 @@ 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
 
 app = FastAPI()
@@ -24,38 +25,51 @@ async def index():
 
 @app.post("/")
 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",
-        pickle.dumps(
+        key=job_id,
+        value=pickle.dumps(
             {
                 "name": file.filename,
                 "data": await file.read(),
             }
         ),
     )
-    consumer = AIOKafkaConsumer(
-        "transcription.output", bootstrap_servers=address
-    )
-    await consumer.start()
-    message = await consumer.getone()
+    text = ""
+    while True:
+        message = await consumer.getone()
+        if message.key == job_id:
+            text = message.value.decode()
+            break
     await consumer.stop()
-    text = message.value.decode()
 
     print("summarizing text")
-    await producer.send_and_wait("summarization.input", str.encode(text))
     consumer = AIOKafkaConsumer(
         "summarization.output", bootstrap_servers=address
     )
     await consumer.start()
-    message = await consumer.getone()
+    await producer.send_and_wait(
+        "summarization.input", key=job_id, value=str.encode(text)
+    )
+
+    summary = ""
+    while True:
+        message = await consumer.getone()
+        if message.key == job_id:
+            summary = message.value.decode()
+            break
     await consumer.stop()
-    summary = message.value.decode()
 
     await producer.stop()
 
diff --git a/summarize/main.py b/summarize/main.py
@@ -29,7 +29,9 @@ async def summarizer():
             model="gpt-4-turbo",
         )
         summary = summarization.choices[0].message.content
-        await producer.send_and_wait("summarization.output", str.encode(summary))
+        await producer.send_and_wait(
+            "summarization.output", key=message.key, value=str.encode(summary)
+        )
 
     await producer.stop()
     await consumer.stop()
diff --git a/transcribe/main.py b/transcribe/main.py
@@ -34,10 +34,13 @@ async def transcribe():
                 model="whisper-1", file=audio_file
             )
             text += transcription.text
-        await producer.send_and_wait("transcription.output", str.encode(text))
+        await producer.send_and_wait(
+            "transcription.output", key=message.key, value=str.encode(text)
+        )
 
     await producer.stop()
     await consumer.stop()
 
+
 if __name__ == "__main__":
     asyncio.run(transcribe())