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,
        vision_model: bool = False,
        attachments: Optional[list] = None
    ) -> Dict[str, Any]:
        """Build the API payload for Ollama."""
        payload = {
            "model": self.ollama_model,
            "prompt": 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/'):
                    if any(v in self.ollama_model.lower() for v in 
                           ['llava', 'bakllava', 'moondream', 'llama3.2-vision', 'minicpm-v']):
                        # Note: actual image handling would require download
                        pass
        
        return payload
    
    async def call_ollama(self, payload: Dict[str, Any]) -> str:
        """Call Ollama API and accumulate response."""
        session = await self.get_session()

        async with session.post(
            self.ollama_api_url,
            json=payload,
            timeout=aiohttp.ClientTimeout(total=self.ollama_timeout)
        ) as response:
            response.raise_for_status()
            
            # Si on ne streame pas, on récupère tout d'un coup
            if not payload.get("stream", True):
                data = await response.json()
                return data.get("response", "")

            accumulated_reply = ""
            async for line in response.content:
                line = line.decode('utf-8').strip()
                if not line:
                    continue
                    
                try:
                    data = json.loads(line)
                    
                    if 'thinking' in data and data['thinking']:
                        print(f"🤔 Le superviseur réfléchit : {data['thinking']}", flush=True)
                        continue

                    if 'response' in data:
                        chunk = data['response']
                        accumulated_reply += chunk
                        if chunk.strip():
                            print(f"🤖 Réponse du Superviseur (fragment) : {chunk}", flush=True)
                    
                    elif 'action' in data:
                        accumulated_reply = line
                        break

                except json.JSONDecodeError:
                    accumulated_reply += line

            return accumulated_reply
    
    def extract_json_actions(self, text: str) -> Optional[str]:
        """Extract JSON action objects from text more robustly."""
        actions = []
        
        # 1. 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]
        
        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:
                            if item not in actions:
                                actions.append(item)
                    if actions: return json.dumps(actions)
                elif isinstance(parsed, dict) and 'action' in parsed:
                    if parsed not in actions:
                        actions.append(parsed)
                    return json.dumps(actions)
            except:
                pass

            # Fallback: scan for { ... } using brace matching
            start_indexes = [i for i, c in enumerate(candidate_text) if c == '{']
            for start in start_indexes:
                brace_count = 0
                for i in range(start, len(candidate_text)):
                    if candidate_text[i] == '{': 
                        brace_count += 1
                    elif candidate_text[i] == '}': 
                        brace_count -= 1
                        if brace_count == 0:
                            json_str = candidate_text[start:i+1]
                            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
                            break # Move to next start brace

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