brain.py

import asyncio
import os
from knowledge_base import KnowledgeBase
from scraper import search_and_extract
import worldmonitor_client
from errors import get_logger, GenerAIError, ErrorCode, fmt_exc

log = get_logger("brain")

LOCAL_THRESHOLD = 0.40
HF_MODEL = os.environ.get("HF_MODEL", "amogaddy/GenerAI")
HF_TOKEN = os.environ.get("HF_TOKEN", "")


def _rerank(results: list, query: str) -> list:
    words = set(query.lower().split())
    for r in results:
        stored_q = r["metadata"].get("query", "")
        overlap = len(words & set(stored_q.lower().split()))
        r["score"] = r["distance"] - overlap * 0.12
    return sorted(results, key=lambda x: x["score"])


def _extract_sentences(text: str, query: str, max_chars: int = 700) -> str:
    words = set(query.lower().split())
    sentences = [s.strip() for s in text.replace("\n", ". ").split(".") if len(s.strip()) > 20]
    scored = [(sum(1 for w in words if w in s.lower()), s) for s in sentences]
    scored.sort(key=lambda x: -x[0])
    result = ""
    for _, s in scored:
        if len(result) + len(s) + 2 > max_chars:
            break
        result += s + ". "
    return result.strip() or text[:max_chars]


class _HFClient:
    SYSTEM = (
        "Sei GenerAI, un assistente AI specializzato in lingua italiana. "
        "Rispondi in modo chiaro, preciso e sempre in italiano. "
        "Se ti viene fornito del contesto, basati su quello per rispondere."
    )

    def __init__(self, model: str, token: str):
        try:
            from huggingface_hub import InferenceClient
            self._client = InferenceClient(model=model, token=token or None)
            self._model = model
            log.info("HuggingFace LLM pronto: %s", model)
        except ImportError:
            raise GenerAIError(ErrorCode.BRAIN_ASK_FAILED, "huggingface_hub non installato.")
        except Exception as e:
            raise GenerAIError(ErrorCode.BRAIN_ASK_FAILED,
                               f"Impossibile connettersi a HF ({model}): {fmt_exc(e)}", cause=e)

    def generate(self, question: str, context: str = "") -> str:
        user_msg = f"Contesto:\n{context}\n\nDomanda: {question}" if context else question
        messages = [
            {"role": "system", "content": self.SYSTEM},
            {"role": "user",   "content": user_msg},
        ]
        try:
            r = self._client.chat_completion(messages=messages, max_tokens=512, temperature=0.3)
            return r.choices[0].message.content.strip()
        except Exception as e:
            raise GenerAIError(ErrorCode.BRAIN_ASK_FAILED,
                               f"Generazione HF fallita: {fmt_exc(e)}", cause=e)


