234 lines
9.4 KiB
Python
234 lines
9.4 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,
|
|
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
|
|
|