Retriever

from gaussia.core.retriever import Retriever
from gaussia.schemas.common import Dataset

class MyRetriever(Retriever):
    def load_dataset(self) -> list[Dataset]:
        # Load and return your data
        ...

class Retriever(ABC):
    def __init__(self, **kwargs):
        self.kwargs = kwargs

    @property
    def iteration_level(self) -> IterationLevel:
        return IterationLevel.FULL_DATASET

    @abstractmethod
    def load_dataset(self) -> list[Dataset] | Iterator[Dataset] | Iterator[StreamedBatch]:
        ...

class FullRetriever(Retriever):
    def load_dataset(self) -> list[Dataset]:
        return [Dataset(...), Dataset(...)]

from gaussia.schemas.common import IterationLevel

class StreamRetriever(Retriever):
    @property
    def iteration_level(self) -> IterationLevel:
        return IterationLevel.STREAM_SESSIONS

    def load_dataset(self):
        for row in database.fetch_sessions():
            yield Dataset(
                session_id=row["id"],
                assistant_id=row["assistant"],
                context=row["context"],
                conversation=[Batch(**b) for b in row["batches"]],
            )

from gaussia.schemas.common import IterationLevel, SessionMetadata, StreamedBatch

class EventRetriever(Retriever):
    @property
    def iteration_level(self) -> IterationLevel:
        return IterationLevel.STREAM_BATCHES

    def load_dataset(self):
        for event in message_queue.consume():
            yield StreamedBatch(
                metadata=SessionMetadata(
                    session_id=event["session_id"],
                    assistant_id=event["assistant_id"],
                    context=event["context"],
                ),
                batch=Batch(**event["interaction"]),
            )

# These kwargs are passed to MyRetriever.__init__
results = Context.run(
    MyRetriever,
    model=model,
    db_url="postgresql://...",  # Forwarded to retriever
    limit=100,                  # Forwarded to retriever
)

class MyRetriever(Retriever):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.db_url = kwargs["db_url"]
        self.limit = kwargs.get("limit", 50)