Beta/superviseur/llm.py

"""LLM integration for Superviseur bot."""

import json
import aiohttp
import logging
import re
from typing import Optional, Dict, Any

logger = logging.getLogger('Superviseur')


class LLMManager:
    """Manages LLM interactions with Ollama."""
    
    def __init__(
        self,
        ollama_api_url: str,
        ollama_model: str,
        ollama_timeout: int = 60,
        ollama_temperature: float = 0.7
    ):
        self.ollama_api_url = ollama_api_url
        self.ollama_model = ollama_model
        self.ollama_timeout = ollama_timeout
        self.ollama_temperature = ollama_temperature
        self.http_session: Optional[aiohttp.ClientSession] = None
    
    async def get_session(self) -> aiohttp.ClientSession:
        """Get or create HTTP session."""
        if not self.http_session or self.http_session.closed:
            self.http_session = aiohttp.ClientSession()
        return self.http_session
    
    async def close_session(self) -> None:
        """Close HTTP session."""
        if self.http_session and not self.http_session.closed:
            await self.http_session.close()
    
    def build_payload(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        structured_history: Optional[list] = None,
        vision_model: bool = False,
        attachments: Optional[list] = None
    ) -> Dict[str, Any]:
        """Build the API payload for Ollama or OpenAI Chat API (llama.cpp)."""
        is_openai = "/v1/chat/completions" in self.ollama_api_url
        
        if is_openai:
            # OpenAI Chat API format
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            
            if structured_history:
                messages.extend(structured_history)
                
            messages.append({"role": "user", "content": prompt})
            
            payload = {
                "model": self.ollama_model,
                "messages": messages,
                "stream": True,
                "temperature": self.ollama_temperature,
                "max_tokens": 2048
            }
        else:
            # Ollama /api/generate payload (Fallback)
            full_prompt = prompt
            if structured_history:
                hist_str = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in structured_history])
                full_prompt = f"Historique:\n{hist_str}\n\nUser: {prompt}"
                
            payload = {
                "model": self.ollama_model,
                "prompt": full_prompt,
                "stream": True,
                "options": {"temperature": self.ollama_temperature}
            }
            if system_prompt:
                payload["system"] = system_prompt
        
        # Vision support (simplified)
        if vision_model and attachments:
            for attachment in attachments:
                if attachment.content_type and attachment.content_type.startswith('image/'):
                    # actual image handling would require download and base64 for OpenAI
                    pass
        
        return payload
    
    async def call_ollama(self, payload: Dict[str, Any]) -> str:
        """Call LLM API and accumulate response."""
        session = await self.get_session()
        is_openai = "/v1/chat/completions" in self.ollama_api_url

        async with session.post(
            self.ollama_api_url,
            json=payload,
            timeout=aiohttp.ClientTimeout(total=self.ollama_timeout)
        ) as response:
            response.raise_for_status()
            
            accumulated_reply = ""
            accumulated_thoughts = ""
            async for line_bytes in response.content:
                line = line_bytes.decode('utf-8').strip()
                if not line:
                    continue
                
                # Strip SSE prefix if present
                if line.startswith("data: "):
                    line = line[6:]
                elif line.startswith("event:"):
                    continue
                
                if line == "[DONE]":
                    break
                    
                try:
                    data = json.loads(line)
                    
                    # 1. Handle Thinking (OpenAI style reasoning or special field)
                    thought = None
                    if 'thinking' in data: # Custom llama.cpp field
                        thought = data['thinking']
                    elif 'choices' in data and len(data['choices']) > 0:
                        delta = data['choices'][0].get('delta', {})
                        thought = delta.get('reasoning_content')
                    
                    if thought:
                        accumulated_thoughts += thought
                        print(f"🤔 Le superviseur réfléchit : {thought}", flush=True)

                    # 2. Handle Content
                    chunk = ""
                    if 'response' in data: # Ollama
                        chunk = data['response']
                    elif 'choices' in data and len(data['choices']) > 0: # OpenAI
                        delta = data['choices'][0].get('delta', {})
                        if delta.get('content') is not None:
                            chunk = delta.get('content', "")
                    elif 'content' in data: # llama.cpp legacy
                        chunk = data['content']
                    
                    if chunk:
                        accumulated_reply += chunk
                        if chunk.strip():
                            print(f"🤖 Réponse du Superviseur (fragment) : {chunk}", flush=True)
                    
                    elif 'action' in data: # Action already in JSON (unlikely for chat API)
                        accumulated_reply = line
                        break

                except json.JSONDecodeError:
                    # In case of malformed line or raw text
                    if not line.startswith("{"):
                         accumulated_reply += line

            # Fallback for models that put the final JSON response inside the thought block
            if not accumulated_reply.strip() and accumulated_thoughts.strip():
                # Test if the thought block contains our strict format using robust search
                if re.search(r'\{[\s\n]*[\"\']action[\"\']', accumulated_thoughts, re.IGNORECASE):
                    print("⚠️ Aucun texte de contenu final mais du JSON détecté dans les pensées. Récupération.", flush=True)
                    accumulated_reply = accumulated_thoughts

            return accumulated_reply
    
    def extract_json_actions(self, text: str) -> Optional[str]:
        """Extract JSON action objects from text more robustly using an O(N) brace matcher."""
        
        # 1. First, try to extract from markdown blocks if present
        markdown_blocks = re.findall(r'```(?:json)?\s*(.*?)\s*```', text, re.DOTALL)
        candidates = markdown_blocks if markdown_blocks else [text]
        
        actions = []
        for candidate_text in candidates:
            # Try to parse the whole text as a list or dict first
            try:
                parsed = json.loads(candidate_text.strip())
                if isinstance(parsed, list):
                    for item in parsed:
                        if isinstance(item, dict) and 'action' in item:
                            actions.append(item)
                    if actions: return json.dumps(actions)
                elif isinstance(parsed, dict) and 'action' in parsed:
                    actions.append(parsed)
                    return json.dumps(actions)
            except:
                pass

            # Fallback: O(N) scan for top-level { ... } blocks
            brace_count = 0
            current_block = []
            
            for char in candidate_text:
                if char == '{':
                    if brace_count < 0:
                        brace_count = 0
                        current_block = []
                    brace_count += 1
                    
                if brace_count > 0:
                    current_block.append(char)
                    
                if char == '}':
                    if brace_count > 0:
                        brace_count -= 1
                        if brace_count == 0 and current_block:
                            json_str = "".join(current_block)
                            current_block = []
                            try:
                                parsed = json.loads(json_str)
                                if isinstance(parsed, dict) and 'action' in parsed:
                                    if parsed not in actions: # avoid duplicates
                                        actions.append(parsed)
                            except:
                                # Fallback with single quote replace
                                try:
                                    parsed = json.loads(json_str.replace("'", '"'))
                                    if isinstance(parsed, dict) and 'action' in parsed:
                                        if parsed not in actions:
                                            actions.append(parsed)
                                except:
                                    pass
                    else:
                        brace_count = 0
                        current_block = []

        if actions:
            return json.dumps(actions)
        return None