File size: 3,455 Bytes
5374a2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | from llama_index.core.indices.base import BaseIndex
from llama_index.core.retrievers import VectorIndexRetriever
from .base import BaseRetrieverWrapper
from evoagentx.core.logging import logger
from evoagentx.rag.schema import Query, RagResult, Corpus
class VectorRetriever(BaseRetrieverWrapper):
"""Wrapper for vector-based retrieval."""
def __init__(self, index: BaseIndex, top_k: int = 5, chunk_class=None):
super().__init__()
self.index = index
self.top_k = top_k
self.chunk_class = chunk_class # Use chunk class passed from RAGEngine based on config
self.retriever = VectorIndexRetriever(
index=self.index,
similarity_top_k=self.top_k
)
async def aretrieve(self, query: Query) -> RagResult:
try:
# config the top_k
self.retriever.similarity_top_k = query.top_k
nodes = await self.retriever.aretrieve(query.query_str)
corpus = Corpus()
scores = []
if nodes is None:
return RagResult(corpus=corpus, scores=scores, metadata={"query": query.query_str, "retriever": "vector"})
for score_node in nodes:
if self.chunk_class is None:
raise ValueError("chunk_class not set - RAGEngine must pass chunk class based on config")
chunk = self.chunk_class.from_llama_node(score_node.node)
chunk.metadata.similarity_score = score_node.score or 0.0
corpus.add_chunk(chunk)
scores.extend([score_node.score or 0.0])
result = RagResult(
corpus=corpus,
scores=scores,
metadata={"query": query.query_str, "retriever": "vector"}
)
logger.info(f"Vector retrieved {len(corpus.chunks)} chunks")
return result
except Exception as e:
logger.error(f"Vector retrieval failed: {str(e)}")
raise
def retrieve(self, query: Query) -> RagResult:
try:
# config the top_k
self.retriever.similarity_top_k = query.top_k
nodes = self.retriever.retrieve(query.query_str)
corpus = Corpus()
scores = []
if nodes is None:
return RagResult(corpus=corpus, scores=scores, metadata={"query": query.query_str, "retriever": "vector"})
for score_node in nodes:
if self.chunk_class is None:
raise ValueError("chunk_class not set - RAGEngine must pass chunk class based on config")
chunk = self.chunk_class.from_llama_node(score_node.node)
chunk.metadata.similarity_score = score_node.score or 0.0
corpus.add_chunk(chunk)
scores.extend([score_node.score or 0.0])
result = RagResult(
corpus=corpus,
scores=scores,
metadata={"query": query.query_str, "retriever": "vector"}
)
logger.info(f"Vector retrieved {len(corpus.chunks)} chunks")
return result
except Exception as e:
logger.error(f"Vector retrieval failed: {str(e)}")
raise
def get_retriever(self) -> VectorIndexRetriever:
logger.debug("Returning vector retriever")
return self.retriever |