Merge pull request #30 from DefangLabs/linda-segment-api

Add analytics tracking (Segment API) to server

Author: raphaeltm

app/app.py

@@ -1,9 +1,13 @@
-from flask import Flask, request, jsonify, render_template, Response, stream_with_context
+from flask import Flask, request, jsonify, render_template, Response, stream_with_context, session
 from flask_wtf.csrf import CSRFProtect
 from rag_system import rag_system
 import hashlib
 import subprocess
 import os
+import segment.analytics as analytics
+import uuid
+
+analytics.write_key = os.getenv('SEGMENT_WRITE_KEY')
 
 app = Flask(__name__, static_folder='templates/images')
 app.config['SECRET_KEY'] = os.getenv('SECRET_KEY')
@@ -33,17 +37,35 @@ def ask():
 
     data = request.get_json()
     query = data.get('query')
+
     if not query:
         return jsonify({"error": "No query provided"}), 400
+    
+    # For analytics tracking, generates an anonymous id and uses it for the session
+    if 'anonymous_id' not in session:
+        session['anonymous_id'] = str(uuid.uuid4())
+    anonymous_id = session['anonymous_id']
 
     def generate():
+        full_response = ""
         try:
             for token in rag_system.answer_query_stream(query):
                 yield token
+                full_response += token
         except Exception as e:
             print(f"Error in /ask endpoint: {e}")
             yield "Internal Server Error"
 
+        if not full_response:
+            full_response = "No response generated"
+
+        # Track the query and response 
+        analytics.track(
+            anonymous_id=anonymous_id, 
+            event='Chatbot Question submitted', 
+            properties={'query': query, 'response': full_response, 'source': 'Ask Defang'}
+        )
+
     return Response(stream_with_context(generate()), content_type='text/markdown')
 
 @app.route('/trigger-rebuild', methods=['POST'])
@@ -75,4 +97,4 @@ def trigger_rebuild():
         return jsonify({"error": "Internal Server Error"}), 500
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+    app.run(host='0.0.0.0', port=5000)

app/requirements.txt

@@ -2,6 +2,7 @@ Flask==2.0.1
 Flask-WTF==1.2.2
 Werkzeug==2.0.3
 scikit-learn==0.24.2
+segment-analytics-python==2.3.3
 numpy==1.22.0
 sentence-transformers==2.1.0
 torch==1.10.0

compose.dev.yaml

@@ -11,6 +11,7 @@ services:
     environment:
       FLASK_APP: app.py
       SECRET_KEY: supersecret
+      SEGMENT_WRITE_KEY: ${SEGMENT_WRITE_KEY} # Set your Segment write key here or in the .env file
       SESSION_COOKIE_SECURE: 0
       OPENAI_API_KEY: ${OPENAI_API_KEY} # Set your OpenAI API key here or in the .env file
     command: flask run --host=0.0.0.0

compose.yaml

@@ -14,6 +14,7 @@ services:
     environment:
       FLASK_APP: app.py
       SECRET_KEY: 
+      SEGMENT_WRITE_KEY: 
       SESSION_COOKIE_SECURE: 1
       OPENAI_API_KEY: ${OPENAI_API_KEY} # Set your OpenAI API key here or in the .env file
     command: uwsgi --http 0.0.0.0:5000 --wsgi-file app.py --callable app --processes 4 --threads 2