(no commit message)

service/retrievers.py

@@ -0,0 +1,73 @@
+"""Retriever implementations used by the JTBD DSPy agent."""
+
+from __future__ import annotations
+
+from typing import Iterable, List
+
+from modaic import PrecompiledConfig, Retriever
+
+
+class NullRetrieverConfig(PrecompiledConfig):
+    """Configuration placeholder for the null retriever."""
+
+
+class NotesRetrieverConfig(PrecompiledConfig):
+    """Serializable configuration for the in-memory notes retriever."""
+
+    notes: List[str] = []
+    top_k: int = 3
+
+
+class NullRetriever(Retriever):
+    """No-op retriever for environments without contextual data."""
+
+    config: NullRetrieverConfig
+
+    def __init__(self, config: NullRetrieverConfig | None = None, **kwargs):
+        super().__init__(config or NullRetrieverConfig(), **kwargs)
+
+    def retrieve(self, query: str) -> str:  # type: ignore[override]
+        return ""
+
+
+class NotesRetriever(Retriever):
+    """Very small keyword-based retriever backed by an in-memory list of notes."""
+
+    config: NotesRetrieverConfig
+
+    def __init__(
+        self,
+        notes: Iterable[str] | None = None,
+        top_k: int | None = None,
+        config: NotesRetrieverConfig | None = None,
+        **kwargs,
+    ):
+        if config is None:
+            cfg = NotesRetrieverConfig()
+            cfg.notes = list(notes or [])
+            if top_k is not None:
+                cfg.top_k = int(top_k)
+        else:
+            cfg = config
+            if notes is not None:
+                cfg.notes = list(notes)
+            if top_k is not None:
+                cfg.top_k = int(top_k)
+
+        super().__init__(cfg, **kwargs)
+
+    def retrieve(self, query: str) -> str:  # type: ignore[override]
+        terms = {token for token in query.lower().split() if token}
+        if not terms:
+            return ""
+
+        scored: List[tuple[int, str]] = []
+        for note in self.config.notes:
+            tokens = {token for token in note.lower().split() if token}
+            score = len(terms & tokens)
+            if score > 0:
+                scored.append((score, note))
+
+        scored.sort(key=lambda item: item[0], reverse=True)
+        top_matches = [note for _, note in scored[: self.config.top_k]]
+        return "\n".join(top_matches)