phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-28 09:10:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat (RAG): Implement configurable search mode in RAGQueryConfig

Browse source 
Signed-off-by: Varsha Prasad Narsing <varshaprasad96@gmail.com>
This commit is contained in:
Varsha Prasad Narsing 2025-04-14 16:53:17 -07:00
parent 85b5f3172b
commit e2a7022d3c
14 changed files with 210 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/providers/utils/memory/vector_store.py
View file

@ -177,7 +177,9 @@ class EmbeddingIndex(ABC):
raise NotImplementedError()
@abstractmethod
async def query(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
async def query(
self, embedding: NDArray, query_string: Optional[str], k: int, score_threshold: float, mode: Optional[str]
) -> QueryChunksResponse:
raise NotImplementedError()
@abstractmethod
@ -210,9 +212,9 @@ class VectorDBWithIndex:
if params is None:
params = {}
k = params.get("max_chunks", 3)
mode = params.get("mode")
score_threshold = params.get("score_threshold", 0.0)
query_str = interleaved_content_as_str(query)
embeddings_response = await self.inference_api.embeddings(self.vector_db.embedding_model, [query_str])
query_string = interleaved_content_as_str(query)
embeddings_response = await self.inference_api.embeddings(self.vector_db.embedding_model, [query_string])
query_vector = np.array(embeddings_response.embeddings[0], dtype=np.float32)
return await self.index.query(query_vector, k, score_threshold)
return await self.index.query(query_vector, query_string, k, score_threshold, mode)