core-crew/local-agent.py
2024-06-18 22:54:36 -04:00

84 lines
2.9 KiB
Python

from flask import Flask, jsonify, render_template, request
from flask_socketio import SocketIO, emit
from pymongo import MongoClient
from bson.objectid import ObjectId
import threading
from tasks import generate_content
import html
app = Flask(__name__)
# MongoDB Atlas connection
client = MongoClient("mongodb+srv://maheshkommareddi:Yu2L6pQKyJgcTb9a@cluster0.qadl40g.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0")
db = client.content_generation
app.config['SECRET_KEY'] = 'secret_key'
socketio = SocketIO(app)
# Periodically fetch and emit task updates
def fetch_and_emit_updates():
while True:
updates = db.task_updates.find({})
for update in updates:
socketio.emit('agent_update', {
'task_id': update['task_id'],
'status': update['status'],
'message': update['message']
})
db.task_updates.delete_one({"_id": update["_id"]}) # Remove the update after emitting
socketio.sleep(5)
# Start the periodic update thread
threading.Thread(target=fetch_and_emit_updates).start()
# Route for the chat interface
@app.route('/chat')
def chat():
return render_template('chat.html')
# SocketIO event handler for new content generation requests
@socketio.on('generate_content')
def handle_generate_content(data):
agenda = data['agenda']
task = generate_content.apply_async(args=[agenda])
task_id = task.id
emit('task_queued', {'task_id': task_id, 'agenda': agenda})
@app.route('/', methods=['GET', 'POST'])
def index():
# Fetch all existing blog entries from MongoDB Atlas
blog_entries = list(db.content.find({}, {"_id": 1, "agenda": 1, "blog_post": 1}))
if request.method == 'POST':
agenda = request.form.get('agenda', 'quantum gravitational sensing for underground detection')
# Queue the task for background processing
task = generate_content.apply_async(args=[agenda])
# Save the task ID to MongoDB Atlas
content = {
"agenda": agenda,
"task_id": str(task.id),
"status": "PENDING"
}
db.content.insert_one(content)
return render_template('processing.html', agenda=agenda, task_id=task.id)
return render_template('index.html', blog_entries=blog_entries)
@app.route('/blog/<id>', methods=['GET'])
def blog_post(id):
# Fetch the blog entry from MongoDB Atlas
blog_entry = db.content.find_one({"_id": ObjectId(id)})
if blog_entry:
if blog_entry.get('status') == 'COMPLETED':
return render_template('result.html', report=html.unescape(blog_entry['report']), blog_post=blog_entry['blog_post'])
else:
return render_template('result.html', report=html.unescape(blog_entry['report']), blog_post=blog_entry['blog_post'])
else:
return "Blog entry not found", 404
if __name__ == '__main__':
socketio.run(app, debug=True, host="0.0.0.0", port=5001)