diff --git a/.github/actions/setup-test-environment/action.yml b/.github/actions/setup-test-environment/action.yml index 30b9b0130..4465fe159 100644 --- a/.github/actions/setup-test-environment/action.yml +++ b/.github/actions/setup-test-environment/action.yml @@ -39,6 +39,17 @@ runs: if: ${{ inputs.provider == 'vllm' && inputs.inference-mode == 'record' }} uses: ./.github/actions/setup-vllm + - name: Set provider URLs for replay mode + if: ${{ inputs.inference-mode == 'replay' }} + shell: bash + run: | + # setting so providers get registered in replay mode + if [ "${{ inputs.provider }}" == "ollama" ]; then + echo "OLLAMA_URL=http://localhost:11434" >> $GITHUB_ENV + elif [ "${{ inputs.provider }}" == "vllm" ]; then + echo "VLLM_URL=http://localhost:8000/v1" >> $GITHUB_ENV + fi + - name: Build Llama Stack shell: bash run: | diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html index 0549dda21..319490eff 100644 --- a/docs/_static/llama-stack-spec.html +++ b/docs/_static/llama-stack-spec.html @@ -14997,6 +14997,47 @@ "text": { "type": "string", "description": "The actual text content" + }, + "embedding": { + "type": "array", + "items": { + "type": "number" + }, + "description": "(Optional) Embedding vector for the content, if available" + }, + "created_timestamp": { + "type": "integer", + "description": "(Optional) Timestamp when the content was created" + }, + "metadata": { + "type": "object", + "additionalProperties": { + "oneOf": [ + { + "type": "null" + }, + { + "type": "boolean" + }, + { + "type": "number" + }, + { + "type": "string" + }, + { + "type": "array" + }, + { + "type": "object" + } + ] + }, + "description": "(Optional) Metadata associated with the content, such as source, author, etc." + }, + "chunk_metadata": { + "$ref": "#/components/schemas/ChunkMetadata", + "description": "(Optional) Metadata associated with the chunk, such as document ID, source, etc." } }, "additionalProperties": false, diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml index aa47cd58d..701c7e328 100644 --- a/docs/_static/llama-stack-spec.yaml +++ b/docs/_static/llama-stack-spec.yaml @@ -11143,6 +11143,34 @@ components: text: type: string description: The actual text content + embedding: + type: array + items: + type: number + description: >- + (Optional) Embedding vector for the content, if available + created_timestamp: + type: integer + description: >- + (Optional) Timestamp when the content was created + metadata: + type: object + additionalProperties: + oneOf: + - type: 'null' + - type: boolean + - type: number + - type: string + - type: array + - type: object + description: >- + (Optional) Metadata associated with the content, such as source, author, + etc. + chunk_metadata: + $ref: '#/components/schemas/ChunkMetadata' + description: >- + (Optional) Metadata associated with the chunk, such as document ID, source, + etc. additionalProperties: false required: - type diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index 3e8065cfb..436a96254 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -226,10 +226,18 @@ class VectorStoreContent(BaseModel): :param type: Content type, currently only "text" is supported :param text: The actual text content + :param embedding: (Optional) Embedding vector for the content, if available + :param created_timestamp: (Optional) Timestamp when the content was created + :param metadata: (Optional) Metadata associated with the content, such as source, author, etc. + :param chunk_metadata: (Optional) Metadata associated with the chunk, such as document ID, source, etc. """ type: Literal["text"] text: str + embedding: list[float] | None = None + created_timestamp: int | None = None + metadata: dict[str, Any] | None = None + chunk_metadata: ChunkMetadata | None = None @json_schema_type diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 120d0d4fc..7c7775691 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -17,6 +17,7 @@ from llama_stack.apis.files import Files, OpenAIFileObject from llama_stack.apis.vector_dbs import VectorDB from llama_stack.apis.vector_io import ( Chunk, + ChunkMetadata, QueryChunksResponse, SearchRankingOptions, VectorStoreChunkingStrategy, @@ -520,31 +521,68 @@ class OpenAIVectorStoreMixin(ABC): raise ValueError(f"Unsupported filter type: {filter_type}") def _chunk_to_vector_store_content(self, chunk: Chunk) -> list[VectorStoreContent]: + created_ts = None + if chunk.chunk_metadata is not None: + created_ts = getattr(chunk.chunk_metadata, "created_timestamp", None) + + metadata_dict = {} + if chunk.chunk_metadata: + if hasattr(chunk.chunk_metadata, "model_dump"): + metadata_dict = chunk.chunk_metadata.model_dump() + else: + metadata_dict = vars(chunk.chunk_metadata) + + user_metadata = chunk.metadata or {} + base_meta = {**metadata_dict, **user_metadata} + # content is InterleavedContent if isinstance(chunk.content, str): content = [ VectorStoreContent( type="text", text=chunk.content, + embedding=chunk.embedding, + created_timestamp=created_ts, + metadata=user_metadata, + chunk_metadata=ChunkMetadata(**base_meta) if base_meta else None, ) ] elif isinstance(chunk.content, list): # TODO: Add support for other types of content - content = [ - VectorStoreContent( - type="text", - text=item.text, - ) - for item in chunk.content - if item.type == "text" - ] + content = [] + for item in chunk.content: + if hasattr(item, "type") and item.type == "text": + item_meta = {**base_meta} + item_user_meta = getattr(item, "metadata", {}) or {} + if item_user_meta: + item_meta.update(item_user_meta) + + content.append( + VectorStoreContent( + type="text", + text=item.text, + embedding=getattr(item, "embedding", None), + created_timestamp=created_ts, + metadata=item_user_meta, + chunk_metadata=ChunkMetadata(**item_meta) if item_meta else None, + ) + ) else: - if chunk.content.type != "text": - raise ValueError(f"Unsupported content type: {chunk.content.type}") + content_item = chunk.content + if content_item.type != "text": + raise ValueError(f"Unsupported content type: {content_item.type}") + + item_user_meta = getattr(content_item, "metadata", {}) or {} + combined_meta = {**base_meta, **item_user_meta} + content = [ VectorStoreContent( type="text", - text=chunk.content.text, + text=content_item.text, + embedding=getattr(content_item, "embedding", None), + created_timestamp=created_ts, + metadata=item_user_meta, + chunk_metadata=ChunkMetadata(**combined_meta) if combined_meta else None, ) ] return content diff --git a/tests/integration/README.md b/tests/integration/README.md index 664116bea..7addd3eff 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -108,9 +108,7 @@ pytest -s -v tests/integration/inference/ \ Running Vector IO tests for a number of embedding models: ```bash -EMBEDDING_MODELS=all-MiniLM-L6-v2 - -pytest -s -v tests/integration/vector_io/ \ - --stack-config=inference=sentence-transformers,vector_io=sqlite-vec \ - --embedding-model=$EMBEDDING_MODELS +uv run pytest -sv --stack-config="inference=inline::sentence-transformers,vector_io=inline::sqlite-vec,files=localfs" \ +tests/integration/vector_io --embedding-model \ +sentence-transformers/all-MiniLM-L6-v2 ``` diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index 7ccca9077..defeead87 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -6,6 +6,7 @@ import logging import time +import uuid from io import BytesIO import pytest @@ -897,3 +898,76 @@ def test_openai_vector_store_search_modes(llama_stack_client, client_with_models search_mode=search_mode, ) assert search_response is not None + + +def test_openai_vector_store_file_contents_with_extended_fields(compat_client_with_empty_stores, client_with_models): + skip_if_provider_doesnt_support_openai_vector_stores(client_with_models) + + compat_client = compat_client_with_empty_stores + vector_store = compat_client.vector_stores.create( + name="extended_fields_test_store", metadata={"purpose": "extended_fields_testing"} + ) + + test_content = b"This is a test document." + file_name = f"extended_fields_test_{uuid.uuid4().hex}.txt" + attributes = {"test_type": "extended_fields", "version": "1.0"} + + with BytesIO(test_content) as file_buffer: + file_buffer.name = file_name + file = compat_client.files.create(file=file_buffer, purpose="assistants") + + file_attach_response = compat_client.vector_stores.files.create( + vector_store_id=vector_store.id, + file_id=file.id, + attributes=attributes, + ) + + assert file_attach_response.status == "completed", f"File attach failed: {file_attach_response.last_error}" + assert file_attach_response.attributes == attributes + + file_contents = compat_client.vector_stores.files.content( + vector_store_id=vector_store.id, + file_id=file.id, + ) + + assert file_contents + assert file_contents.filename == file_name + assert file_contents.attributes == attributes + assert len(file_contents.content) > 0 + + for content_item in file_contents.content: + if isinstance(compat_client, LlamaStackClient): + content_item = content_item.to_dict() + assert content_item["type"] == "text" + assert "text" in content_item + assert isinstance(content_item["text"], str) + assert len(content_item["text"]) > 0 + + if "embedding" in content_item: + assert isinstance(content_item["embedding"], list) + assert all(isinstance(x, (int | float)) for x in content_item["embedding"]) + + if "created_timestamp" in content_item: + assert isinstance(content_item["created_timestamp"], int) + assert content_item["created_timestamp"] > 0 + + if "chunk_metadata" in content_item: + assert isinstance(content_item["chunk_metadata"], dict) + if "chunk_id" in content_item["chunk_metadata"]: + assert isinstance(content_item["chunk_metadata"]["chunk_id"], str) + if "chunk_window" in content_item["chunk_metadata"]: + assert isinstance(content_item["chunk_metadata"]["chunk_window"], str) + + search_response = compat_client.vector_stores.search( + vector_store_id=vector_store.id, query="test document", max_num_results=5 + ) + + assert search_response is not None + assert len(search_response.data) > 0 + + for result_object in search_response.data: + result = result_object.to_dict() + assert "content" in result + assert len(result["content"]) > 0 + assert result["content"][0]["type"] == "text" + assert "text" in result["content"][0] diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index 98889f38e..249df3659 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -12,7 +12,7 @@ import numpy as np import pytest from llama_stack.apis.vector_dbs import VectorDB -from llama_stack.apis.vector_io import Chunk, QueryChunksResponse +from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse, VectorStoreContent from llama_stack.providers.remote.vector_io.milvus.milvus import VECTOR_DBS_PREFIX # This test is a unit test for the inline VectoerIO providers. This should only contain @@ -294,3 +294,35 @@ async def test_delete_openai_vector_store_file_from_storage(vector_io_adapter, t assert loaded_file_info == {} loaded_contents = await vector_io_adapter._load_openai_vector_store_file_contents(store_id, file_id) assert loaded_contents == [] + + +async def test_chunk_to_vector_store_content_with_new_fields(vector_io_adapter): + sample_chunk_metadata = ChunkMetadata( + chunk_id="chunk123", + document_id="doc456", + source="test_source", + created_timestamp=1625133600, + updated_timestamp=1625133600, + chunk_window="0-100", + chunk_tokenizer="test_tokenizer", + chunk_embedding_model="dummy_model", + chunk_embedding_dimension=384, + content_token_count=100, + metadata_token_count=100, + ) + + sample_chunk = Chunk( + content="hello world", metadata={"lang": "en"}, embedding=[0.5, 0.7, 0.9], chunk_metadata=sample_chunk_metadata + ) + + vsc_list: VectorStoreContent = vector_io_adapter._chunk_to_vector_store_content(sample_chunk) + assert isinstance(vsc_list, list) + assert len(vsc_list) > 0 + + vsc = vsc_list[0] + assert vsc.text == "hello world" + assert vsc.type == "text" + assert vsc.metadata == {"lang": "en"} + assert vsc.chunk_metadata == sample_chunk_metadata + assert vsc.embedding == [0.5, 0.7, 0.9] + assert vsc.created_timestamp == 1625133600