Obinna/s3 en 1111 support neo4j db

pyproject.toml

@@ -18,7 +18,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 dependencies = [
-  'trace-attributes==7.2.0',
+  'trace-attributes==7.2.1',
   'opentelemetry-api>=1.25.0',
   'opentelemetry-sdk>=1.25.0',
   'opentelemetry-instrumentation>=0.47b0',

src/langtrace_python_sdk/constants/instrumentation/neo4j.py

@@ -0,0 +1,36 @@
+from langtrace.trace_attributes import Neo4jMethods
+
+APIS = {
+    "RUN": {
+        "METHOD": Neo4jMethods.RUN.value,
+        "OPERATION": "run",
+    },
+    "BEGIN_TRANSACTION": {
+        "METHOD": Neo4jMethods.BEGIN_TRANSACTION.value,
+        "OPERATION": "begin_transaction",
+    },
+    "READ_TRANSACTION": {
+        "METHOD": Neo4jMethods.READ_TRANSACTION.value,
+        "OPERATION": "read_transaction",
+    },
+    "WRITE_TRANSACTION": {
+        "METHOD": Neo4jMethods.WRITE_TRANSACTION.value,
+        "OPERATION": "write_transaction",
+    },
+    "EXECUTE_READ": {
+        "METHOD": Neo4jMethods.EXECUTE_READ.value,
+        "OPERATION": "execute_read",
+    },
+    "EXECUTE_WRITE": {
+        "METHOD": Neo4jMethods.EXECUTE_WRITE.value,
+        "OPERATION": "execute_write",
+    },
+    "EXECUTE_QUERY": {
+        "METHOD": Neo4jMethods.EXECUTE_QUERY.value,
+        "OPERATION": "execute_query",
+    },
+    "TX_RUN": {
+        "METHOD": Neo4jMethods.TX_RUN.value,
+        "OPERATION": "tx_run",
+    },
+}

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -22,6 +22,7 @@
 from .llamaindex import LlamaindexInstrumentation
 from .milvus import MilvusInstrumentation
 from .mistral import MistralInstrumentation
+from .neo4j import Neo4jInstrumentation
 from .neo4j_graphrag import Neo4jGraphRAGInstrumentation
 from .ollama import OllamaInstrumentor
 from .openai import OpenAIInstrumentation

src/langtrace_python_sdk/instrumentation/neo4j/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import Neo4jInstrumentation
+
+__all__ = ["Neo4jInstrumentation"]

src/langtrace_python_sdk/instrumentation/neo4j/instrumentation.py

@@ -0,0 +1,51 @@
+"""
+Copyright (c) 2025 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+"""
+
+import importlib.metadata
+import logging
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+from langtrace_python_sdk.constants.instrumentation.neo4j import APIS
+from langtrace_python_sdk.instrumentation.neo4j.patch import driver_patch
+
+logging.basicConfig(level=logging.FATAL)
+
+
+class Neo4jInstrumentation(BaseInstrumentor):
+    """
+    The Neo4jInstrumentation class represents the Neo4j graph database instrumentation
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["neo4j >= 5.25.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version("neo4j")
+
+        wrap_function_wrapper(
+            "neo4j._sync.driver",
+            "Driver.execute_query",
+            driver_patch("EXECUTE_QUERY", version, tracer),
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/neo4j/patch.py

@@ -0,0 +1,111 @@
+"""
+Copyright (c) 2025 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+"""
+
+import json
+
+from langtrace_python_sdk.utils.llm import get_span_name
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from langtrace.trace_attributes import DatabaseSpanAttributes
+from langtrace_python_sdk.utils import set_span_attribute
+from opentelemetry import baggage, trace
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.trace.propagation import set_span_in_context
+
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
+from langtrace_python_sdk.constants.instrumentation.neo4j import APIS
+from importlib.metadata import version as v
+
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+
+
+def driver_patch(operation_name, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+
+        api = APIS[operation_name]
+        service_provider = SERVICE_PROVIDERS.get("NEO4J", "neo4j")
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "vectordb",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "db.system": "neo4j",
+            "db.operation": api["OPERATION"],
+            "db.query": json.dumps(kwargs),
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        attributes = DatabaseSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(api["METHOD"]),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            for field, value in attributes.model_dump(by_alias=True).items():
+                if value is not None:
+                    span.set_attribute(field, value)
+
+            if operation_name == "EXECUTE_QUERY":
+                _set_execute_query_attributes(span, args, kwargs)
+
+            try:
+                result = wrapped(*args, **kwargs)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as err:
+                span.record_exception(err)
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+                raise
+
+    return traced_method
+
+
+@silently_fail
+def _set_execute_query_attributes(span, args, kwargs):
+    query = args[0] if args else kwargs.get("query_", None)
+    if query:
+        if hasattr(query, "text"):
+            set_span_attribute(span, "db.statement", query.text)
+            set_span_attribute(span, "db.query", query.text)
+            if hasattr(query, "metadata") and query.metadata:
+                set_span_attribute(span, "db.query.metadata", json.dumps(query.metadata))
+            if hasattr(query, "timeout") and query.timeout:
+                set_span_attribute(span, "db.query.timeout", query.timeout)
+        else:
+            set_span_attribute(span, "db.statement", query)
+            set_span_attribute(span, "db.query", query)
+
+    parameters = kwargs.get("parameters_", None)
+    if parameters:
+        try:
+            set_span_attribute(span, "db.statement.parameters", json.dumps(parameters))
+        except (TypeError, ValueError):
+            pass
+
+    database = kwargs.get("database_", None)
+    if database:
+        set_span_attribute(span, "db.name", database)
+
+    routing = kwargs.get("routing_", None)
+    if routing:
+        set_span_attribute(span, "db.routing", str(routing))

src/langtrace_python_sdk/langtrace.py

@@ -47,9 +47,9 @@
     CrewaiToolsInstrumentation, DspyInstrumentation, EmbedchainInstrumentation,
     GeminiInstrumentation, GoogleGenaiInstrumentation, GraphlitInstrumentation,
     GroqInstrumentation, LangchainCommunityInstrumentation,
-    LangchainCoreInstrumentation, LangchainInstrumentation,
-    LanggraphInstrumentation, LiteLLMInstrumentation, LlamaindexInstrumentation,
-    MilvusInstrumentation, MistralInstrumentation, Neo4jGraphRAGInstrumentation,
+    LangchainCoreInstrumentation, LangchainInstrumentation, LanggraphInstrumentation, 
+    LiteLLMInstrumentation, LlamaindexInstrumentation, MilvusInstrumentation, 
+    MistralInstrumentation, Neo4jInstrumentation, Neo4jGraphRAGInstrumentation,
     OllamaInstrumentor, OpenAIAgentsInstrumentation, OpenAIInstrumentation,
     PhiDataInstrumentation, PineconeInstrumentation, PyMongoInstrumentation,
     QdrantInstrumentation, VertexAIInstrumentation, WeaviateInstrumentation)
@@ -284,6 +284,7 @@ def init(
         "phidata": PhiDataInstrumentation(),
         "agno": AgnoInstrumentation(),
         "mistralai": MistralInstrumentation(),
+        "neo4j": Neo4jInstrumentation(),
         "neo4j-graphrag": Neo4jGraphRAGInstrumentation(),
         "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.8.8"
+__version__ = "3.8.9"