summarizer

commit f7ac254c936de7a6fb09f5ed0ffa11ba8247c631
parent 05ec278686aa51c5046004be1dba3df6bff7fce6
Author: Damian Myrda <damian@prime8.dev>
Date:   Sun, 16 Jun 2024 16:25:02 -0500

use send_and_wait instead of only send

Diffstat:
M
main.py
 | 
15
++++++---------
M
summarize/main.py
 | 
9
++++-----
M
transcribe/main.py
 | 
14
++++++--------

3 files changed, 16 insertions(+), 22 deletions(-)
diff --git a/main.py b/main.py
@@ -24,12 +24,13 @@ async def index():
 
 @app.post("/")
 async def summarize(file: UploadFile):
-    print("transribing file")
     producer = AIOKafkaProducer(
         max_request_size=26214400, bootstrap_servers=address
     )
     await producer.start()
-    await producer.send(
+
+    print("transribing file")
+    await producer.send_and_wait(
         "transcription.input",
         pickle.dumps(
             {
@@ -38,7 +39,6 @@ async def summarize(file: UploadFile):
             }
         ),
     )
-    await producer.stop()
     consumer = AIOKafkaConsumer(
         "transcription.output", bootstrap_servers=address
     )
@@ -48,12 +48,7 @@ async def summarize(file: UploadFile):
     text = message.value.decode()
 
     print("summarizing text")
-    producer = AIOKafkaProducer(
-        max_request_size=26214400, bootstrap_servers=address
-    )
-    await producer.start()
-    await producer.send("summarization.input", str.encode(text))
-    await producer.stop()
+    await producer.send_and_wait("summarization.input", str.encode(text))
     consumer = AIOKafkaConsumer(
         "summarization.output", bootstrap_servers=address
     )
@@ -62,6 +57,8 @@ async def summarize(file: UploadFile):
     await consumer.stop()
     summary = message.value.decode()
 
+    await producer.stop()
+
     print("request completed")
     return {
         "text": text,
diff --git a/summarize/main.py b/summarize/main.py
@@ -8,6 +8,8 @@ address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
 
 
 async def summarizer():
+    producer = AIOKafkaProducer(bootstrap_servers=address)
+    await producer.start()
     consumer = AIOKafkaConsumer(
         "summarization.input", bootstrap_servers=address
     )
@@ -27,12 +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))
 
-        producer = AIOKafkaProducer(bootstrap_servers=address)
-        await producer.start()
-        await producer.send("summarization.output", str.encode(summary))
-        await producer.stop()
-
+    await producer.stop()
     await consumer.stop()
 
 
diff --git a/transcribe/main.py b/transcribe/main.py
@@ -10,6 +10,10 @@ address = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
 
 
 async def transcribe():
+    producer = AIOKafkaProducer(
+        max_request_size=26214400, bootstrap_servers=address
+    )
+    await producer.start()
     consumer = AIOKafkaConsumer(
         "transcription.input", bootstrap_servers=address
     )
@@ -30,16 +34,10 @@ async def transcribe():
                 model="whisper-1", file=audio_file
             )
             text += transcription.text
+        await producer.send_and_wait("transcription.output", str.encode(text))
 
-        producer = AIOKafkaProducer(
-            max_request_size=26214400, bootstrap_servers=address
-        )
-        await producer.start()
-        await producer.send("transcription.output", str.encode(text))
-        await producer.stop()
-
+    await producer.stop()
     await consumer.stop()
 
-
 if __name__ == "__main__":
     asyncio.run(transcribe())