Spaces:
Paused
Paused
File size: 5,766 Bytes
7ad7699 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
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 []
|