add summary listener

summary/main.py

@@ -0,0 +1,67 @@
+import sqlite3
+import json
+import os
+import paho.mqtt.client as mqtt
+
+DB_FILE = os.getenv('DB_FILE', 'agent-summary.db3')
+MQTT_HOST = os.getenv('MQTT_HOST', 'iris.dgtlu.net')
+MQTT_TOPIC = os.getenv('MQTT_TOPIC','agent/summary')
+MQTT_USER = os.getenv('MQTT_USER','agent')
+MQTT_PASS = os.getenv('MQTT_PASS','agent')
+
+def db_init():
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    c.execute('''CREATE TABLE IF NOT EXISTS AgentSummary (
+camera_id int,
+dt datetime default current_timestamp,
+file char(500),
+model char(100),
+summary text,
+primary key (dt, file)
+)''')
+    conn.commit()
+    return conn
+
+
+def db_save_input(conn: sqlite3.Connection, json_string):
+    dt = json.loads(json_string)
+    c = conn.cursor()
+    c.execute('INSERT INTO AgentSummary (camera_id, file, model, summary) VALUES (?, ?, ?, ?)',
+              (dt["id"], dt["file"], dt["data"]["model"], dt["data"]["choices"][0]["message"]["content"].strip()))
+    conn.commit()
+
+
+def mqtt_init() -> mqtt:
+    def on_connect(mqt_client, userdata, flags, rc, properties):
+        if rc == 0:
+            print("Connected to MQTT Broker!")
+        else:
+            print("Failed to connect, return code %d\n", rc)
+
+    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
+    client.on_connect = on_connect
+    client.username_pw_set(username=MQTT_USER, password=MQTT_PASS)
+    client.connect(host=MQTT_HOST, keepalive=60)
+    return client
+
+
+def mqtt_subscribe(client: mqtt, db_conn: sqlite3.Connection):
+    def on_message(mqt_client, userdata, msg):
+        print(f"Received `{msg.payload.decode()}` from `{msg.topic}` topic")
+        db_save_input(db_conn, msg.payload.decode())
+
+    client.subscribe(MQTT_TOPIC)
+    client.on_message = on_message
+
+
+def main():
+    db_conn = db_init()
+    mqtt_client = mqtt_init()
+    mqtt_subscribe(mqtt_client, db_conn)
+    mqtt_client.loop_forever()
+    db_conn.close()
+
+
+if __name__ == '__main__':
+    main()