RAGGER/agent.py

from typing import Optional
from enum import Enum

from vectorstore import vectorstore
from ollama_client import ollama
from logger import get_logger

log = get_logger("agent")


class AgentMode(str, Enum):
    CHAT = "chat"
    LETTER = "lettera"
    PROJECT = "progetto"


AGENT_SYSTEM_PROMPT = """Sei un assistente AI intelligente per un'associazione di volontariato.
Il tuo compito è analizzare le richieste degli utenti e fornire risposte accurate.

Hai accesso ai seguenti strumenti:
- Ricerca documenti (vectorstore)
- Generazione lettere
- Generazione progetti

Analizza sempre la richiesta dell'utente e usa lo strumento più appropriato."""


class Agent:
    def __init__(self):
        self.mode = AgentMode.CHAT
        log.info("Agent initialized")
    
    def set_mode(self, mode: str):
        """Imposta la modalità di lavoro dell'agente"""
        try:
            self.mode = AgentMode(mode)
            log.info(f"Agent mode set to: {self.mode.value}")
        except ValueError:
            log.warning(f"Unknown mode: {mode}, defaulting to chat")
            self.mode = AgentMode.CHAT
    
    def process(
        self,
        user_input: str,
        mode: Optional[str] = None,
        context: Optional[str] = None
    ) -> dict:
        """Elabora l'input dell'utente in base alla modalità"""
        if mode:
            self.set_mode(mode)
        
        log.info(f"Agent processing: mode={self.mode.value}, input='{user_input[:50]}...'")
        
        if self.mode == AgentMode.CHAT:
            return self._process_chat(user_input)
        elif self.mode == AgentMode.LETTER:
            return self._process_letter(user_input, context)
        elif self.mode == AgentMode.PROJECT:
            return self._process_project(user_input, context)
        
        return {"error": "Modalità non supportata"}
    
    def _process_chat(self, question: str) -> dict:
        """Elabora una domanda in chat"""
        log.debug(f"Chat mode: '{question}'")
        
        results = self._multi_pass_search(question)
        
        if not results:
            return {
                "answer": "Non ho trovato informazioni rilevanti nei documenti.",
                "sources": [],
                "mode": "chat"
            }
        
        context = self._build_context(results)
        
        prompt = f"""Basandoti sul seguente contesto estratto dai documenti, rispondi alla domanda.

CONTESTO:
{context}

DOMANDA: {question}

Rispondi in italiano, citando le fonti quando appropriato."""
        
        answer = ollama.generate(
            prompt=prompt,
            system=AGENT_SYSTEM_PROMPT
        )
        
        sources = self._format_sources(results)
        
        return {
            "answer": answer,
            "sources": sources,
            "mode": "chat"
        }
    
    def _process_letter(self, topic: str, extra_context: Optional[str] = None) -> dict:
        """Elabora una richiesta di generazione lettera"""
        log.debug(f"Letter mode: '{topic}'")
        
        results = self._multi_pass_search(topic)
        
        doc_context = ""
        if results:
            doc_context = "\n\n".join(r["text"] for r in results)
        
        full_context = ""
        if doc_context:
            full_context += f"Documenti rilevanti:\n{doc_context}\n\n"
        if extra_context:
            full_context += f"Contesto aggiuntivo:\n{extra_context}\n\n"
        
        prompt = f"""Genera una lettera formale per un'associazione di volontariato.

Argomento: {topic}

{full_context}

Scrivi una lettera professionale in italiano, con:
- Intestazione (mittente, destinatario, data)
- Oggetto
- Corpo della lettera
- Cordiali saluti
- Firma"""
        
        letter = ollama.generate(prompt=prompt, system=AGENT_SYSTEM_PROMPT)
        
        sources = self._format_sources(results) if results else []
        
        return {
            "answer": letter,
            "sources": sources,
            "mode": "lettera"
        }
    
    def _process_project(
        self,
        description: str,
        title: Optional[str] = None
    ) -> dict:
        """Elabora una richiesta di generazione progetto"""
        log.debug(f"Project mode: '{description}'")
        
        search_query = f"{title or ''} {description}"
        results = self._multi_pass_search(search_query)
        
        doc_context = ""
        if results:
            doc_context = "\n\n".join(r["text"] for r in results)
        
        prompt = f"""Genera una proposta di progetto per un'associazione di volontariato.

Titolo: {title or 'Da definire'}
Descrizione: {description}

{f'Documenti rilevanti:\n{doc_context}' if doc_context else ''}

Scrivi una proposta di progetto completa con:
1. Titolo e breve introduzione
2. Obiettivi del progetto
3. Attività previste
4. Risorse necessarie (umane, finanziarie, materiali)
5. Timeline e milestone
6. Risultati attesi e indicatori di successo
7. Partner e stakeholder"""
        
        project = ollama.generate(prompt=prompt, system=AGENT_SYSTEM_PROMPT)
        
        sources = self._format_sources(results) if results else []
        
        return {
            "answer": project,
            "sources": sources,
            "mode": "progetto"
        }
    
    def _multi_pass_search(self, query: str, passes: int = 2) -> list:
        """
        Multi-pass retrieval: ricerche multiple per query complesse.
        Prima ricerca con query originale, poi con versione riformulata.
        """
        log.debug(f"Multi-pass search: {passes} passes")
        
        results = vectorstore.search(query, limit=5)
        
        if passes > 1 and len(results) < 3:
            log.debug("Low results, trying reformulated query")
            reformulated = self._reformulate_query(query)
            if reformulated != query:
                results2 = vectorstore.search(reformulated, limit=3)
                existing_ids = {r.get("id") for r in results}
                for r in results2:
                    if r.get("id") not in existing_ids:
                        results.append(r)
        
        log.debug(f"Total results: {len(results)}")
        return results
    
    def _reformulate_query(self, query: str) -> str:
        """Riformula la query per migliorare il retrieval"""
        prompt = f"""Riformula questa domanda per migliorare la ricerca in un database vettoriale.
Restituisci solo la query riformulata, nient'altro.

Query originale: {query}

Query riformulata:"""
        
        try:
            reformulated = ollama.generate(prompt=prompt, system="Sei un assistente di ricerca.")
            return reformulated.strip()
        except Exception as e:
            log.warning(f"Query reformulation failed: {e}")
            return query
    
    def _build_context(self, results: list) -> str:
        """Costruisce il contesto dai risultati della ricerca"""
        return "\n\n---\n\n".join(
            f"[Fonte: {r['source']}" + 
            (f", pagina {r['page']}" if r.get('page') else "") + 
            f"]\n{r['text']}"
            for r in results
        )
    
    def _format_sources(self, results: list) -> list:
        """Formatta le fonti per la risposta"""
        sources = []
        for r in results:
            sources.append({
                "source": r.get("source", "Sconosciuto"),
                "page": r.get("page"),
                "score": r.get("score", 0.0),
                "excerpt": r["text"][:200] + "..." if len(r.get("text", "")) > 200 else r.get("text", "")
            })
        return sources


agent = Agent()