Loading
Loading
LlamaIndex's ChatMessage carries additional_kwargs — the right shape for atlaso's dispersion fields.
LangChain's BaseMemory exchanges dict[str, str]with the chain — dispersion flags drop at the boundary. LlamaIndex's BaseMemory exchanges ChatMessage objects which carry an additional_kwargs: dict[str, Any]— atlaso's flags survive end-to-end into the prompt.
from llama_index.core.memory import BaseMemory
from llama_index.core.llms import ChatMessage, MessageRole
from atlaso import Memory
class AtlasoLlamaMemory(BaseMemory):
def __init__(self, user_id: str):
self.user = Memory().for_user(user_id)
def get(self, input_str: str | None = None, **kwargs) -> list[ChatMessage]:
if not input_str:
return []
hits = self.user.recall(input_str, limit=8)
msgs = []
for h in hits:
msgs.append(
ChatMessage(
role=MessageRole.SYSTEM,
content=h.content,
additional_kwargs={
"is_confident": h.is_confident,
"has_disagreement": h.has_disagreement,
"agreement_score": h.agreement_score,
"conflict_peers": list(h.conflict_peers),
},
)
)
return msgs
def put(self, message: ChatMessage) -> None:
if message.role == MessageRole.USER:
self.user.add(message.content)
def get_all(self) -> list[ChatMessage]:
return []
def reset(self) -> None:
raise NotImplementedError("Atlaso deposits are immutable.")Was this page helpful?