summarizer

summarizer
git clone https://git.prime8.dev/summarizer.git
Log | Files | Refs | README | LICENSE

main.py (2250B)

      1 import asyncio
      2 import io
      3 import os
      4 import pickle
      5 import signal
      6 
      7 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
      8 from openai import AsyncOpenAI
      9 
     10 CHUNK_SIZE = 26214400 - 1024
     11 
     12 openAIClient = AsyncOpenAI()
     13 
     14 kafkaAddress = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
     15 
     16 running = True
     17 
     18 
     19 def signal_handler(signum, frame):
     20     global running
     21     print(f"Received signal {signum}, shutting down...")
     22     running = False
     23 
     24 
     25 async def transcribe():
     26     global running
     27 
     28     signal.signal(signal.SIGTERM, signal_handler)
     29     signal.signal(signal.SIGINT, signal_handler)
     30 
     31     producer = None
     32     consumer = None
     33 
     34     try:
     35         producer = AIOKafkaProducer(
     36             max_request_size=26214400, bootstrap_servers=kafkaAddress
     37         )
     38         await producer.start()
     39 
     40         consumer = AIOKafkaConsumer(
     41             "transcription.input", bootstrap_servers=kafkaAddress
     42         )
     43         await consumer.start()
     44 
     45         while running:
     46             try:
     47                 message = await asyncio.wait_for(consumer.getone(), timeout=1.0)
     48             except asyncio.TimeoutError:
     49                 continue
     50 
     51             try:
     52                 file = pickle.loads(message.value)
     53             except Exception as e:
     54                 print(f"Failed to deserialize message: {e}")
     55                 continue
     56 
     57             text = ""
     58 
     59             try:
     60                 for i in range(0, len(file["data"]), CHUNK_SIZE):
     61                     audio_file = io.BytesIO(file["data"][i : i + CHUNK_SIZE])
     62                     audio_file.name = file["name"]
     63 
     64                     transcription = await openAIClient.audio.transcriptions.create(
     65                         model="whisper-1", file=audio_file
     66                     )
     67                     text += transcription.text
     68             except Exception as e:
     69                 print(f"Transcription failed: {e}")
     70                 text = f"[ERROR: Transcription failed - {e}]"
     71 
     72             await producer.send_and_wait(
     73                 "transcription.output", key=message.key, value=str.encode(text)
     74             )
     75 
     76     finally:
     77         if consumer:
     78             await consumer.stop()
     79         if producer:
     80             await producer.stop()
     81         print("Transcriber shut down")
     82 
     83 
     84 if __name__ == "__main__":
     85     asyncio.run(transcribe())