162 lines
6 KiB
Python
162 lines
6 KiB
Python
"""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 = []
|
|
|
|
# Regex plus laxiste pour trouver des blocs JSON potentiels
|
|
# On cherche des blocs commençant par {"action" ou {'action'
|
|
matches = re.finditer(r'\{[\s\n]*[\"\']action[\"\'].*?\}', text, re.DOTALL)
|
|
|
|
for match in matches:
|
|
json_str = match.group(0)
|
|
try:
|
|
# Tentative de parsing standard
|
|
parsed = json.loads(json_str)
|
|
if isinstance(parsed, dict) and 'action' in parsed:
|
|
actions.append(parsed)
|
|
except json.JSONDecodeError:
|
|
# Si erreur, on tente de remplacer les simples quotes par des doubles
|
|
# (Cas fréquent si le LLM mélange les styles)
|
|
try:
|
|
# Remplacement très basique (peut échouer si contenu complexe)
|
|
json_corrected = json_str.replace("'", '"')
|
|
parsed = json.loads(json_corrected)
|
|
if isinstance(parsed, dict) and 'action' in parsed:
|
|
actions.append(parsed)
|
|
continue
|
|
except: pass
|
|
|
|
# Méthode récursive des accolades si le bloc est tronqué ou complexe
|
|
start = match.start()
|
|
brace_count = 0
|
|
for i in range(start, len(text)):
|
|
if text[i] == '{': brace_count += 1
|
|
elif text[i] == '}':
|
|
brace_count -= 1
|
|
if brace_count == 0:
|
|
try:
|
|
candidate = text[start:i+1]
|
|
try:
|
|
parsed = json.loads(candidate)
|
|
except:
|
|
parsed = json.loads(candidate.replace("'", '"'))
|
|
|
|
if isinstance(parsed, dict) and 'action' in parsed:
|
|
actions.append(parsed)
|
|
break
|
|
except: pass
|
|
|
|
if actions:
|
|
return json.dumps(actions)
|
|
return None
|
|
|