Obinna/s3 en 1111 update neo4j span result attributes

src/examples/neo4j_example/__init__.py

@@ -0,0 +1,8 @@
+from .basic import execute_query
+from langtrace_python_sdk import with_langtrace_root_span
+
+
+class Neo4jRunner:
+    @with_langtrace_root_span("Neo4jRunner")
+    def run(self):
+        execute_query()

src/examples/neo4j_example/basic.py

@@ -0,0 +1,26 @@
+import os
+from langtrace_python_sdk import langtrace
+from neo4j import GraphDatabase
+
+langtrace.init()
+
+def execute_query():
+    driver = GraphDatabase.driver(
+        os.getenv("NEO4J_URI"), 
+        auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD"))
+    )
+
+    records, summary, keys = driver.execute_query(
+        "MATCH (p:Person {age: $age}) RETURN p.name AS name",
+        age=42,
+        database_=os.getenv("NEO4J_DATABASE"),
+    )
+
+    # Loop through results and do something with them
+    for person in records:
+        print(person)
+    # Summary information
+    print("The query `{query}` returned {records_count} records in {time} ms.".format(
+        query=summary.query, records_count=len(records),
+        time=summary.result_available_after,
+    ))

src/examples/neo4j_graphrag_example/__init__.py

@@ -0,0 +1,9 @@
+import asyncio
+from .basic import search
+from langtrace_python_sdk import with_langtrace_root_span
+
+
+class Neo4jGraphRagRunner:
+    @with_langtrace_root_span("Neo4jGraphRagRunner")
+    def run(self):
+        asyncio.run(search())

src/examples/neo4j_graphrag_example/basic.py

@@ -0,0 +1,52 @@
+import os
+from langtrace_python_sdk import langtrace
+from langtrace_python_sdk.utils.with_root_span import with_langtrace_root_span
+from neo4j import GraphDatabase
+from neo4j_graphrag.generation import GraphRAG
+from neo4j_graphrag.indexes import create_vector_index
+from neo4j_graphrag.llm import OpenAILLM as LLM
+from neo4j_graphrag.embeddings import OpenAIEmbeddings as Embeddings
+from neo4j_graphrag.retrievers import VectorRetriever
+from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+
+langtrace.init()
+
+neo4j_driver = GraphDatabase.driver(os.getenv("NEO4J_URI"), auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD")))
+
+ex_llm=LLM(
+   model_name="gpt-4o-mini",
+   model_params={
+       "response_format": {"type": "json_object"},
+       "temperature": 0
+   })
+
+embedder = Embeddings()
+
+@with_langtrace_root_span("run_neo_rag")
+async def search():
+    # 1. Build KG and Store in Neo4j Database
+    kg_builder_pdf = SimpleKGPipeline(
+        llm=ex_llm,
+        driver=neo4j_driver,
+        embedder=embedder,
+        from_pdf=True
+    )
+    await kg_builder_pdf.run_async(file_path='src/examples/neo4j_graphrag_example/data/abramov.pdf')
+
+    create_vector_index(neo4j_driver, name="text_embeddings", label="Chunk",
+                    embedding_property="embedding", dimensions=1536, similarity_fn="cosine")
+
+    # 2. KG Retriever
+    vector_retriever = VectorRetriever(
+        neo4j_driver,
+        index_name="text_embeddings",
+        embedder=embedder
+    )
+
+    # 3. GraphRAG Class
+    llm = LLM(model_name="gpt-4o")
+    rag = GraphRAG(llm=llm, retriever=vector_retriever)
+
+    # 4. Run
+    response = rag.search("What did the author do in college?")
+    print(response.answer)

src/langtrace_python_sdk/instrumentation/neo4j/patch.py

@@ -41,7 +41,6 @@ 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,
@@ -50,7 +49,7 @@ def traced_method(wrapped, instance, args, kwargs):
             "langtrace.version": v(LANGTRACE_SDK_NAME),
             "db.system": "neo4j",
             "db.operation": api["OPERATION"],
-            "db.query": json.dumps(kwargs),
+            "db.query": json.dumps(args[0]) if args and len(args) > 0 else "",
             **(extra_attributes if extra_attributes is not None else {}),
         }
 
@@ -70,6 +69,14 @@ def traced_method(wrapped, instance, args, kwargs):
 
             try:
                 result = wrapped(*args, **kwargs)
+
+                if isinstance(result, tuple) and len(result) == 3:
+                    records, result_summary, keys = result
+                    _set_result_attributes(span, records, result_summary, keys)
+                else:
+                    res = json.dumps(result)
+                    set_span_attribute(span, "neo4j.result.query_response", res)
+
                 span.set_status(StatusCode.OK)
                 return result
             except Exception as err:
