diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/base.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/base.py index 12272f6cfca0b..7fd5a1b4f1784 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/base.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/base.py @@ -38,11 +38,13 @@ def __init__( graph_store: PropertyGraphStore, include_text: bool = True, include_text_preamble: Optional[str] = DEFAULT_PREAMBLE, + include_properties: bool = False, **kwargs: Any, ) -> None: self._graph_store = graph_store self.include_text = include_text self._include_text_preamble = include_text_preamble + self.include_properties = include_properties super().__init__(callback_manager=kwargs.get("callback_manager", None)) def _get_nodes_with_score( @@ -57,7 +59,10 @@ def _get_nodes_with_score( node_id=source_id ) - text = f"{triplet[0]!s} -> {triplet[1]!s} -> {triplet[2]!s}" + if self.include_properties: + text = f"{triplet[0]!s} -> {triplet[1]!s} -> {triplet[2]!s}" + else: + text = f"{triplet[0].id} -> {triplet[1].id} -> {triplet[2].id}" results.append( NodeWithScore( node=TextNode( diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/custom.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/custom.py index 74152ffe760a4..6e6222c0f5fbb 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/custom.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/custom.py @@ -32,9 +32,15 @@ def __init__( self, graph_store: PropertyGraphStore, include_text: bool = False, + include_properties: bool = False, **kwargs: Any, ) -> None: - super().__init__(graph_store=graph_store, include_text=include_text, **kwargs) + super().__init__( + graph_store=graph_store, + include_text=include_text, + include_properties=include_properties, + **kwargs, + ) self.init(**kwargs) @property diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/cypher_template.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/cypher_template.py index 2f575164079db..bc09a84e53335 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/cypher_template.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/cypher_template.py @@ -41,7 +41,9 @@ def __init__( self.output_cls = output_cls self.cypher_query = cypher_query - super().__init__(graph_store=graph_store, include_text=False) + super().__init__( + graph_store=graph_store, include_text=False, include_properties=False + ) def retrieve_from_graph(self, query_bundle: QueryBundle) -> List[NodeWithScore]: question = query_bundle.query_str diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/llm_synonym.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/llm_synonym.py index a464663d74f22..17938c1ed77a5 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/llm_synonym.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/llm_synonym.py @@ -53,6 +53,7 @@ def __init__( self, graph_store: PropertyGraphStore, include_text: bool = True, + include_properties: bool = False, synonym_prompt: Union[ BasePromptTemplate, str ] = DEFAULT_SYNONYM_EXPAND_TEMPLATE, @@ -69,7 +70,12 @@ def __init__( self._output_parsing_fn = output_parsing_fn self._max_keywords = max_keywords self._path_depth = path_depth - super().__init__(graph_store=graph_store, include_text=include_text, **kwargs) + super().__init__( + graph_store=graph_store, + include_text=include_text, + include_properties=include_properties, + **kwargs, + ) def _parse_llm_output(self, output: str) -> List[str]: if self._output_parsing_fn: diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/text_to_cypher.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/text_to_cypher.py index afdd02bda55c9..6b52b9fac6128 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/text_to_cypher.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/text_to_cypher.py @@ -59,7 +59,9 @@ def __init__( ) self.cypher_validator = cypher_validator self.allowed_output_fields = allowed_output_fields - super().__init__(graph_store=graph_store, include_text=False) + super().__init__( + graph_store=graph_store, include_text=False, include_properties=False + ) def _parse_generated_cypher(self, cypher_query: str) -> str: if self.cypher_validator is not None: diff --git a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py index 5ba6d98831aec..ac06729111d83 100644 --- a/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py +++ b/llama-index-core/llama_index/core/indices/property_graph/sub_retrievers/vector.py @@ -43,6 +43,7 @@ def __init__( self, graph_store: PropertyGraphStore, include_text: bool = True, + include_properties: bool = False, embed_model: Optional[BaseEmbedding] = None, vector_store: Optional[BasePydanticVectorStore] = None, similarity_top_k: int = 4, @@ -59,7 +60,12 @@ def __init__( self._similarity_score = similarity_score self._filters = filters - super().__init__(graph_store=graph_store, include_text=include_text, **kwargs) + super().__init__( + graph_store=graph_store, + include_text=include_text, + include_properties=include_properties, + **kwargs, + ) def _get_vector_store_query(self, query_bundle: QueryBundle) -> VectorStoreQuery: if query_bundle.embedding is None: