my-ia/backend/main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import ollama
import asyncio
import json
import uvicorn
from datetime import datetime
import subprocess
import tempfile
import os

app = FastAPI(title="AI Code Assistant API")

# CORS pour permettre le frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Models
class Message(BaseModel):
    role: str
    content: str
    timestamp: Optional[str] = None

class ChatRequest(BaseModel):
    messages: List[Message]
    model: str = "code-expert"
    stream: bool = True

class CodeExecutionRequest(BaseModel):
    code: str
    language: str = "python"

class ConversationHistory:
    def __init__(self):
        self.conversations: Dict[str, List[Message]] = {}
    
    def add_message(self, conversation_id: str, message: Message):
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = []
        self.conversations[conversation_id].append(message)
    
    def get_conversation(self, conversation_id: str) -> List[Message]:
        return self.conversations.get(conversation_id, [])

history = ConversationHistory()

# WebSocket manager
class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

manager = ConnectionManager()

# Routes
@app.get("/")
async def root():
    return {"message": "AI Code Assistant API", "version": "1.0"}

@app.get("/models")
async def list_models():
    """Liste tous les modèles Ollama disponibles."""
    try:
        models = ollama.list()
        return {
            "models": [
                {
                    "name": m['name'],
                    "size": m.get('size', 0),
                    "modified": m.get('modified_at', '')
                }
                for m in models['models']
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat")
async def chat(request: ChatRequest):
    """Endpoint de chat standard (non-streaming ou streaming)."""
    try:
        messages = [{"role": m.role, "content": m.content} for m in request.messages]
        
        if request.stream:
            async def generate():
                stream = ollama.chat(
                    model=request.model,
                    messages=messages,
                    stream=True
                )
                
                for chunk in stream:
                    if 'message' in chunk and 'content' in chunk['message']:
                        content = chunk['message']['content']
                        yield f"data: {json.dumps({'content': content})}\n\n"
                
                yield "data: [DONE]\n\n"
            
            return StreamingResponse(generate(), media_type="text/event-stream")
        else:
            response = ollama.chat(
                model=request.model,
                messages=messages
            )
            return {"response": response['message']['content']}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    """WebSocket pour chat en temps réel."""
    await manager.connect(websocket)
    conversation_id = str(datetime.now().timestamp())
    
    try:
        while True:
            # Recevoir le message du client
            data = await websocket.receive_text()
            request_data = json.loads(data)
            
            user_message = Message(
                role="user",
                content=request_data['message'],
                timestamp=datetime.now().isoformat()
            )
            history.add_message(conversation_id, user_message)
            
            # Récupérer l'historique de conversation
            conversation = history.get_conversation(conversation_id)
            messages = [{"role": m.role, "content": m.content} for m in conversation]
            
            # Envoyer un accusé de réception
            await manager.send_message(
                json.dumps({
                    "type": "status",
                    "content": "Processing..."
                }),
                websocket
            )
            
            # Générer la réponse avec streaming
            model = request_data.get('model', 'code-expert')
            full_response = ""
            
            stream = ollama.chat(
                model=model,
                messages=messages,
                stream=True
            )
            
            for chunk in stream:
                if 'message' in chunk and 'content' in chunk['message']:
                    content = chunk['message']['content']
                    full_response += content
                    
                    await manager.send_message(
                        json.dumps({
                            "type": "stream",
                            "content": content
                        }),
                        websocket
                    )
            
            # Sauvegarder la réponse complète
            assistant_message = Message(
                role="assistant",
                content=full_response,
                timestamp=datetime.now().isoformat()
            )
            history.add_message(conversation_id, assistant_message)
            
            # Signal de fin
            await manager.send_message(
                json.dumps({
                    "type": "done",
                    "content": full_response
                }),
                websocket
            )
    
    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        await manager.send_message(
            json.dumps({
                "type": "error",
                "content": str(e)
            }),
            websocket
        )
        manager.disconnect(websocket)

@app.post("/execute")
async def execute_code(request: CodeExecutionRequest):
    """Exécute du code de manière sécurisée."""
    try:
        # Créer un fichier temporaire
        with tempfile.NamedTemporaryFile(
            mode='w',
            suffix=f'.{request.language}',
            delete=False
        ) as f:
            f.write(request.code)
            temp_file = f.name
        
        # Exécuter selon le langage
        if request.language == "python":
            result = subprocess.run(
                ['python3', temp_file],
                capture_output=True,
                text=True,
                timeout=10
            )
        elif request.language == "javascript":
            result = subprocess.run(
                ['node', temp_file],
                capture_output=True,
                text=True,
                timeout=10
            )
        elif request.language == "bash":
            result = subprocess.run(
                ['bash', temp_file],
                capture_output=True,
                text=True,
                timeout=10
            )
        else:
            raise HTTPException(status_code=400, detail=f"Language {request.language} not supported")
        
        # Nettoyer
        os.unlink(temp_file)
        
        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "returncode": result.returncode,
            "success": result.returncode == 0
        }
    
    except subprocess.TimeoutExpired:
        os.unlink(temp_file)
        raise HTTPException(status_code=408, detail="Code execution timeout")
    except Exception as e:
        if os.path.exists(temp_file):
            os.unlink(temp_file)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/analyze-code")
async def analyze_code(request: CodeExecutionRequest):
    """Analyse du code avec l'IA."""
    try:
        response = ollama.chat(
            model='code-reviewer',
            messages=[{
                'role': 'user',
                'content': f'''Analyse ce code {request.language}:

```{request.language}
{request.code}
```

Fournis:
1. Une brève description de ce que fait le code
2. Les problèmes potentiels (bugs, sécurité, performance)
3. Des suggestions d'amélioration
'''
            }]
        )
        
        return {"analysis": response['message']['content']}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/fix-code")
async def fix_code(code: str, error: str, language: str = "python"):
    """Corrige du code avec une erreur."""
    try:
        response = ollama.chat(
            model='debugger',
            messages=[{
                'role': 'user',
                'content': f'''Code {language} avec erreur:

```{language}
{code}
```

Erreur:
```
{error}
```

Identifie le problème et fournis le code corrigé.
'''
            }]
        )
        
        return {"fixed_code": response['message']['content']}
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Vérifie que Ollama est accessible."""
    try:
        ollama.list()
        return {"status": "healthy", "ollama": "connected"}
    except Exception as e:
        return {"status": "unhealthy", "error": str(e)}

if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=9001,
        reload=True,
        ws_ping_interval=20,
        ws_ping_timeout=20
    )