Camadas de Memória
em Ambiente Distribuído

from langchain_core.messages import trim_messages

def supervisor_node(state: AgentState):
    mensagens_trimadas = trim_messages(
        state["messages"],
        strategy="last",       # mantém as mensagens mais recentes
        token_counter=llm,       # conta tokens reais do modelo
        max_tokens=8000,         # reserva espaço para a resposta
        start_on="human",        # sempre começa com mensagem humana
        include_system=True,     # mantém o system prompt
        allow_partial=False,     # sem mensagens cortadas no meio
    )
    response = supervisor_llm.invoke(mensagens_trimadas)
    return {"messages": [response]}

from langgraph.checkpoint.redis import AsyncRedisSaver
import redis.asyncio as aioredis

redis_client = aioredis.from_url(
    "redis://elasticache-redis.interno:6379",
    decode_responses=True
)
checkpointer = AsyncRedisSaver(redis_client)

# thread_id = identificador único da sessão do usuário
config = {"configurable": {"thread_id": "user_123_session_456"}}

# Invocação 1: usuário faz pergunta
result1 = await app.ainvoke(
    {"messages": [HumanMessage("Pesquise sobre X")]},
    config
)

# Invocação 2: usuário continua, 5 min depois
# O grafo carrega o estado anterior do Redis automaticamente
result2 = await app.ainvoke(
    {"messages": [HumanMessage("Agora analise os dados")]},
    config  # mesmo thread_id → continua de onde parou
)

class ContextEnricher:
    def __init__(self, redis_client):
        self.redis = redis_client

    async def build_payload(self, session_id: str, agent_name: str, instrucoes: str) -> dict:
        raw = await self.redis.get(f"working:{session_id}")
        working_mem = json.loads(raw) if raw else {}

        # Filtra contexto relevante por agente — evita payload desnecessário
        context_map = {
            "agent_pesquisa":  ["queries_anteriores", "intent"],
            "agent_analise":   ["pesquisa_resultado", "intent", "dados_coletados"],
            "agent_relatorio": ["pesquisa_resultado", "analise_resultado", "intent"],
        }

        keys = context_map.get(agent_name, [])
        contexto = {k: working_mem[k] for k in keys if k in working_mem}

        return {"session_id": session_id, "instrucoes": instrucoes, "contexto": contexto}

    async def persist_result(self, session_id: str, agent_name: str, resultado: str):
        # Supervisor é o único que escreve no Redis
        raw = await self.redis.get(f"working:{session_id}")
        working_mem = json.loads(raw) if raw else {}
        result_key = agent_name.replace("agent_", "") + "_resultado"
        working_mem[result_key] = resultado
        await self.redis.setex(f"working:{session_id}", 3600, json.dumps(working_mem))

from langchain_community.vectorstores import OpenSearchVectorSearch
from langchain_aws import BedrockEmbeddings

def criar_retriever(opensearch_url: str, index: str, aws_region: str, strategy: str):
    embeddings = BedrockEmbeddings(
        model_id="amazon.titan-embed-text-v2:0",
        region_name=aws_region
    )
    vectorstore = OpenSearchVectorSearch(
        opensearch_url=opensearch_url,
        index_name=index,
        embedding_function=embeddings,
        http_auth=get_aws_auth(opensearch_url, aws_region),  # SigV4
        use_ssl=True,
        verify_certs=True,
    )

    configs = {
        "pesquisa":  {"search_type": "similarity", "k": 8,  "score_threshold": 0.65},
        "analise":   {"search_type": "mmr",        "k": 5,  "fetch_k": 20, "lambda_mult": 0.7},
        "relatorio": {"search_type": "similarity_score_threshold", "score_threshold": 0.85, "k": 3},
    }
    return vectorstore.as_retriever(search_kwargs=configs[strategy])

class LongTermMemoryManager:
    async def salvar_fato(self, user_id: str, fato: str):
        embedding = await gerar_embedding(fato)
        await self.os.index(
            index=f"user-memory-{user_id}",
            body={
                "fato":       fato,
                "embedding":  embedding,
                "timestamp":  datetime.utcnow().isoformat(),
            }
        )

    async def recuperar_contexto(self, user_id: str, query: str, k: int = 3) -> List[str]:
        query_emb = await gerar_embedding(query)
        resp = await self.os.search(
            index=f"user-memory-{user_id}",
            body={"knn": {"embedding": {"vector": query_emb, "k": k}}}
        )
        return [hit["_source"]["fato"] for hit in resp["hits"]["hits"]]

# Uso no supervisor: enriquece o system prompt com memória do usuário
ltm_fatos = await ltm.recuperar_contexto(
    user_id=state["user_id"],
    query=state["messages"][-1].content
)
system = f"""Contexto sobre este usuário:
{chr(10).join(ltm_fatos)}
"""