assistentes.py

import os
from dotenv import load_dotenv
from pinecone import Pinecone
from pinecone_plugins.assistant.models.chat import Message
from rich.console import Console
from rich.panel import Panel
from dataclasses import dataclass
from typing import List, Optional, Dict, Generator
from pathlib import Path
from datetime import datetime


load_dotenv()

@dataclass
class AssistenteConfig:
    """Configuração do assistente"""
    nome: str
    instrucoes: str

class CriaAssistente:
    """Classe responsável por criar e configurar assistentes"""
    
    def __init__(self):
        self.pc = Pinecone(api_key=os.getenv('PINECONE_API_KEY'))
        self.console = Console()
    
    def criar(self, config: AssistenteConfig) -> Dict:
        """Cria um novo assistente com as configurações especificadas"""
        try:
            assistant = self.pc.assistant.create_assistant(
                assistant_name=config.nome,
                instructions=config.instrucoes,
                timeout=30
            )
            
            self.console.print(
                f"[green]✓ Assistente '{config.nome}' criado com sucesso![/green]"
            )
            return assistant
            
        except Exception as e:
            self.console.print(f"[red]✗ Erro ao criar assistente: {str(e)}[/red]")
            raise

class Assistente:
    """Classe para interagir com assistentes existentes"""
    
    def __init__(self, assistant_name: str):
        self.pc = Pinecone(api_key=os.getenv('PINECONE_API_KEY'))
        self.console = Console()
        self.nome = assistant_name
        self.assistente = self.pc.assistant.Assistant(assistant_name=assistant_name)
    
    def get_context(self, query: str) -> List[Dict]:
        """Busca snippets relevantes para uma query"""
        try:
            response = self.assistente.context(query=query)
            if response is None:
                return []
            
            # Verificar se response é um dicionário
            if isinstance(response, dict):
                return response.get("snippets", [])
            # Se for uma lista, retornar diretamente
            elif isinstance(response, list):
                return response
            return []
        
        except Exception as e:
            self.console.print(f"[red]✗ Erro ao buscar contexto: {str(e)}[/red]")
            return []

    def fazer_pergunta(self, pergunta: str) -> Dict:
        """Faz uma pergunta ao assistente com contexto"""
        try:
            # Primeiro buscar contexto relevante
            snippets = self.get_context(pergunta)
            
            # Fazer a pergunta
            msg = Message(content=pergunta)
            response = self.assistente.chat(messages=[msg])
            
            # Extrair conteúdo e citações
            content = response["message"]["content"]
            citations = response.get("citations", [])
            
            return {
                "content": content,
                "citations": citations,
                "snippets": snippets  # Incluir snippets na resposta
            }
            
        except Exception as e:
            self.console.print(f"[red]✗ Erro ao fazer pergunta: {str(e)}[/red]")
            raise
    
    def _process_stream(self, chunks: Generator) -> Dict:
        """Processa o stream de resposta"""
        content = ""
        citations = []
        
        for chunk in chunks:
            if not chunk:
                continue
                
            # Processar diferentes tipos de chunks
            if isinstance(chunk, dict):
                chunk_type = chunk.get("type")
                
                if chunk_type == "message_start":
                    continue
                elif chunk_type == "content_chunk":
                    if "delta" in chunk and "content" in chunk["delta"]:
                        content += chunk["delta"]["content"]
                elif chunk_type == "citation":
                    citations.append(chunk.get("citation", {}))
                elif chunk_type == "message_end":
                    if chunk.get("finish_reason") == "stop":
                        break
            else:
                content += str(chunk)
        
        return {
            "content": content,
            "citations": citations
        }
    
    def upload_file(self, file_path: str, timeout: Optional[int] = None) -> Dict:
        """Faz upload de um arquivo para o assistente"""
        try:
            # Adicionar metadados com nome original do arquivo
            original_name = Path(file_path).name
            metadata = {
                "original_name": original_name,
                "uploaded_on": datetime.now().isoformat()
            }
            
            response = self.assistente.upload_file(
                file_path=file_path,
                timeout=timeout,
                metadata=metadata
            )
            return response
        except Exception as e:
            self.console.print(f"[red]✗ Erro ao fazer upload: {str(e)}[/red]")
            raise

    def get_files(self) -> List[Dict]:
        """Retorna a lista de arquivos carregados no assistente"""
        try:
            response = self.assistente.list_files()
            if isinstance(response, list):
                return response
            elif isinstance(response, dict):
                return response.get("files", [])
            return []
        except Exception as e:
            self.console.print(f"[red]✗ Erro ao listar arquivos: {str(e)}[/red]")
            return []