milvus: create TextNode if no '_node_content' set (#19772)

Author: RasmusAntons

llama-index-integrations/vector_stores/llama-index-vector-stores-milvus/llama_index/vector_stores/milvus/base.py

@@ -12,7 +12,7 @@
 
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.indices.query.embedding_utils import get_top_k_mmr_embeddings
-from llama_index.core.schema import BaseNode
+from llama_index.core.schema import BaseNode, TextNode
 from llama_index.core.utils import iter_batch
 from llama_index.vector_stores.milvus.utils import (
     get_default_sparse_embedding_function,
@@ -688,7 +688,17 @@ def get_nodes(
                     "The passed in text_key value does not exist "
                     "in the retrieved entity."
                 )
-            node = metadata_dict_to_node(item, text=text_content)
+            if "_node_content" in item:
+                node = metadata_dict_to_node(item, text=text_content)
+            elif text_content:
+                node = TextNode(
+                    text=text_content,
+                    metadata={key: item.get(key) for key in self.output_fields},
+                )
+            else:
+                raise ValueError(
+                    "Node content not found in metadata dict and no text content found."
+                )
             node.embedding = item.get(self.embedding_field, None)
             nodes.append(node)
         return nodes
@@ -721,7 +731,17 @@ async def aget_nodes(
                     "The passed in text_key value does not exist "
                     "in the retrieved entity."
                 )
-            node = metadata_dict_to_node(item, text=text_content)
+            if "_node_content" in item:
+                node = metadata_dict_to_node(item, text=text_content)
+            elif text_content:
+                node = TextNode(
+                    text=text_content,
+                    metadata={key: item.get(key) for key in self.output_fields},
+                )
+            else:
+                raise ValueError(
+                    "Node content not found in metadata dict and no text content found."
+                )
             node.embedding = item.get(self.embedding_field, None)
             nodes.append(node)
         return nodes
@@ -1412,13 +1432,18 @@ def _parse_from_milvus_results(
         ids = []
         # Parse the results
         for hit in results[0]:
-            metadata = {
-                "_node_content": hit["entity"].get("_node_content", None),
-                "_node_type": hit["entity"].get("_node_type", None),
-            }
-            for key in self.output_fields:
-                metadata[key] = hit["entity"].get(key)
-            node = metadata_dict_to_node(metadata)
+            if "_node_content" in hit["entity"]:
+                metadata = {
+                    "_node_content": hit["entity"].get("_node_content", None),
+                    "_node_type": hit["entity"].get("_node_type", None),
+                }
+                for key in self.output_fields:
+                    metadata[key] = hit["entity"].get(key)
+                node = metadata_dict_to_node(metadata)
+            else:
+                node = TextNode(
+                    metadata={key: hit["entity"].get(key) for key in self.output_fields}
+                )
 
             # Set the text field if it exists
             if self.text_key in hit["entity"]:

llama-index-integrations/vector_stores/llama-index-vector-stores-milvus/pyproject.toml

@@ -27,7 +27,7 @@ dev = [
 
 [project]
 name = "llama-index-vector-stores-milvus"
-version = "0.9.0"
+version = "0.9.1"
 description = "llama-index vector_stores milvus integration"
 authors = []
 requires-python = ">=3.9,<4.0"

llama-index-integrations/vector_stores/llama-index-vector-stores-milvus/tests/test_vector_stores_milvus.py

@@ -1,5 +1,6 @@
 from typing import List, Dict
 
+import pymilvus
 import pytest
 
 from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo
@@ -765,3 +766,56 @@ def test_query_hybrid_mode():
     assert len(result.nodes) == 1
     assert result.nodes[0].id_ == "n1"
     assert result.nodes[0].text == "n1_text"
+
+
+def test_custom_node_format():
+    milvus_db_uri = "./milvus_llamaindex_custom_format.db"
+    client = pymilvus.MilvusClient(
+        uri=milvus_db_uri,
+    )
+    if COLLECTION_NAME in client.list_collections():
+        client.drop_collection(COLLECTION_NAME)
+    client.create_collection(
+        collection_name=COLLECTION_NAME,
+        dimension=DIM,
+        primary_field_name="id",
+        vector_field_name="embedding",
+        metric_type="COSINE",
+        auto_id=True,
+    )
+    vector_store = MilvusVectorStore(
+        uri=milvus_db_uri,
+        collection_name=COLLECTION_NAME,
+        dim=DIM,
+        doc_id_field="id",
+        embedding_field="embedding",
+        text_key="custom_text",
+        output_fields=["custom_meta"],
+    )
+    vector_store.client.insert(
+        COLLECTION_NAME,
+        data=[
+            {
+                "embedding": [0.5] * DIM,
+                "custom_text": "n1_text",
+                "custom_meta": "n1_meta",
+            },
+            {
+                "embedding": [-0.5] * DIM,
+                "custom_text": "n2_text",
+                "custom_meta": "n2_meta",
+            },
+        ],
+    )
+
+    query = VectorStoreQuery(query_embedding=[0.5] * DIM, similarity_top_k=1)
+    result = vector_store.query(query=query)
+    assert len(result.nodes) == 1
+    assert result.nodes[0].text == "n1_text"
+    assert result.nodes[0].metadata.get("custom_meta") == "n1_meta"
+
+    query = VectorStoreQuery(query_embedding=[-0.5] * DIM, similarity_top_k=1)
+    result = vector_store.query(query=query)
+    assert len(result.nodes) == 1
+    assert result.nodes[0].text == "n2_text"
+    assert result.nodes[0].metadata.get("custom_meta") == "n2_meta"