app.py

from flask import Flask, render_template, request, jsonify, session
import google.generativeai as genai
import os
from dotenv import load_dotenv
import http.client
import json
from werkzeug.utils import secure_filename
import tempfile
from datetime import datetime

app = Flask(__name__)
app.secret_key = os.urandom(24)
load_dotenv()

# Configure Google AI
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Safety settings
safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# System prompt
SYSTEM_PROMPT = """
# Prompt System pour Mariam, IA conçu par youssouf 
[Your existing system prompt content here]
"""

# Initialize Gemini model
model = genai.GenerativeModel('gemini-2.0-flash-exp', 
                            tools='code_execution',
                            safety_settings=safety_settings,
                            system_instruction=SYSTEM_PROMPT)

# Store chat sessions with user isolation
chat_sessions = {}
chat_histories = {}

def get_client_ip():
    """Get a unique identifier for the client"""
    if request.headers.getlist("X-Forwarded-For"):
        ip = request.headers.getlist("X-Forwarded-For")[0]
    else:
        ip = request.remote_addr
    return ip

def perform_web_search(query):
    conn = http.client.HTTPSConnection("google.serper.dev")
    payload = json.dumps({"q": query})
    headers = {
        'X-API-KEY': '9b90a274d9e704ff5b21c0367f9ae1161779b573',
        'Content-Type': 'application/json'
    }
    try:
        conn.request("POST", "/search", payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode("utf-8"))
        return data
    except Exception as e:
        print(f"Web search error: {e}")
        return None
    finally:
        conn.close()

def format_search_results(data):
    if not data:
        return "Aucun résultat trouvé"
    
    result = []
    
    if 'knowledgeGraph' in data:
        kg = data['knowledgeGraph']
        result.append({
            'type': 'knowledge',
            'title': kg.get('title', ''),
            'description': kg.get('description', ''),
            'category': kg.get('type', '')
        })
    
    if 'organic' in data:
        for item in data['organic'][:3]:
            result.append({
                'type': 'organic',
                'title': item['title'],
                'snippet': item['snippet'],
                'link': item['link']
            })
    
    return result

UPLOAD_FOLDER = 'temp'
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'pdf', 'txt', 'mp3', 'mp4'}

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def home():
    client_id = get_client_ip()
    if client_id not in chat_histories:
        chat_histories[client_id] = []
    return render_template('index.html', messages=chat_histories[client_id])

@app.route('/send_message', methods=['POST'])
def send_message():
    try:
        client_id = get_client_ip()
        data = request.json
        message = data.get('message')
        web_search_enabled = data.get('web_search', False)
        
        if not message:
            return jsonify({'error': 'No message provided'}), 400

        # Initialize or get existing chat session for this client
        if client_id not in chat_sessions:
            chat_sessions[client_id] = model.start_chat(history=[])
            chat_histories[client_id] = []

        # Perform web search if enabled
        if web_search_enabled:
            web_results = perform_web_search(message)
            if web_results:
                formatted_results = format_search_results(web_results)
                message = f"""Question: {message}\n\nRésultats de recherche web:\n{formatted_results}\n\nPourrais-tu analyser ces informations et me donner une réponse complète?"""

        # Send message to Gemini
        response = chat_sessions[client_id].send_message(message)
        
        # Update message history for this client
        chat_histories[client_id].append({
            'role': 'user',
            'content': message,
            'timestamp': datetime.now().isoformat()
        })
        chat_histories[client_id].append({
            'role': 'assistant',
            'content': response.text,
            'timestamp': datetime.now().isoformat()
        })
        
        return jsonify({
            'response': response.text
        })

    except Exception as e:
        print(f"Error in send_message: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        try:
            gemini_file = genai.upload_file(filepath)
            return jsonify({'success': True, 'filename': filename})
        except Exception as e:
            return jsonify({'error': str(e)}), 500
    
    return jsonify({'error': 'Invalid file type'}), 400

@app.route('/clear_chat', methods=['POST'])
def clear_chat():
    client_id = get_client_ip()
    if client_id in chat_sessions:
        del chat_sessions[client_id]
    if client_id in chat_histories:
        del chat_histories[client_id]
    return jsonify({'success': True})

if __name__ == '__main__':
    app.run()