diff --git a/src/examples/neo4j_example/__init__.py b/src/examples/neo4j_example/__init__.py new file mode 100644 index 00000000..3a1740c9 --- /dev/null +++ b/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() diff --git a/src/examples/neo4j_example/basic.py b/src/examples/neo4j_example/basic.py new file mode 100644 index 00000000..12c56c44 --- /dev/null +++ b/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, + )) diff --git a/src/examples/neo4j_graphrag_example/__init__.py b/src/examples/neo4j_graphrag_example/__init__.py new file mode 100644 index 00000000..198c13c5 --- /dev/null +++ b/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()) diff --git a/src/examples/neo4j_graphrag_example/basic.py b/src/examples/neo4j_graphrag_example/basic.py new file mode 100644 index 00000000..79c4fa7b --- /dev/null +++ b/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) diff --git a/src/examples/neo4j_graphrag_example/data/abramov.pdf b/src/examples/neo4j_graphrag_example/data/abramov.pdf new file mode 100644 index 00000000..dc7d0a61 Binary files /dev/null and b/src/examples/neo4j_graphrag_example/data/abramov.pdf differ diff --git a/src/langtrace_python_sdk/instrumentation/neo4j/patch.py b/src/langtrace_python_sdk/instrumentation/neo4j/patch.py index a3bd9ca2..c3633472 100644 --- a/src/langtrace_python_sdk/instrumentation/neo4j/patch.py +++ b/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)) \ No newline at end of file + 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 \ No newline at end of file diff --git a/src/run_example.py b/src/run_example.py index 32c7e0f5..cf108cb2 100644 --- a/src/run_example.py +++ b/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()