diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/schema/vespa/vespa_node.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/schema/vespa/vespa_node.py new file mode 100644 index 0000000000000..892c366e0bd25 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/schema/vespa/vespa_node.py @@ -0,0 +1,5 @@ +from llama_index.core.schema import TextNode + + +class VespaNode(TextNode): + vespa_fields: dict \ No newline at end of file diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py index 74ec8d657dc06..6a280049ff0bd 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/base.py @@ -15,6 +15,7 @@ metadata_dict_to_node, ) +from llama_index.schema.vespa.vespa_node import VespaNode from llama_index.vector_stores.vespa.templates import hybrid_template import asyncio @@ -157,7 +158,14 @@ def client(self) -> Vespa: def _try_get_running_app(self) -> Vespa: app = Vespa(url=f"{self.url}:{self.port}") + status = app.get_application_status() + try: + status.status_code + except AttributeError: + raise ConnectionError( + f"Vespa application not running on url {self.url} and port {self.port}. Please start Vespa application first." + ) if status.status_code == 200: return app else: @@ -213,14 +221,11 @@ def add( node, remove_text=False, flat_metadata=self.flat_metadata ) logger.debug(f"Metadata: {metadata}") - entry = { - "id": node.node_id, - "fields": { - "id": node.node_id, - "text": node.get_content(metadata_mode=MetadataMode.NONE) or "", - "metadata": json.dumps(metadata), - }, - } + + if isinstance(node, VespaNode): + entry = self._get_vespa_node_entry(node) + else: + entry = self._get_base_node_entry(node, metadata) if self.embeddings_outside_vespa: entry["fields"]["embedding"] = node.get_embedding() data_to_insert.append(entry) @@ -235,6 +240,28 @@ def add( ) return ids + def _get_base_node_entry(self, node: BaseNode, metadata:dict): + entry = { + "id": node.node_id, + "fields": { + "id": node.node_id, + "text": node.get_content(metadata_mode=MetadataMode.NONE) or "", + "metadata": json.dumps(metadata), + }, + } + return entry + def _get_vespa_node_entry(self,node: VespaNode): + vespa_fields = node.vespa_fields + metadata = node_to_metadata_dict( + node, remove_text=False, flat_metadata=self.flat_metadata + ) + vespa_fields.update({ + "text": node.get_content(metadata_mode=MetadataMode.NONE) or "", + "metadata": json.dumps(metadata), + }) + entry = {"id": node.node_id, "fields": vespa_fields} + return entry + async def async_add( self, nodes: List[BaseNode], @@ -465,18 +492,38 @@ def query( ids: List[str] = [] similarities: List[float] = [] for hit in response.hits: - response_fields: dict = hit.get("fields", {}) - metadata = response_fields.get("metadata", {}) - metadata = json.loads(metadata) - logger.debug(f"Metadata: {metadata}") - node = metadata_dict_to_node(metadata) - text = response_fields.get("body", "") - node.set_content(text) + node = self._vespa_hit_to_node(hit) nodes.append(node) - ids.append(response_fields.get("id")) + id = hit["fields"].get("id") + ids.append(id) similarities.append(hit["relevance"]) return VectorStoreQueryResult(nodes=nodes, ids=ids, similarities=similarities) + def _vespa_hit_to_node(self, hit: dict) -> BaseNode | VespaNode: + response_fields = hit.get("fields", {}) + if self._is_vespa_node(response_fields): + node = self._get_vespa_node_from_fields(response_fields) + else: + node = self._get_base_node_from_fields(response_fields) + text = response_fields.get("body", "") + node.set_content(text) + return node + + def _is_vespa_node(self, fields: dict): + # Check if there are more fields than text and metadata + base_node_fields = ["text", "metadata"] + return len([f for f in fields if f not in base_node_fields]) > 0 + + def _get_vespa_node_from_fields(self, fields: dict): + vespa_fields = fields + return VespaNode(vespa_fields=vespa_fields) + + def _get_base_node_from_fields(self, fields: dict): + metadata = fields.get("metadata", {}) + node = metadata_dict_to_node(metadata) + return node + + async def aquery( self, query: VectorStoreQuery, diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py index 18570bd9e620c..9a123ca89edf2 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/llama_index/vector_stores/vespa/templates.py @@ -94,3 +94,87 @@ ) ], ) + +# I am not sure what to name this, so I decided a name that highlights the difference with hybrid_template +# This needs to be renamed to something more meaningful +with_fields_template = ApplicationPackage( + name="withfields", + schema=[ + Schema( + name="doc", + document=Document( + fields=[ + Field(name="id", type="string", indexing=["summary"]), + Field(name="title", type="string", indexing=["summary"]), + Field(name="author", type="string", indexing=["summary"]), + Field(name="theme", type="string", indexing=["summary"]), + Field(name="year", type="int", indexing=["summary"]), + Field(name="metadata", type="string", indexing=["summary"]), + Field( + name="text", + type="string", + indexing=["index", "summary"], + index="enable-bm25", + bolding=True, + ), + Field( + name="embedding", + type="tensor(x[384])", + indexing=[ + "input text", + "embed", + "index", + "attribute", + ], + ann=HNSW(distance_metric="angular"), + is_document_field=False, + ), + ] + ), + fieldsets=[FieldSet(name="default", fields=["text", "metadata"])], + rank_profiles=[ + RankProfile( + name="bm25", + inputs=[("query(q)", "tensor(x[384])")], + functions=[Function(name="bm25sum", expression="bm25(text)")], + first_phase="bm25sum", + ), + RankProfile( + name="semantic", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + ), + RankProfile( + name="fusion", + inherits="bm25", + inputs=[("query(q)", "tensor(x[384])")], + first_phase="closeness(field, embedding)", + global_phase=GlobalPhaseRanking( + expression="reciprocal_rank_fusion(bm25sum, closeness(field, embedding))", + rerank_count=1000, + ), + ), + ], + ) + ], + components=[ + Component( + id="e5", + type="hugging-face-embedder", + parameters=[ + Parameter( + "transformer-model", + { + "url": "https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx" + }, + ), + Parameter( + "tokenizer-model", + { + "url": "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json" + }, + ), + ], + ) + ], +) diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespa_vector_store_with_vespa_nodes.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespa_vector_store_with_vespa_nodes.py new file mode 100644 index 0000000000000..993f3ed1039e5 --- /dev/null +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespa_vector_store_with_vespa_nodes.py @@ -0,0 +1,198 @@ +import pytest +from llama_index.core.vector_stores import VectorStoreQuery +from llama_index.core.vector_stores.types import VectorStoreQueryMode +from vespa.package import ApplicationPackage + +from llama_index.schema.vespa.vespa_node import VespaNode +from llama_index.vector_stores.vespa import VespaVectorStore +from llama_index.vector_stores.vespa.templates import with_fields_template + + +try: + # Should be installed as pyvespa-dependency + import docker + + client = docker.from_env() + docker_available = client.ping() +except Exception: + docker_available = False + +@pytest.fixture(scope="session") +def vespa_app(): + app_package: ApplicationPackage = with_fields_template + try: + # Try getting the local instance if available + return VespaVectorStore(url="http://localhost", application_package=app_package, deployment_target="local") + except ConnectionError: + return VespaVectorStore(application_package=app_package, deployment_target="local") +@pytest.fixture(scope="session") +def vespa_nodes() -> list: + return [ + VespaNode( + vespa_fields={ + "title": "The Shawshank Redemption", + "author": "Stephen King", + "theme": "Friendship", + "year": 1994, + "id": "1" + }, + text="The Shawshank Redemption", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "The Godfather", + "director": "Francis Ford Coppola", + "theme": "Mafia", + "year": 1972, + "id": "2" + }, + text="The Godfather", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "Inception", + "director": "Christopher Nolan", + "theme": "Fiction", + "year": 2010, + "id": "3" + }, + text="Inception", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "To Kill a Mockingbird", + "author": "Harper Lee", + "theme": "Mafia", + "year": 1960, + "id": 4 + }, + text="To Kill a Mockingbird", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "1984", + "author": "George Orwell", + "theme": "Totalitarianism", + "year": 1949, + "id": "5" + }, + text="1984", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "The Great Gatsby", + "author": "F. Scott Fitzgerald", + "theme": "The American Dream", + "year": 1925, + "id": "6" + }, + text="The Great Gatsby", + metadata={"added_by": "gokturkDev"}, + ), + VespaNode( + vespa_fields={ + "title": "Harry Potter and the Sorcerer's Stone", + "author": "J.K. Rowling", + "theme": "Fiction", + "year": 1997, + "id": "7" + }, + text="Harry Potter and the Sorcerer's Stone", + metadata={"added_by": "gokturkDev"}, + ), + + ] + +@pytest.fixture() +def added_node_ids(vespa_app, vespa_nodes): + yield vespa_app.add(vespa_nodes) + for node in vespa_nodes: + vespa_app.delete(node.node_id) + +class TestTextQuery: + def setup_method(self): + self.text_query = VectorStoreQuery( + query_str="1984", # Ensure the query matches the case used in the nodes + mode=VectorStoreQueryMode.TEXT_SEARCH, + similarity_top_k=1, + ) + @pytest.mark.skipif(not docker_available, reason="Docker not available") + def test_returns_hit(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + assert len(result.nodes) == 1 + + def test_returns_vespa_node(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + node = result.nodes[0] + assert isinstance(node, VespaNode) + def test_correctly_assigns_vespa_node_fields(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + node = result.nodes[0] + assert node.vespa_fields["title"] == "1984" + assert node.vespa_fields["author"] == "George Orwell" + assert node.vespa_fields["theme"] == "Totalitarianism" + assert node.vespa_fields["year"] == 1949 + assert node.vespa_fields["id"] == "5" + +class TestSemanticQuery: + def setup_method(self): + self.text_query = VectorStoreQuery( + query_str="1984", # Ensure the query matches the case used in the nodes + mode=VectorStoreQueryMode.SEMANTIC_HYBRID, + similarity_top_k=1, + ) + @pytest.mark.skipif(not docker_available, reason="Docker not available") + def test_returns_hit(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + assert len(result.nodes) == 1 + + def test_returns_vespa_node(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + node = result.nodes[0] + assert isinstance(node, VespaNode) + def test_correctly_assigns_vespa_node_fields(self, vespa_app, added_node_ids): + result = vespa_app.query(self.text_query) + node = result.nodes[0] + assert node.vespa_fields["title"] == "1984" + assert node.vespa_fields["author"] == "George Orwell" + assert node.vespa_fields["theme"] == "Totalitarianism" + assert node.vespa_fields["year"] == 1949 + assert node.vespa_fields["id"] == "5" + + +class TestDeleteNode: + def setup_method(self): + self.test_node = VespaNode( + vespa_fields={ + "title": "Clean Code: A Handbook of Agile Software Craftsmanship", + "author": "Robert C Martin", + "theme": "Software Development", + "year": 2008, + "id": "8" + }, + text="Clean Code: A Handbook of Agile Software Craftsmanship ", + metadata={"added_by": "gokturkDev"}, + ) + + @pytest.mark.skipif(not docker_available, reason="Docker not available") + def test_deletes_node(self, vespa_app): + vespa_app.add([self.test_node]) + query = VectorStoreQuery( + query_str="Clean Code: A Handbook of Agile Software Craftsmanship", # Ensure the query matches the case used in the nodes + mode=VectorStoreQueryMode.TEXT_SEARCH, + similarity_top_k=1, + ) + result = vespa_app.query(query) + assert len(result.nodes) == 1 + + vespa_app.delete(self.test_node.node_id) + result = vespa_app.query(query) + assert len(result.nodes) == 0 + + + diff --git a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py index 44bbc6b90d5b4..d980c8c10f5bb 100644 --- a/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py +++ b/llama-index-integrations/vector_stores/llama-index-vector-stores-vespa/tests/test_vespavectorstore.py @@ -23,7 +23,11 @@ @pytest.fixture(scope="session") def vespa_app(): app_package: ApplicationPackage = hybrid_template - return VespaVectorStore(application_package=app_package, deployment_target="local") + try: + # Try getting the local instance if available + return VespaVectorStore(url="http://localhost", application_package=app_package, deployment_target="local") + except ConnectionError: + return VespaVectorStore(application_package=app_package, deployment_target="local") @pytest.fixture(scope="session") @@ -105,7 +109,7 @@ def added_node_ids(vespa_app, nodes): def test_query_text_search(vespa_app, added_node_ids): query = VectorStoreQuery( query_str="Inception", # Ensure the query matches the case used in the nodes - mode="text_search", + mode=VectorStoreQueryMode.TEXT_SEARCH, similarity_top_k=1, ) result = vespa_app.query(query) @@ -118,7 +122,7 @@ def test_query_text_search(vespa_app, added_node_ids): def test_query_vector_search(vespa_app, added_node_ids): query = VectorStoreQuery( query_str="magic, wizardry", - mode="semantic_hybrid", + mode=VectorStoreQueryMode.SEMANTIC_HYBRID, similarity_top_k=1, ) result = vespa_app.query(query) @@ -130,6 +134,7 @@ def test_query_vector_search(vespa_app, added_node_ids): @pytest.mark.skipif(not docker_available, reason="Docker not available") def test_delete_node(vespa_app, added_node_ids): + print("added node ids: ", added_node_ids) # Testing the deletion of a node vespa_app.delete(ref_doc_id=added_node_ids[1]) query = VectorStoreQuery(