Beta/core/llm.py

"""Ollama integration for Superviseur bot - replaces llama-cpp-python for better stability."""

import json
import logging
import asyncio
import aiohttp
import re
import os
from typing import Optional, Dict, Any, List

logger = logging.getLogger('Superviseur')

class LLMManager:
    """Manages LLM interactions via Ollama API."""
    
    def __init__(
        self,
        model_name: str = "gpt-oss:20b",
        base_url: str = None,
        temperature: float = 0.5
    ):
        self.model_name = model_name
        self.base_url = base_url or os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
        self.temperature = temperature
    
    def build_payload(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        vision_model: bool = False,
        attachments: Optional[list] = None,
        force_json: bool = True
    ) -> Dict[str, Any]:
        """Build the payload for Ollama Generate API."""
        return {
            "model": self.model_name,
            "prompt": prompt,
            "system": system_prompt or "Tu es une IA serviable.",
            "stream": True,  # Streaming pour voir la réponse en direct dans le terminal
            "format": "json" if force_json else "",
            "options": {
                "temperature": self.temperature,
                "num_ctx": 16384,
                "repeat_penalty": 1.2,
                "num_predict": 4096,
                "top_p": 0.9,
                "top_k": 40
            }
        }
    
    async def call_llama(self, payload: Dict[str, Any]) -> str:
        """Call Ollama Generate API and stream response."""
        try:
            # Automatic payload conversion for legacy compatibility (e.g. from moderation.py)
            ollama_payload = {
                "model": payload.get("model", self.model_name),
                "prompt": payload["prompt"],
                "system": payload.get("system", payload.get("system_prompt", "Tu es une IA serviable.")),
                "stream": True,  # Streaming pour voir la réponse en direct
                "think": False,  # Désactive le reasoning GPT-OSS
                "format": "",
                "options": {
                    "temperature": payload.get("temperature", self.temperature),
                    "num_ctx": payload.get("n_ctx", 16384),
                    "num_predict": payload.get("max_tokens") or payload.get("num_predict", 8192),
                    "repeat_penalty": payload.get("repeat_penalty") or payload.get("repeat_last_n", 1.2),
                    "stop": payload.get("stop", [])
                }
            }

            accumulated_reply = ""
            api_url = f"{self.base_url}/api/generate"
            
            # Build headers with optional API key
            headers = {}
            api_key = os.environ.get("OLLAMA_API_KEY", "")
            if api_key:
                headers["Authorization"] = f"Bearer {api_key}"

            async with aiohttp.ClientSession() as session:
                async with session.post(api_url, json=ollama_payload, headers=headers) as response:
                    if response.status != 200:
                        err_text = await response.text()
                        logger.error(f"Ollama API Error ({response.status}): {err_text}")
                        return ""

                    # Streaming en direct dans le terminal
                    async for line in response.content:
                        if line:
                            try:
                                chunk = json.loads(line)
                                resp = chunk.get("response", "")
                                if resp:
                                    accumulated_reply += resp
                                    print(resp, end="", flush=True)
                                if chunk.get("done"):
                                    break
                            except json.JSONDecodeError:
                                continue
                    print()
                    return accumulated_reply.strip()
                
        except Exception as e:
            logger.error(f"Error calling Ollama: {e}")
            return ""

    async def close_session(self):
        pass

    def extract_json_actions(self, text: str) -> Optional[str]:
        """
        Extract JSON action objects from text with surgical precision.
        Handles both single objects {...} and arrays [...].
        Also handles // comments that GPT-OSS sometimes inserts.
        """
        if not text:
            return None

        text = text.strip()

        # Remove // comments (GPT-OSS sometimes adds them)
        text = re.sub(r'//.*?\n', '\n', text)
        text = re.sub(r'//.*?$', '', text, flags=re.MULTILINE)

        # 1. Try to find a JSON array [...]
        first_bracket = text.find('[')
        if first_bracket != -1:
            last_bracket = text.rfind(']')
            if last_bracket > first_bracket:
                array_str = text[first_bracket:last_bracket+1]
                try:
                    json.loads(array_str)
                    return array_str
                except json.JSONDecodeError:
                    pass

        # 2. Try to find a JSON object {...}
        first_brace = text.find('{')
        if first_brace == -1:
            return None
            
        last_brace = text.rfind('}')
        if last_brace == -1 or last_brace < first_brace:
            return None

        stack = []
        for i in range(first_brace, len(text)):
            if text[i] == '{':
                stack.append('{')
            elif text[i] == '}':
                if stack:
                    stack.pop()
                    if not stack:
                        match = text[first_brace:i+1]
                        if '"action"' in match.lower():
                            return match

        return None