File size: 5,393 Bytes
831e835 |
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 |
# Arquivo: jade/shorestone.py
import chromadb
import uuid
import logging
import time
import joblib
import numpy as np
from sentence_transformers import SentenceTransformer
# Configura um logger para nosso m贸dulo de mem贸ria
logging.basicConfig(level=logging.INFO, format="%(asctime)s - SHORESTONE - %(levelname)s - %(message)s")
class ShoreStoneMemory:
"""
Sistema de mem贸ria vetorial persistente e enriquecido com metadados para curadoria.
Salva os dados em disco e gerencia o ciclo de vida das mem贸rias.
"""
def __init__(self, model_name='all-MiniLM-L6-v2', storage_path="./jade_memory_db", pca_model_path="pca_model.joblib"):
"""
Inicializa o ShoreStone.
:param model_name: O modelo de embedding a ser usado.
:param storage_path: O diret贸rio no disco onde a mem贸ria ser谩 salva.
:param pca_model_path: (Opcional) Caminho para um modelo PCA treinado para compress茫o de vetores.
"""
logging.info("Iniciando o motor de mem贸ria ShoreStone...")
self.model = SentenceTransformer(model_name)
# Usa PersistentClient para salvar os dados em disco
self.client = chromadb.PersistentClient(path=storage_path)
self.collection = None
# Tenta carregar o modelo PCA para compress茫o de vetores
try:
logging.info(f"Tentando carregar modelo PCA de '{pca_model_path}'...")
self.pca = joblib.load(pca_model_path)
logging.info("Modelo PCA carregado. A compress茫o de vetores est谩 ATIVA.")
except FileNotFoundError:
logging.warning(f"Arquivo do modelo PCA '{pca_model_path}' n茫o encontrado. A compress茫o de vetores est谩 DESATIVADA.")
self.pca = None
logging.info("Motor ShoreStone pronto.")
def load_or_create_session(self, session_name: str):
"""
Carrega uma sess茫o de mem贸ria existente ou cria uma nova.
:param session_name: O nome da sess茫o (cole莽茫o no ChromaDB).
"""
logging.info(f"Carregando ou criando a sess茫o de mem贸ria: '{session_name}'")
self.collection = self.client.get_or_create_collection(name=session_name)
logging.info(f"Sess茫o '{session_name}' pronta com {self.collection.count()} mem贸rias.")
def _get_compressed_embedding(self, text: str) -> list:
"""Gera um embedding e o comprime com PCA, se o modelo PCA estiver carregado."""
embedding = self.model.encode(text)
if self.pca:
embedding_2d = embedding.reshape(1, -1)
compressed_embedding = self.pca.transform(embedding_2d)
return compressed_embedding.flatten().tolist()
return embedding.tolist()
def memorize(self, user_input: str, assistant_response: str):
"""Adiciona uma nova mem贸ria 脿 sess茫o ativa, enriquecida com metadados para curadoria."""
if not self.collection:
logging.error("Nenhuma sess茫o de mem贸ria ativa. Use 'load_or_create_session' primeiro.")
return
text = f"O usu谩rio disse: '{user_input}'. A IA respondeu: '{assistant_response}'"
embedding = self._get_compressed_embedding(text)
doc_id = str(uuid.uuid4())
current_timestamp = time.time()
metadata = {
"created_at": current_timestamp,
"last_accessed_at": current_timestamp,
"access_count": 0,
"text": text # Armazena o texto completo nos metadados para o Curador
}
self.collection.add(
embeddings=[embedding],
metadatas=[metadata],
documents=[text.split('.')[0] + "..."], # Documento pode ser um resumo para buscas r谩pidas
ids=[doc_id]
)
logging.info(f"Nova mem贸ria adicionada 脿 sess茫o '{self.collection.name}'.")
def remember(self, query: str, n_results: int = 3) -> str | None:
"""
Busca as mem贸rias mais relevantes e ATUALIZA seus metadados de acesso.
"""
if not self.collection or self.collection.count() == 0:
return None
query_embedding = self._get_compressed_embedding(query)
results = self.collection.query(
query_embeddings=[query_embedding],
n_results=n_results,
include=["metadatas", "documents"] # Essencial para obter os metadados para atualiza莽茫o
)
if not (results and results.get('ids') and results['ids'][0]):
return None
# Atualiza os metadados das mem贸rias que foram recuperadas (acessadas)
ids_to_update = results['ids'][0]
metadatas_to_update = results['metadatas'][0]
current_timestamp = time.time()
updated_metadatas = []
for meta in metadatas_to_update:
meta['access_count'] = meta.get('access_count', 0) + 1
meta['last_accessed_at'] = current_timestamp
updated_metadatas.append(meta)
if ids_to_update:
self.collection.update(ids=ids_to_update, metadatas=updated_metadatas)
logging.info(f"{len(ids_to_update)} mem贸rias tiveram seus metadados de acesso atualizados.")
retrieved_docs = "\n".join(results['documents'][0])
return retrieved_docs |