diff --git a/graphrag/config/models/vector_store_schema_config.py b/graphrag/config/models/vector_store_schema_config.py index 2da0cdc17a..ccb91b3cbb 100644 --- a/graphrag/config/models/vector_store_schema_config.py +++ b/graphrag/config/models/vector_store_schema_config.py @@ -30,16 +30,6 @@ class VectorStoreSchemaConfig(BaseModel): default="vector", ) - text_field: str = Field( - description="The text field to use.", - default="text", - ) - - attributes_field: str = Field( - description="The attributes field to use.", - default="attributes", - ) - vector_size: int = Field( description="The vector size to use.", default=DEFAULT_VECTOR_SIZE, @@ -52,8 +42,6 @@ def _validate_schema(self) -> None: for field in [ self.id_field, self.vector_field, - self.text_field, - self.attributes_field, ]: if not is_valid_field_name(field): msg = f"Unsafe or invalid field name: {field}" diff --git a/graphrag/index/operations/embed_text/embed_text.py b/graphrag/index/operations/embed_text/embed_text.py index c1a4b796e3..46b81ce212 100644 --- a/graphrag/index/operations/embed_text/embed_text.py +++ b/graphrag/index/operations/embed_text/embed_text.py @@ -152,7 +152,6 @@ async def _text_embed_with_vector_store( ) batch = input.iloc[insert_batch_size * i : insert_batch_size * (i + 1)] texts: list[str] = batch[embed_column].tolist() - titles: list[str] = batch[title].tolist() ids: list[str] = batch[id_column].tolist() result = await strategy_exec(texts, callbacks, cache, strategy_config) if result.embeddings: @@ -163,16 +162,12 @@ async def _text_embed_with_vector_store( vectors = result.embeddings or [] documents: list[VectorStoreDocument] = [] - for doc_id, doc_text, doc_title, doc_vector in zip( - ids, texts, titles, vectors, strict=True - ): + for doc_id, doc_vector in zip(ids, vectors, strict=True): if type(doc_vector) is np.ndarray: doc_vector = doc_vector.tolist() document = VectorStoreDocument( id=doc_id, - text=doc_text, vector=doc_vector, - attributes={"title": doc_title}, ) documents.append(document) diff --git a/graphrag/query/structured_search/basic_search/basic_context.py b/graphrag/query/structured_search/basic_search/basic_context.py index 8296de3c89..215b26a31e 100644 --- a/graphrag/query/structured_search/basic_search/basic_context.py +++ b/graphrag/query/structured_search/basic_search/basic_context.py @@ -59,14 +59,15 @@ def build_context( text_embedder=lambda t: self.text_embedder.embed(t), k=k, ) - related_text_list = [ - { - text_id_col: self.text_id_map[f"{chunk.document.id}"], - text_col: chunk.document.text, - } - for chunk in related_texts + + text_unit_ids = {t.document.id for t in related_texts} + text_units_filtered = [] + text_units_filtered = [ + {text_id_col: t.id, text_col: t.text} + for t in self.text_units or [] + if t.id in text_unit_ids ] - related_text_df = pd.DataFrame(related_text_list) + related_text_df = pd.DataFrame(text_units_filtered) else: related_text_df = pd.DataFrame({ text_id_col: [], diff --git a/graphrag/vector_stores/azure_ai_search.py b/graphrag/vector_stores/azure_ai_search.py index 037cc28754..c1eb063d64 100644 --- a/graphrag/vector_stores/azure_ai_search.py +++ b/graphrag/vector_stores/azure_ai_search.py @@ -3,7 +3,6 @@ """A package containing the Azure AI Search vector store implementation.""" -import json from typing import Any from azure.core.credentials import AzureKeyCredential @@ -13,7 +12,6 @@ from azure.search.documents.indexes.models import ( HnswAlgorithmConfiguration, HnswParameters, - SearchableField, SearchField, SearchFieldDataType, SearchIndex, @@ -121,13 +119,6 @@ def load_documents( vector_search_dimensions=self.vector_size, vector_search_profile_name=self.vector_search_profile_name, ), - SearchableField( - name=self.text_field, type=SearchFieldDataType.String - ), - SimpleField( - name=self.attributes_field, - type=SearchFieldDataType.String, - ), ], vector_search=vector_search, ) @@ -139,8 +130,6 @@ def load_documents( { self.id_field: doc.id, self.vector_field: doc.vector, - self.text_field: doc.text, - self.attributes_field: json.dumps(doc.attributes), } for doc in documents if doc.vector is not None @@ -165,9 +154,7 @@ def similarity_search_by_vector( VectorStoreSearchResult( document=VectorStoreDocument( id=doc.get(self.id_field, ""), - text=doc.get(self.text_field, ""), vector=doc.get(self.vector_field, []), - attributes=(json.loads(doc.get(self.attributes_field, "{}"))), ), # Cosine similarity between 0.333 and 1.000 # https://learn.microsoft.com/en-us/azure/search/hybrid-search-ranking#scores-in-a-hybrid-search-results @@ -192,7 +179,5 @@ def search_by_id(self, id: str) -> VectorStoreDocument: response = self.db_connection.get_document(id) return VectorStoreDocument( id=response.get(self.id_field, ""), - text=response.get(self.text_field, ""), vector=response.get(self.vector_field, []), - attributes=(json.loads(response.get(self.attributes_field, "{}"))), ) diff --git a/graphrag/vector_stores/base.py b/graphrag/vector_stores/base.py index 2008fb472d..1bee5910af 100644 --- a/graphrag/vector_stores/base.py +++ b/graphrag/vector_stores/base.py @@ -4,7 +4,7 @@ """Base classes for vector stores.""" from abc import ABC, abstractmethod -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig @@ -18,12 +18,8 @@ class VectorStoreDocument: id: str | int """unique id for the document""" - text: str | None vector: list[float] | None - attributes: dict[str, Any] = field(default_factory=dict) - """store any additional metadata, e.g. title, date ranges, etc""" - @dataclass class VectorStoreSearchResult: @@ -54,9 +50,7 @@ def __init__( self.index_name = vector_store_schema_config.index_name self.id_field = vector_store_schema_config.id_field - self.text_field = vector_store_schema_config.text_field self.vector_field = vector_store_schema_config.vector_field - self.attributes_field = vector_store_schema_config.attributes_field self.vector_size = vector_store_schema_config.vector_size @abstractmethod diff --git a/graphrag/vector_stores/cosmosdb.py b/graphrag/vector_stores/cosmosdb.py index 186412d6ba..c609aa30b2 100644 --- a/graphrag/vector_stores/cosmosdb.py +++ b/graphrag/vector_stores/cosmosdb.py @@ -3,7 +3,6 @@ """A package containing the CosmosDB vector store implementation.""" -import json from typing import Any from azure.cosmos import ContainerProxy, CosmosClient, DatabaseProxy @@ -171,8 +170,6 @@ def load_documents( doc_json = { self.id_field: doc.id, self.vector_field: doc.vector, - self.text_field: doc.text, - self.attributes_field: json.dumps(doc.attributes), } print("Storing document in CosmosDB:") # noqa: T201 print(doc_json) # noqa: T201 @@ -187,7 +184,7 @@ def similarity_search_by_vector( raise ValueError(msg) try: - query = f"SELECT TOP {k} c.{self.id_field}, c.{self.text_field}, c.{self.vector_field}, c.{self.attributes_field}, VectorDistance(c.{self.vector_field}, @embedding) AS SimilarityScore FROM c ORDER BY VectorDistance(c.{self.vector_field}, @embedding)" # noqa: S608 + query = f"SELECT TOP {k} c.{self.id_field}, c.{self.vector_field}, VectorDistance(c.{self.vector_field}, @embedding) AS SimilarityScore FROM c ORDER BY VectorDistance(c.{self.vector_field}, @embedding)" # noqa: S608 query_params = [{"name": "@embedding", "value": query_embedding}] items = list( self._container_client.query_items( @@ -199,7 +196,7 @@ def similarity_search_by_vector( except (CosmosHttpResponseError, ValueError): # Currently, the CosmosDB emulator does not support the VectorDistance function. # For emulator or test environments - fetch all items and calculate distance locally - query = f"SELECT c.{self.id_field}, c.{self.text_field}, c.{self.vector_field}, c.{self.attributes_field} FROM c" # noqa: S608 + query = f"SELECT c.{self.id_field}, c.{self.vector_field} FROM c" # noqa: S608 items = list( self._container_client.query_items( query=query, @@ -231,9 +228,7 @@ def cosine_similarity(a, b): VectorStoreSearchResult( document=VectorStoreDocument( id=item.get(self.id_field, ""), - text=item.get(self.text_field, ""), vector=item.get(self.vector_field, []), - attributes=(json.loads(item.get(self.attributes_field, "{}"))), ), score=item.get("SimilarityScore", 0.0), ) @@ -261,8 +256,6 @@ def search_by_id(self, id: str) -> VectorStoreDocument: return VectorStoreDocument( id=item.get(self.id_field, ""), vector=item.get(self.vector_field, []), - text=item.get(self.text_field, ""), - attributes=(json.loads(item.get(self.attributes_field, "{}"))), ) def clear(self) -> None: diff --git a/graphrag/vector_stores/lancedb.py b/graphrag/vector_stores/lancedb.py index 7c03d97386..fe53ca7182 100644 --- a/graphrag/vector_stores/lancedb.py +++ b/graphrag/vector_stores/lancedb.py @@ -3,19 +3,19 @@ """The LanceDB vector storage implementation package.""" -import json # noqa: I001 from typing import Any -import pyarrow as pa + +import lancedb import numpy as np +import pyarrow as pa + from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig from graphrag.data_model.types import TextEmbedder - from graphrag.vector_stores.base import ( BaseVectorStore, VectorStoreDocument, VectorStoreSearchResult, ) -import lancedb class LanceDBVectorStore(BaseVectorStore): @@ -41,9 +41,7 @@ def load_documents( """Load documents into vector storage.""" # Step 1: Prepare data columns manually ids = [] - texts = [] vectors = [] - attributes = [] for document in documents: self.vector_size = ( @@ -51,9 +49,7 @@ def load_documents( ) if document.vector is not None and len(document.vector) == self.vector_size: ids.append(document.id) - texts.append(document.text) vectors.append(np.array(document.vector, dtype=np.float32)) - attributes.append(json.dumps(document.attributes)) # Step 2: Handle empty case if len(ids) == 0: @@ -69,9 +65,7 @@ def load_documents( # Step 4: Create PyArrow table (let schema be inferred) data = pa.table({ self.id_field: pa.array(ids, type=pa.string()), - self.text_field: pa.array(texts, type=pa.string()), self.vector_field: vector_column, - self.attributes_field: pa.array(attributes, type=pa.string()), }) # NOTE: If modifying the next section of code, ensure that the schema remains the same. @@ -127,9 +121,7 @@ def similarity_search_by_vector( VectorStoreSearchResult( document=VectorStoreDocument( id=doc[self.id_field], - text=doc[self.text_field], vector=doc[self.vector_field], - attributes=json.loads(doc[self.attributes_field]), ), score=1 - abs(float(doc["_distance"])), ) @@ -155,8 +147,6 @@ def search_by_id(self, id: str) -> VectorStoreDocument: if doc: return VectorStoreDocument( id=doc[0][self.id_field], - text=doc[0][self.text_field], vector=doc[0][self.vector_field], - attributes=json.loads(doc[0][self.attributes_field]), ) - return VectorStoreDocument(id=id, text=None, vector=None) + return VectorStoreDocument(id=id, vector=None) diff --git a/tests/integration/vector_stores/test_azure_ai_search.py b/tests/integration/vector_stores/test_azure_ai_search.py index 2d67c62ea8..aaf6f00374 100644 --- a/tests/integration/vector_stores/test_azure_ai_search.py +++ b/tests/integration/vector_stores/test_azure_ai_search.py @@ -63,8 +63,6 @@ def vector_store_custom(self, mock_search_client, mock_index_client): vector_store_schema_config=VectorStoreSchemaConfig( index_name="test_vectors", id_field="id_custom", - text_field="text_custom", - attributes_field="attributes_custom", vector_field="vector_custom", vector_size=5, ), @@ -86,15 +84,11 @@ def sample_documents(self): return [ VectorStoreDocument( id="doc1", - text="This is document 1", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Doc 1", "category": "test"}, ), VectorStoreDocument( id="doc2", - text="This is document 2", vector=[0.2, 0.3, 0.4, 0.5, 0.6], - attributes={"title": "Doc 2", "category": "test"}, ), ] @@ -110,16 +104,12 @@ async def test_vector_store_operations( search_results = [ { "id": "doc1", - "text": "This is document 1", "vector": [0.1, 0.2, 0.3, 0.4, 0.5], - "attributes": '{"title": "Doc 1", "category": "test"}', "@search.score": 0.9, }, { "id": "doc2", - "text": "This is document 2", "vector": [0.2, 0.3, 0.4, 0.5, 0.6], - "attributes": '{"title": "Doc 2", "category": "test"}', "@search.score": 0.8, }, ] @@ -127,9 +117,7 @@ async def test_vector_store_operations( mock_search_client.get_document.return_value = { "id": "doc1", - "text": "This is document 1", "vector": [0.1, 0.2, 0.3, 0.4, 0.5], - "attributes": '{"title": "Doc 1", "category": "test"}', } vector_store.load_documents(sample_documents) @@ -154,8 +142,6 @@ def mock_embedder(text: str) -> list[float]: doc = vector_store.search_by_id("doc1") assert doc.id == "doc1" - assert doc.text == "This is document 1" - assert doc.attributes["title"] == "Doc 1" async def test_empty_embedding(self, vector_store, mock_search_client): """Test similarity search by text with empty embedding.""" @@ -186,16 +172,12 @@ async def test_vector_store_customization( search_results = [ { vector_store_custom.id_field: "doc1", - vector_store_custom.text_field: "This is document 1", vector_store_custom.vector_field: [0.1, 0.2, 0.3, 0.4, 0.5], - vector_store_custom.attributes_field: '{"title": "Doc 1", "category": "test"}', "@search.score": 0.9, }, { vector_store_custom.id_field: "doc2", - vector_store_custom.text_field: "This is document 2", vector_store_custom.vector_field: [0.2, 0.3, 0.4, 0.5, 0.6], - vector_store_custom.attributes_field: '{"title": "Doc 2", "category": "test"}', "@search.score": 0.8, }, ] @@ -203,9 +185,7 @@ async def test_vector_store_customization( mock_search_client.get_document.return_value = { vector_store_custom.id_field: "doc1", - vector_store_custom.text_field: "This is document 1", vector_store_custom.vector_field: [0.1, 0.2, 0.3, 0.4, 0.5], - vector_store_custom.attributes_field: '{"title": "Doc 1", "category": "test"}', } vector_store_custom.load_documents(sample_documents) @@ -230,5 +210,3 @@ def mock_embedder(text: str) -> list[float]: doc = vector_store_custom.search_by_id("doc1") assert doc.id == "doc1" - assert doc.text == "This is document 1" - assert doc.attributes["title"] == "Doc 1" diff --git a/tests/integration/vector_stores/test_cosmosdb.py b/tests/integration/vector_stores/test_cosmosdb.py index f9515beed9..364a3d85d8 100644 --- a/tests/integration/vector_stores/test_cosmosdb.py +++ b/tests/integration/vector_stores/test_cosmosdb.py @@ -37,25 +37,19 @@ def test_vector_store_operations(): docs = [ VectorStoreDocument( id="doc1", - text="This is document 1", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Doc 1", "category": "test"}, ), VectorStoreDocument( id="doc2", - text="This is document 2", vector=[0.2, 0.3, 0.4, 0.5, 0.6], - attributes={"title": "Doc 2", "category": "test"}, ), ] vector_store.load_documents(docs) doc = vector_store.search_by_id("doc1") assert doc.id == "doc1" - assert doc.text == "This is document 1" assert doc.vector is not None assert np.allclose(doc.vector, [0.1, 0.2, 0.3, 0.4, 0.5]) - assert doc.attributes["title"] == "Doc 1" # Define a simple text embedder function for testing def mock_embedder(text: str) -> list[float]: @@ -87,9 +81,7 @@ def test_clear(): doc = VectorStoreDocument( id="test", - text="Test document", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Test Doc"}, ) vector_store.load_documents([doc]) @@ -109,9 +101,7 @@ def test_vector_store_customization(): vector_store_schema_config=VectorStoreSchemaConfig( index_name="text-embeddings", id_field="id", - text_field="text_custom", vector_field="vector_custom", - attributes_field="attributes_custom", vector_size=5, ), ) @@ -125,25 +115,19 @@ def test_vector_store_customization(): docs = [ VectorStoreDocument( id="doc1", - text="This is document 1", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Doc 1", "category": "test"}, ), VectorStoreDocument( id="doc2", - text="This is document 2", vector=[0.2, 0.3, 0.4, 0.5, 0.6], - attributes={"title": "Doc 2", "category": "test"}, ), ] vector_store.load_documents(docs) doc = vector_store.search_by_id("doc1") assert doc.id == "doc1" - assert doc.text == "This is document 1" assert doc.vector is not None assert np.allclose(doc.vector, [0.1, 0.2, 0.3, 0.4, 0.5]) - assert doc.attributes["title"] == "Doc 1" # Define a simple text embedder function for testing def mock_embedder(text: str) -> list[float]: diff --git a/tests/integration/vector_stores/test_factory.py b/tests/integration/vector_stores/test_factory.py index 724c1f8e2d..de819a6ab0 100644 --- a/tests/integration/vector_stores/test_factory.py +++ b/tests/integration/vector_stores/test_factory.py @@ -146,7 +146,7 @@ def filter_by_id(self, include_ids): def search_by_id(self, id): from graphrag.vector_stores.base import VectorStoreDocument - return VectorStoreDocument(id=id, text="test", vector=None) + return VectorStoreDocument(id=id, vector=None) # VectorStoreFactory allows registering classes directly (no TypeError) VectorStoreFactory.register("custom_class", CustomVectorStore) diff --git a/tests/integration/vector_stores/test_lancedb.py b/tests/integration/vector_stores/test_lancedb.py index cca48a1271..d0fb19eaaa 100644 --- a/tests/integration/vector_stores/test_lancedb.py +++ b/tests/integration/vector_stores/test_lancedb.py @@ -23,21 +23,15 @@ def sample_documents(self): return [ VectorStoreDocument( id="1", - text="This is document 1", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Doc 1", "category": "test"}, ), VectorStoreDocument( id="2", - text="This is document 2", vector=[0.2, 0.3, 0.4, 0.5, 0.6], - attributes={"title": "Doc 2", "category": "test"}, ), VectorStoreDocument( id="3", - text="This is document 3", vector=[0.3, 0.4, 0.5, 0.6, 0.7], - attributes={"title": "Doc 3", "category": "test"}, ), ] @@ -47,21 +41,15 @@ def sample_documents_categories(self): return [ VectorStoreDocument( id="1", - text="Document about cats", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"category": "animals"}, ), VectorStoreDocument( id="2", - text="Document about dogs", vector=[0.2, 0.3, 0.4, 0.5, 0.6], - attributes={"category": "animals"}, ), VectorStoreDocument( id="3", - text="Document about cars", vector=[0.3, 0.4, 0.5, 0.6, 0.7], - attributes={"category": "vehicles"}, ), ] @@ -85,11 +73,8 @@ def test_vector_store_operations(self, sample_documents): doc = vector_store.search_by_id("1") assert doc.id == "1" - assert doc.text == "This is document 1" - assert doc.vector is not None assert np.allclose(doc.vector, [0.1, 0.2, 0.3, 0.4, 0.5]) - assert doc.attributes["title"] == "Doc 1" results = vector_store.similarity_search_by_vector( [0.1, 0.2, 0.3, 0.4, 0.5], k=2 @@ -101,7 +86,6 @@ def test_vector_store_operations(self, sample_documents): vector_store.load_documents([sample_documents[2]], overwrite=False) result = vector_store.search_by_id("3") assert result.id == "3" - assert result.text == "This is document 3" # Define a simple text embedder function for testing def mock_embedder(text: str) -> list[float]: @@ -116,7 +100,6 @@ def mock_embedder(text: str) -> list[float]: # Test non-existent document non_existent = vector_store.search_by_id("nonexistent") assert non_existent.id == "nonexistent" - assert non_existent.text is None assert non_existent.vector is None finally: shutil.rmtree(temp_dir) @@ -136,9 +119,7 @@ def test_empty_collection(self): # Load the vector store with a document, then delete it sample_doc = VectorStoreDocument( id="tmp", - text="Temporary document to create schema", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Tmp"}, ) vector_store.load_documents([sample_doc]) vector_store.db_connection.open_table( @@ -154,15 +135,12 @@ def test_empty_collection(self): # Add a document after creating an empty collection doc = VectorStoreDocument( id="1", - text="This is document 1", vector=[0.1, 0.2, 0.3, 0.4, 0.5], - attributes={"title": "Doc 1"}, ) vector_store.load_documents([doc], overwrite=False) result = vector_store.search_by_id("1") assert result.id == "1" - assert result.text == "This is document 1" finally: # Clean up - remove the temporary directory shutil.rmtree(temp_dir) @@ -203,9 +181,7 @@ def test_vector_store_customization(self, sample_documents): vector_store_schema_config=VectorStoreSchemaConfig( index_name="text-embeddings", id_field="id_custom", - text_field="text_custom", vector_field="vector_custom", - attributes_field="attributes_custom", vector_size=5, ), ) @@ -219,11 +195,8 @@ def test_vector_store_customization(self, sample_documents): doc = vector_store.search_by_id("1") assert doc.id == "1" - assert doc.text == "This is document 1" - assert doc.vector is not None assert np.allclose(doc.vector, [0.1, 0.2, 0.3, 0.4, 0.5]) - assert doc.attributes["title"] == "Doc 1" results = vector_store.similarity_search_by_vector( [0.1, 0.2, 0.3, 0.4, 0.5], k=2 @@ -235,7 +208,6 @@ def test_vector_store_customization(self, sample_documents): vector_store.load_documents([sample_documents[2]], overwrite=False) result = vector_store.search_by_id("3") assert result.id == "3" - assert result.text == "This is document 3" # Define a simple text embedder function for testing def mock_embedder(text: str) -> list[float]: @@ -250,7 +222,6 @@ def mock_embedder(text: str) -> list[float]: # Test non-existent document non_existent = vector_store.search_by_id("nonexistent") assert non_existent.id == "nonexistent" - assert non_existent.text is None assert non_existent.vector is None finally: shutil.rmtree(temp_dir) diff --git a/tests/unit/query/context_builder/test_entity_extraction.py b/tests/unit/query/context_builder/test_entity_extraction.py index 8a15e8e929..aa2ff29b09 100644 --- a/tests/unit/query/context_builder/test_entity_extraction.py +++ b/tests/unit/query/context_builder/test_entity_extraction.py @@ -47,7 +47,8 @@ def similarity_search_by_text( return sorted( [ VectorStoreSearchResult( - document=document, score=abs(len(text) - len(document.text or "")) + document=document, + score=abs(len(text) - len(str(document.id) or "")), ) for document in self.documents ], @@ -94,8 +95,7 @@ def test_map_query_to_entities(): assert map_query_to_entities( query="t22", text_embedding_vectorstore=MockBaseVectorStore([ - VectorStoreDocument(id=entity.id, text=entity.title, vector=None) - for entity in entities + VectorStoreDocument(id=entity.id, vector=None) for entity in entities ]), text_embedder=ModelManager().get_or_create_embedding_model( model_type="mock_embedding", name="mock" @@ -116,8 +116,7 @@ def test_map_query_to_entities(): assert map_query_to_entities( query="t22", text_embedding_vectorstore=MockBaseVectorStore([ - VectorStoreDocument(id=entity.title, text=entity.title, vector=None) - for entity in entities + VectorStoreDocument(id=entity.title, vector=None) for entity in entities ]), text_embedder=ModelManager().get_or_create_embedding_model( model_type="mock_embedding", name="mock" @@ -138,8 +137,7 @@ def test_map_query_to_entities(): assert map_query_to_entities( query="", text_embedding_vectorstore=MockBaseVectorStore([ - VectorStoreDocument(id=entity.id, text=entity.title, vector=None) - for entity in entities + VectorStoreDocument(id=entity.id, vector=None) for entity in entities ]), text_embedder=ModelManager().get_or_create_embedding_model( model_type="mock_embedding", name="mock" @@ -165,8 +163,7 @@ def test_map_query_to_entities(): assert map_query_to_entities( query="", text_embedding_vectorstore=MockBaseVectorStore([ - VectorStoreDocument(id=entity.id, text=entity.title, vector=None) - for entity in entities + VectorStoreDocument(id=entity.id, vector=None) for entity in entities ]), text_embedder=ModelManager().get_or_create_embedding_model( model_type="mock_embedding", name="mock"