@@ -85,14 +92,12 @@ 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)
@@ -104,8 +109,72 @@ def _set_execute_query_attributes(span, args, kwargs):
 
     database = kwargs.get("database_", None)
     if database:
-        set_span_attribute(span, "db.name", database)
+        set_span_attribute(span, "neo4j.db.name", database)
 
     routing = kwargs.get("routing_", None)
     if routing:
-        set_span_attribute(span, "db.routing", str(routing))
+        set_span_attribute(span, "neo4j.db.routing", str(routing))
+
+
+@silently_fail
+def _set_result_attributes(span, records, result_summary, keys):
+    """
+    Set attributes related to the query result and summary
+    """
+    if records is not None:
+        record_count = len(records)
+        set_span_attribute(span, "neo4j.result.record_count", record_count)
+        if record_count > 0:
+            set_span_attribute(span, "neo4j.result.records", json.dumps(records))
+
+    if keys is not None:
+        set_span_attribute(span, "neo4j.result.keys", json.dumps(keys))
+
+    if result_summary:
+        if hasattr(result_summary, "database") and result_summary.database:
+            set_span_attribute(span, "neo4j.db.name", result_summary.database)
+
+        if hasattr(result_summary, "query_type") and result_summary.query_type:
+            set_span_attribute(span, "neo4j.result.query_type", result_summary.query_type)
+
+        if hasattr(result_summary, "parameters") and result_summary.parameters:
+            try:
+                set_span_attribute(span, "neo4j.result.parameters", json.dumps(result_summary.parameters))
+            except (TypeError, ValueError):
+                pass
+
+        if hasattr(result_summary, "result_available_after") and result_summary.result_available_after is not None:
+            set_span_attribute(span, "neo4j.result.available_after_ms", result_summary.result_available_after)
+
+        if hasattr(result_summary, "result_consumed_after") and result_summary.result_consumed_after is not None:
+            set_span_attribute(span, "neo4j.result.consumed_after_ms", result_summary.result_consumed_after)
+
+        if hasattr(result_summary, "counters") and result_summary.counters:
+            counters = result_summary.counters
+            if hasattr(counters, "nodes_created") and counters.nodes_created:
+                set_span_attribute(span, "neo4j.result.nodes_created", counters.nodes_created)
+
+            if hasattr(counters, "nodes_deleted") and counters.nodes_deleted:
+                set_span_attribute(span, "neo4j.result.nodes_deleted", counters.nodes_deleted)
+
+            if hasattr(counters, "relationships_created") and counters.relationships_created:
+                set_span_attribute(span, "neo4j.result.relationships_created", counters.relationships_created)
+
+            if hasattr(counters, "relationships_deleted") and counters.relationships_deleted:
+                set_span_attribute(span, "neo4j.result.relationships_deleted", counters.relationships_deleted)
+
+            if hasattr(counters, "properties_set") and counters.properties_set:
+                set_span_attribute(span, "neo4j.result.properties_set", counters.properties_set)
+
+        if hasattr(result_summary, "plan") and result_summary.plan:
+            try:
+                set_span_attribute(span, "neo4j.result.plan", json.dumps(result_summary.plan))
+            except (TypeError, ValueError):
+                pass
+
+        if hasattr(result_summary, "notifications") and result_summary.notifications:
+            try:
+                set_span_attribute(span, "neo4j.result.notification_count", len(result_summary.notifications))
+                set_span_attribute(span, "neo4j.result.notifications", json.dumps(result_summary.notifications))
+            except (AttributeError, TypeError):
+                pass

src/run_example.py

@@ -24,7 +24,9 @@
     "cerebras": False,
     "google_genai": False,
     "graphlit": False,
-    "phidata": True,
+    "phidata": False,
+    "neo4j": True,
+    "neo4jgraphrag": True,
 }
 
 if ENABLED_EXAMPLES["anthropic"]:
@@ -158,3 +160,15 @@
 
     print(Fore.BLUE + "Running PhiData example" + Fore.RESET)
     PhiDataRunner().run()
+
+if ENABLED_EXAMPLES["neo4j"]:
+    from examples.neo4j_example import Neo4jRunner
+
+    print(Fore.BLUE + "Running Neo4j example" + Fore.RESET)
+    Neo4jRunner().run()
+
+if ENABLED_EXAMPLES["neo4jgraphrag"]:
+    from examples.neo4j_graphrag_example import Neo4jGraphRagRunner
+
+    print(Fore.BLUE + "Running Neo4jGraphRag example" + Fore.RESET)
+    Neo4jGraphRagRunner().run()