llm_lab_perso/api_server.py

#!/usr/bin/env python3
"""
Serveur API pour intégrer LLM Lab avec Cursor et Obsidian
"""
from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import json
import os
import logging
import time
import sys

# Ajouter le répertoire courant au chemin de recherche Python
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Importer les modules LLM Lab
from utils.agent_manager import AgentManager

# Configuration du logging
os.makedirs("logs", exist_ok=True)
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("logs/api_server.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("api_server")

# Initialisation de l'application Flask
app = Flask(__name__)
CORS(app)  # Permet les requêtes cross-origin

@app.route('/v1/chat/completions', methods=['POST'])
def chat_completion():
    """
    Endpoint compatible avec l'API OpenAI Chat pour Cursor
    """
    try:
        # Vérifier que la requête contient du JSON valide
        if not request.is_json:
            return jsonify({"error": "La requête doit contenir du JSON valide"}), 400
            
        data = request.json or {}  # Utiliser un dictionnaire vide par défaut si None
        logger.info(f"Requête reçue: {json.dumps(data)}")
        
        # Extraire les messages et les paramètres
        messages = data.get('messages', [])
        model = data.get('model', 'codellama:13b-python')
        temperature = data.get('temperature', 0.7)
        
        # Construire le prompt à partir des messages
        system_message = next((msg['content'] for msg in messages if msg['role'] == 'system'), None)
        user_messages = [msg['content'] for msg in messages if msg['role'] == 'user']
        
        # Utiliser le dernier message utilisateur comme prompt
        prompt = user_messages[-1] if user_messages else ""
        
        # Détecter le type de tâche pour choisir l'agent approprié
        agent_name = "cursor"  # Par défaut
        
        # Logique de sélection d'agent en fonction du contenu
        if "obsidian" in prompt.lower() or "markdown" in prompt.lower() or "note" in prompt.lower():
            agent_name = "obsidian"
        elif "javascript" in prompt.lower() or "js" in prompt.lower() or "html" in prompt.lower() or "css" in prompt.lower():
            agent_name = "webdev"
        elif "python" in prompt.lower():
            agent_name = "python"
        
        logger.info(f"Agent sélectionné: {agent_name}")
        
        # Créer et configurer l'agent
        agent = AgentManager.create(agent_name)
        
        # Remplacer le system prompt si fourni
        if system_message:
            agent.system_prompt = system_message
        
        # Ajuster les paramètres
        agent.params["temperature"] = temperature
        
        # Générer la réponse
        start_time = time.time()
        response = agent.generate(prompt)
        end_time = time.time()
        
        generation_time = end_time - start_time
        logger.info(f"Réponse générée pour l'agent {agent_name} en {generation_time:.2f} secondes")
        
        # Formatage compatible avec l'API OpenAI
        return jsonify({
            "id": f"llmlab-{agent_name}-{hash(prompt) % 10000}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": agent.model,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response
                    },
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "prompt_tokens": len(prompt.split()),
                "completion_tokens": len(response.split()),
                "total_tokens": len(prompt.split()) + len(response.split())
            }
        })
    
    except Exception as e:
        logger.error(f"Erreur: {str(e)}", exc_info=True)
        return jsonify({
            "error": {
                "message": str(e),
                "type": "server_error",
                "code": 500
            }
        }), 500

@app.route('/v1/models', methods=['GET'])
def list_models():
    """
    Liste les modèles disponibles (compatible OpenAI)
    """
    agents = AgentManager.list_agents()
    models = []
    
    for agent_name, info in agents.items():
        models.append({
            "id": info['model'],
            "object": "model",
            "created": int(time.time()),
            "owned_by": "llmlab",
            "permission": [{"id": agent_name, "object": "model_permission"}],
            "root": info['model'],
            "parent": None
        })
    
    return jsonify({
        "object": "list",
        "data": models
    })

@app.route('/health', methods=['GET'])
def health_check():
    """
    Endpoint de vérification de l'état du serveur
    """
    return jsonify({
        "status": "healthy",
        "version": "1.0.0",
        "timestamp": int(time.time())
    })

@app.route('/agents', methods=['GET'])
def list_agents():
    """
    Liste les agents disponibles (endpoint personnalisé)
    """
    agents = AgentManager.list_agents()
    return jsonify({
        "agents": [
            {
                "name": name,
                "model": info['model'],
                "description": info['description']
            }
            for name, info in agents.items()
        ]
    })

@app.route('/generate', methods=['POST'])
def generate():
    """
    Endpoint simplifié pour les applications personnalisées
    """
    try:
        # Vérifier que la requête contient du JSON valide
        if not request.is_json:
            return jsonify({"error": "La requête doit contenir du JSON valide"}), 400
            
        data = request.json or {}  # Utiliser un dictionnaire vide par défaut si None
        prompt = data.get('prompt', '')
        agent_name = data.get('agent', 'cursor')
        
        # Paramètres optionnels
        system_prompt = data.get('system_prompt', None)
        temperature = data.get('temperature', None)
        
        # Créer l'agent
        agent = AgentManager.create(agent_name)
        
        # Appliquer les paramètres personnalisés si fournis
        if system_prompt:
            agent.system_prompt = system_prompt
            
        if temperature is not None:
            agent.params["temperature"] = temperature
            
        # Générer la réponse
        start_time = time.time()
        response = agent.generate(prompt)
        generation_time = time.time() - start_time
        
        return jsonify({
            "response": response,
            "agent": agent_name,
            "model": agent.model,
            "generation_time": generation_time
        })
        
    except Exception as e:
        logger.error(f"Erreur: {str(e)}", exc_info=True)
        return jsonify({
            "error": str(e)
        }), 500

if __name__ == '__main__':
    print("=== Serveur API LLM Lab pour Cursor et Obsidian ===")
    print("Serveur démarré sur http://localhost:8000")
    print()
    print("Endpoints disponibles:")
    print("  - http://localhost:8000/v1/chat/completions (compatible OpenAI)")
    print("  - http://localhost:8000/v1/models (compatible OpenAI)")
    print("  - http://localhost:8000/generate (API simplifiée)")
    print("  - http://localhost:8000/agents (liste d'agents)")
    print("  - http://localhost:8000/health (statut)")
    print()
    print("Pour Cursor:")
    print("  1. Ouvrez Cursor")
    print("  2. Allez dans Settings > AI")
    print("  3. Sélectionnez 'Custom endpoint'")
    print("  4. Entrez l'URL: http://localhost:8000/v1")
    print()
    print("Agents disponibles:")
    try:
        for agent_name, info in AgentManager.list_agents().items():
            print(f"  - {agent_name}: {info['description']} ({info['model']})")
    except Exception as e:
        print(f"Erreur lors de la liste des agents: {str(e)}")
        print("Assurez-vous que les modules LLM Lab sont correctement installés.")
    print()
    print("Logs: logs/api_server.log")
    print("Appuyez sur Ctrl+C pour arrêter le serveur")
    
    # Démarrer le serveur
    app.run(host='0.0.0.0', port=8000, debug=False)