class Brain:
    def __init__(self):
        log.info("Avvio Brain...")
        try:
            self.kb = KnowledgeBase()
        except GenerAIError as e:
            e.log(log)
            raise
        self._last_doc_id: str | None = None
        self._hf: _HFClient | None = None
        if HF_TOKEN:
            try:
                self._hf = _HFClient(HF_MODEL, HF_TOKEN)
            except GenerAIError as e:
                e.log(log)
                log.warning("Fallback a ricerca web + estrazione (HF non disponibile).")
        else:
            log.info("HF_TOKEN non impostato — uso ricerca web + estrazione testo.")
        log.info("Brain pronto. Modalita: %s", "LLM+KB" if self._hf else "KB+Web")

    async def ask(self, question: str, on_status=None) -> tuple[str, str]:
        async def emit(msg: str):
            if on_status:
                await on_status(msg)

        self._last_doc_id = None
        log.info("Domanda: %r", question)

        await emit("Cerco nella memoria locale...")
        try:
            results = _rerank(self.kb.search(question, n_results=12), question)
        except Exception as e:
            log.warning("KB search fallita: %s", fmt_exc(e))
            results = []

        if results and results[0]["score"] < LOCAL_THRESHOLD:
            best = results[0]
            await emit(f"Trovato in memoria locale (affidabilita: {int((1-results[0]['score'])*100)}%)")
            if self._hf:
                await emit("Elaboro risposta con intelligenza artificiale...")
                return await self._llm_answer(question, best["answer"],
                                              best["metadata"].get("source", ""), "llm")
            answer = _extract_sentences(best["answer"], question)
            src = best["metadata"].get("source", "")
            if src and src != "grammatica_italiana":
                answer += f"\n\n*Fonte: {src}*"
            await emit("Risposta pronta dalla memoria locale.")
            return answer, "local"

        await emit("Non trovato in memoria — consulto World Monitor...")
        try:
            wm_items = await asyncio.to_thread(worldmonitor_client.find_relevant, question, 5)
        except Exception as e:
            log.debug("World Monitor non consultabile: %s", fmt_exc(e))
            wm_items = []

        # World Monitor è il primo posto in cui si cerca: se ha risultati
        # pertinenti a sufficienza, si evita del tutto la ricerca web (più
        # veloce, nessun browser Chromium da avviare). Altrimenti si scende
        # sul web come prima, e gli item WM trovati (se pochi) restano come
        # contesto aggiuntivo.
        web_results = []
        if len(wm_items) >= 2:
            await emit(f"Trovati {len(wm_items)} risultati pertinenti su World Monitor.")
        else:
            await emit("World Monitor non basta — cerco anche sul web...")
            try:
                web_results = await asyncio.to_thread(search_and_extract, question)
            except Exception as e:
                log.warning("Ricerca fallita: %s", fmt_exc(e))
                web_results = []

        if not web_results and not wm_items:
            await emit("Nessuna informazione trovata.")
            return (
                "Non ho trovato informazioni su questo argomento. "
                "Prova a riformulare la domanda.",
                "unknown",
            )

        await emit(f"Trovati {len(web_results) + len(wm_items)} risultati — elaboro...")
        await asyncio.sleep(0.2)

        blocks = [f"[{r['title']}]\n{r['text']}" for r in web_results]
        if wm_items:
            wm_block = "\n".join(f"- {it['title']} ({it['source']})" for it in wm_items)
            blocks.insert(0, f"[World Monitor — notizie recenti]\n{wm_block}")

        combined = "\n\n---\n\n".join(blocks)
        sources = ", ".join(r["url"] for r in web_results if r.get("url")) \
            or ", ".join(it["url"] for it in wm_items if it.get("url"))

        if self._hf:
            await emit("Sintetizzo con intelligenza artificiale...")
            answer, status = await self._llm_answer(question, combined, sources, "llm")
        else:
            await emit("Estraggo i passaggi significativi...")
            answer = _extract_sentences(combined, question)
            if sources:
                answer += f"\n\n*Fonte: {sources}*"
            status = "searched"

        # Auto-apprendimento: salva in KB se domanda frequente
        learned = False
        try:
            if self.kb.increment_and_should_archive(question):
                doc_id = self.kb.add(question, answer, sources)
                self._last_doc_id = doc_id
                learned = True
                await emit("Concetto appreso e salvato.")
        except Exception as e:
            log.warning("Auto-apprendimento fallito: %s", fmt_exc(e))

        await emit("Risposta pronta!")
        return answer, status

    async def _llm_answer(self, question: str, context: str,
                          source: str, status: str) -> tuple[str, str]:
        try:
            answer = await asyncio.to_thread(self._hf.generate, question, context)
            if source and source != "grammatica_italiana":
                answer += f"\n\n*Fonte: {source}*"
            return answer, status
        except GenerAIError as e:
            e.log(log)
            answer = _extract_sentences(context, question)
            if source and source != "grammatica_italiana":
                answer += f"\n\n*Fonte: {source}*"
            return answer, "searched"

    def give_feedback(self, positive: bool):
        if self._last_doc_id:
            self.kb.reinforce(self._last_doc_id, positive)

    @property
    def kb_size(self) -> int:
        return self.kb.count()