phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-17 11:22:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

correct output structure

Browse source
This commit is contained in:
Matthew Farrellee 2025-08-31 09:06:59 -04:00
parent 486cd8679a
commit 15c1f8b885
2 changed files with 12 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -298,8 +298,8 @@ class VectorDBWithIndex:
self.vector_db.embedding_model, self.vector_db.embedding_model,
[c.content for c in chunks_to_embed], [c.content for c in chunks_to_embed],
) )
for c, embedding in zip(chunks_to_embed, resp.data, strict=False): for c, data in zip(chunks_to_embed, resp.data, strict=False):
c.embedding = embedding c.embedding = data.embedding
embeddings = np.array([c.embedding for c in chunks], dtype=np.float32) embeddings = np.array([c.embedding for c in chunks], dtype=np.float32)
await self.index.add_chunks(chunks, embeddings) await self.index.add_chunks(chunks, embeddings)
@ -335,7 +335,7 @@ class VectorDBWithIndex:
return await self.index.query_keyword(query_string, k, score_threshold) return await self.index.query_keyword(query_string, k, score_threshold)
embeddings_response = await self.inference_api.openai_embeddings(self.vector_db.embedding_model, [query_string]) embeddings_response = await self.inference_api.openai_embeddings(self.vector_db.embedding_model, [query_string])
query_vector = np.array(embeddings_response.data[0], dtype=np.float32) query_vector = np.array(embeddings_response.data[0].embedding, dtype=np.float32)
if mode == "hybrid": if mode == "hybrid":
return await self.index.query_hybrid( return await self.index.query_hybrid(
query_vector, query_string, k, score_threshold, reranker_type, reranker_params query_vector, query_string, k, score_threshold, reranker_type, reranker_params

11
tests/unit/rag/test_vector_store.py
View file

@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock
import numpy as np import numpy as np
import pytest import pytest
from llama_stack.apis.inference.inference import OpenAIEmbeddingData
from llama_stack.apis.tools import RAGDocument from llama_stack.apis.tools import RAGDocument
from llama_stack.apis.vector_io import Chunk from llama_stack.apis.vector_io import Chunk
from llama_stack.providers.utils.memory.vector_store import ( from llama_stack.providers.utils.memory.vector_store import (
@ -218,7 +219,10 @@ class TestVectorDBWithIndex:
Chunk(content="Test 2", embedding=None, metadata={}), Chunk(content="Test 2", embedding=None, metadata={}),
] ]
mock_inference_api.openai_embeddings.return_value.data = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] mock_inference_api.openai_embeddings.return_value.data = [
OpenAIEmbeddingData(embedding=[0.1, 0.2, 0.3], index=0),
OpenAIEmbeddingData(embedding=[0.4, 0.5, 0.6], index=1),
]
await vector_db_with_index.insert_chunks(chunks) await vector_db_with_index.insert_chunks(chunks)
@ -310,7 +314,10 @@ class TestVectorDBWithIndex:
Chunk(content="Test 3", embedding=None, metadata={}), Chunk(content="Test 3", embedding=None, metadata={}),
] ]
mock_inference_api.openai_embeddings.return_value.data = [[0.1, 0.1, 0.1], [0.3, 0.3, 0.3]] mock_inference_api.openai_embeddings.return_value.data = [
OpenAIEmbeddingData(embedding=[0.1, 0.1, 0.1], index=0),
OpenAIEmbeddingData(embedding=[0.3, 0.3, 0.3], index=1),
]
await vector_db_with_index.insert_chunks(chunks) await vector_db_with_index.insert_chunks(chunks)