summarizer

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

main.py (2002B)

      1 import asyncio
      2 import os
      3 import signal
      4 
      5 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
      6 from openai import AsyncOpenAI
      7 
      8 openAIClient = AsyncOpenAI()
      9 
     10 kafkaAddress = os.environ.get("KAFKA_ADDRESS", "localhost:9092")
     11 
     12 running = True
     13 
     14 
     15 def signal_handler(signum, frame):
     16     global running
     17     print(f"Received signal {signum}, shutting down...")
     18     running = False
     19 
     20 
     21 async def summarizer():
     22     global running
     23 
     24     signal.signal(signal.SIGTERM, signal_handler)
     25     signal.signal(signal.SIGINT, signal_handler)
     26 
     27     producer = None
     28     consumer = None
     29 
     30     try:
     31         producer = AIOKafkaProducer(bootstrap_servers=kafkaAddress)
     32         await producer.start()
     33 
     34         consumer = AIOKafkaConsumer(
     35             "summarization.input", bootstrap_servers=kafkaAddress
     36         )
     37         await consumer.start()
     38 
     39         while running:
     40             try:
     41                 message = await asyncio.wait_for(consumer.getone(), timeout=1.0)
     42             except asyncio.TimeoutError:
     43                 continue
     44 
     45             text = message.value.decode()
     46 
     47             try:
     48                 summarization = await openAIClient.chat.completions.create(
     49                     messages=[
     50                         {
     51                             "role": "user",
     52                             "content": "summarize the following:\n" + text,
     53                         }
     54                     ],
     55                     model="gpt-4-turbo",
     56                 )
     57                 summary = summarization.choices[0].message.content
     58             except Exception as e:
     59                 print(f"Summarization failed: {e}")
     60                 summary = f"[ERROR: Summarization failed - {e}]"
     61 
     62             await producer.send_and_wait(
     63                 "summarization.output", key=message.key, value=str.encode(summary)
     64             )
     65 
     66     finally:
     67         if consumer:
     68             await consumer.stop()
     69         if producer:
     70             await producer.stop()
     71         print("Summarizer shut down")
     72 
     73 
     74 if __name__ == "__main__":
     75     asyncio.run(summarizer())