Override neo4j user agent when driver is injected

Author: jonbesga

examples/customize/build_graph/components/writers/custom_writer.py

@@ -6,11 +6,12 @@
 from neo4j_graphrag.experimental.components.kg_writer import KGWriter, KGWriterModel
 from neo4j_graphrag.experimental.components.types import LexicalGraphConfig, Neo4jGraph
 from pydantic import validate_call
+from neo4j_graphrag.utils import telemetry
 
 
 class MyWriter(KGWriter):
     def __init__(self, driver: neo4j.Driver) -> None:
-        self.driver = driver
+        self.driver = telemetry.override_user_agent(driver)
 
     @validate_call
     async def run(

src/neo4j_graphrag/__init__.py

@@ -12,3 +12,9 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("neo4j-graphrag")
+except PackageNotFoundError:
+    __version__ = "0.0.0"

src/neo4j_graphrag/experimental/components/kg_writer.py

@@ -34,6 +34,7 @@
     UPSERT_RELATIONSHIP_QUERY,
     UPSERT_RELATIONSHIP_QUERY_VARIABLE_SCOPE_CLAUSE,
 )
+from neo4j_graphrag.utils import telemetry
 
 logger = logging.getLogger(__name__)
 
@@ -113,7 +114,7 @@ def __init__(
         neo4j_database: Optional[str] = None,
         batch_size: int = 1000,
     ):
-        self.driver = driver
+        self.driver = telemetry.override_user_agent(driver)
         self.neo4j_database = neo4j_database
         self.batch_size = batch_size
         self.is_version_5_23_or_above = self._check_if_version_5_23_or_above()

src/neo4j_graphrag/experimental/components/neo4j_reader.py

@@ -25,6 +25,7 @@
     TextChunks,
 )
 from neo4j_graphrag.experimental.pipeline import Component
+from neo4j_graphrag.utils import telemetry
 
 
 class Neo4jChunkReader(Component):
@@ -58,7 +59,7 @@ def __init__(
         fetch_embeddings: bool = False,
         neo4j_database: Optional[str] = None,
     ):
-        self.driver = driver
+        self.driver = telemetry.override_user_agent(driver)
         self.fetch_embeddings = fetch_embeddings
         self.neo4j_database = neo4j_database
 

src/neo4j_graphrag/experimental/components/resolver.py

@@ -19,6 +19,7 @@
 
 from neo4j_graphrag.experimental.components.types import ResolutionStats
 from neo4j_graphrag.experimental.pipeline import Component
+from neo4j_graphrag.utils import telemetry
 
 
 class EntityResolver(Component, abc.ABC):
@@ -34,7 +35,7 @@ def __init__(
         driver: neo4j.Driver,
         filter_query: Optional[str] = None,
     ) -> None:
-        self.driver = driver
+        self.driver = telemetry.override_user_agent(driver)
         self.filter_query = filter_query
 
     @abc.abstractmethod

src/neo4j_graphrag/retrievers/base.py

@@ -24,6 +24,7 @@
 
 from neo4j_graphrag.exceptions import Neo4jVersionError
 from neo4j_graphrag.types import RawSearchResult, RetrieverResult, RetrieverResultItem
+from neo4j_graphrag.utils import telemetry
 
 T = ParamSpec("T")
 P = TypeVar("P")
@@ -82,7 +83,7 @@ class Retriever(ABC, metaclass=RetrieverMetaclass):
     VERIFY_NEO4J_VERSION = True
 
     def __init__(self, driver: neo4j.Driver, neo4j_database: Optional[str] = None):
-        self.driver = driver
+        self.driver = telemetry.override_user_agent(driver)
         self.neo4j_database = neo4j_database
         if self.VERIFY_NEO4J_VERSION:
             self._verify_version()

src/neo4j_graphrag/utils/telemetry.py

@@ -0,0 +1,8 @@
+import neo4j
+from neo4j_graphrag import __version__
+
+
+# Override user-agent used by neo4j package so we can measure usage of the package by version
+def override_user_agent(driver: neo4j.Driver) -> neo4j.Driver:
+    driver._pool.pool_config.user_agent = f"neo4j-graphrag-python/v{__version__}"
+    return driver