Use with DSPy

import dspy
from atlaso import Memory

class AtlasoRetrieve(dspy.Retrieve):
    def __init__(self, user_id: str, k: int = 8):
        super().__init__(k=k)
        self.user = Memory().for_user(user_id)

    def forward(self, query_or_queries, k=None):
        queries = (
            [query_or_queries]
            if isinstance(query_or_queries, str)
            else list(query_or_queries)
        )
        results = []
        for q in queries:
            hits = self.user.recall(q, limit=k or self.k)
            for h in hits:
                results.append(
                    dspy.Example(
                        long_text=h.content,
                        score=h.score,
                        is_confident=h.is_confident,
                        has_disagreement=h.has_disagreement,
                        agreement_score=h.agreement_score,
                    )
                )
        return dspy.Prediction(passages=results)