| """ |
| Lightweight TF-IDF based retriever to replace llama-index VectorStoreIndex. |
| Uses scikit-learn (already a dependency) — no heavyweight RAG framework needed. |
| """ |
|
|
| import numpy as np |
| from sklearn.feature_extraction.text import TfidfVectorizer |
| from sklearn.metrics.pairwise import cosine_similarity |
|
|
|
|
| class TextNode: |
| """Minimal node wrapper matching the llama-index retrieval result interface.""" |
|
|
| def __init__(self, text: str): |
| self.text = text |
|
|
|
|
| class Document: |
| """Minimal document wrapper matching the llama-index Document interface.""" |
|
|
| def __init__(self, text: str): |
| self.text = text |
|
|
|
|
| class SimpleRetriever: |
| """ |
| TF-IDF cosine-similarity retriever that mirrors the subset of the |
| llama-index VectorStoreIndex / BaseRetriever interface actually used: |
| |
| index = SimpleRetriever.from_documents([Document(text=x) for x in docs]) |
| ret = index.as_retriever(similarity_top_k=1) |
| nodes = ret.retrieve(query) # -> list[TextNode] |
| text = nodes[0].text |
| """ |
|
|
| def __init__(self, texts: list[str]): |
| self._docs = texts |
| self._top_k = 1 |
| self._vectorizer = TfidfVectorizer() |
| if texts: |
| self._matrix = self._vectorizer.fit_transform(texts) |
| else: |
| self._matrix = None |
|
|
| @classmethod |
| def from_documents(cls, documents) -> "SimpleRetriever": |
| """Build a retriever from a list of Document objects.""" |
| texts = [doc.text for doc in documents] |
| return cls(texts) |
|
|
| def as_retriever(self, similarity_top_k: int = 1) -> "SimpleRetriever": |
| """Configure top-k and return self (mirrors llama-index fluent API).""" |
| self._top_k = similarity_top_k |
| return self |
|
|
| def retrieve(self, query: str) -> list[TextNode]: |
| """Return the top-k most similar nodes for a query.""" |
| if not self._docs or self._matrix is None: |
| return [TextNode("")] |
|
|
| query_vec = self._vectorizer.transform([query]) |
| scores = cosine_similarity(query_vec, self._matrix)[0] |
| top_indices = np.argsort(scores)[::-1][: self._top_k] |
| return [TextNode(self._docs[i]) for i in top_indices] |
|
|
