langtrace-python-sdk 3.8.7__py3-none-any.whl → 3.8.9__py3-none-any.whl

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()

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,
+    ))

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())

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)

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",
+    },
+}

langtrace_python_sdk/instrumentation/__init__.py

@@ -22,6 +22,8 @@ from .litellm import LiteLLMInstrumentation
 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
 from .openai_agents import OpenAIAgentsInstrumentation
@@ -59,6 +61,7 @@ __all__ = [
     "AWSBedrockInstrumentation",
     "CerebrasInstrumentation",
     "MilvusInstrumentation",
+    "Neo4jGraphRAGInstrumentation",
     "GoogleGenaiInstrumentation",
     "CrewaiToolsInstrumentation",
     "GraphlitInstrumentation",

langtrace_python_sdk/instrumentation/neo4j/__init__.py

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

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

langtrace_python_sdk/instrumentation/neo4j/patch.py

@@ -0,0 +1,180 @@
+"""
+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(args[0]) if args and len(args) > 0 else "",
+            **(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)
+
+                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:
+                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.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.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, "neo4j.db.name", database)
+
+    routing = kwargs.get("routing_", None)
+    if 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

langtrace_python_sdk/instrumentation/neo4j_graphrag/__init__.py

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

langtrace_python_sdk/instrumentation/neo4j_graphrag/instrumentation.py

@@ -0,0 +1,62 @@
+"""
+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.
+"""
+
+from typing import Collection
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry import trace
+from wrapt import wrap_function_wrapper as _W
+from importlib.metadata import version as v
+from .patch import patch_graphrag_search, patch_kg_pipeline_run, \
+patch_kg_pipeline_run, patch_retriever_search
+
+
+class Neo4jGraphRAGInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["neo4j-graphrag>=1.6.0"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = trace.get_tracer(__name__, "", tracer_provider)
+        graphrag_version = v("neo4j-graphrag")
+        
+        try:
+            # instrument kg builder
+            _W(
+                "neo4j_graphrag.experimental.pipeline.kg_builder",
+                "SimpleKGPipeline.run_async",
+                patch_kg_pipeline_run("run_async", graphrag_version, tracer),
+            )
+            
+            # Instrument GraphRAG
+            _W(
+                "neo4j_graphrag.generation.graphrag",
+                "GraphRAG.search",
+                patch_graphrag_search("search", graphrag_version, tracer),
+            )
+            
+            # Instrument retrievers
+            _W(
+                "neo4j_graphrag.retrievers.vector",
+                "VectorRetriever.get_search_results",
+                patch_retriever_search("vector_search", graphrag_version, tracer),
+            )
+            
+        except Exception as e:
+            print(f"Failed to instrument Neo4j GraphRAG: {e}")
+
+    def _uninstrument(self, **kwargs):
+        pass

langtrace_python_sdk/instrumentation/neo4j_graphrag/patch.py

@@ -0,0 +1,229 @@
+"""
+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 importlib_metadata import version as v
+from langtrace.trace_attributes import FrameworkSpanAttributes
+from opentelemetry import baggage
+from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace.status import Status, StatusCode
+
+from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace_python_sdk.constants.instrumentation.common import (
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, SERVICE_PROVIDERS)
+from langtrace_python_sdk.utils.llm import set_span_attributes
+from langtrace_python_sdk.utils.misc import serialize_args, serialize_kwargs
+
+
+def patch_kg_pipeline_run(operation_name: str, version: str, tracer: Tracer):
+    
+    async def async_traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS.get("NEO4J_GRAPHRAG", "neo4j_graphrag")
+        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": "framework",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "neo4j.pipeline.type": "SimpleKGPipeline",
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        if len(args) > 0:
+            span_attributes["neo4j.pipeline.inputs"] = serialize_args(*args)
+        if kwargs:
+            span_attributes["neo4j.pipeline.kwargs"] = serialize_kwargs(**kwargs)
+
+        file_path = kwargs.get("file_path", args[0] if len(args) > 0 else None)
+        text = kwargs.get("text", args[1] if len(args) > 1 else None)
+        if file_path:
+            span_attributes["neo4j.pipeline.file_path"] = file_path
+        if text:
+            span_attributes["neo4j.pipeline.text_length"] = len(text)
+
+        if hasattr(instance, "runner") and hasattr(instance.runner, "config"):
+            config = instance.runner.config
+            if config:
+                span_attributes["neo4j.pipeline.from_pdf"] = getattr(config, "from_pdf", None)
+                span_attributes["neo4j.pipeline.perform_entity_resolution"] = getattr(config, "perform_entity_resolution", None)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=f"neo4j.pipeline.{operation_name}",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            try:
+                set_span_attributes(span, attributes)
+
+                result = await wrapped(*args, **kwargs)
+
+                if result:
+                    try:
+                        if hasattr(result, "to_dict"):
+                            result_dict = result.to_dict()
+                            span.set_attribute("neo4j.pipeline.result", json.dumps(result_dict))
+                        elif hasattr(result, "model_dump"):
+                            result_dict = result.model_dump()
+                            span.set_attribute("neo4j.pipeline.result", json.dumps(result_dict))
+                    except Exception as e:
+                        span.set_attribute("neo4j.pipeline.result_error", str(e))
+                
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+            except Exception as err:
+                span.record_exception(err)
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+                raise
+
+    return async_traced_method
+
+
+def patch_graphrag_search(operation_name: str, version: str, tracer: Tracer):
+    
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS.get("NEO4J_GRAPHRAG", "neo4j_graphrag")
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+        
+        # Basic attributes
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "framework",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "neo4j_graphrag.operation": operation_name,
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        query_text = kwargs.get("query_text", args[0] if len(args) > 0 else None)
+        if query_text:
+            span_attributes["neo4j_graphrag.query_text"] = query_text
+
+        retriever_config = kwargs.get("retriever_config", None)
+        if retriever_config:
+            span_attributes["neo4j_graphrag.retriever_config"] = json.dumps(retriever_config)
+
+        if hasattr(instance, "retriever"):
+            span_attributes["neo4j_graphrag.retriever_type"] = instance.retriever.__class__.__name__
+
+        if hasattr(instance, "llm"):
+            span_attributes["neo4j_graphrag.llm_type"] = instance.llm.__class__.__name__
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=f"neo4j_graphrag.{operation_name}",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            try:
+                set_span_attributes(span, attributes)
+
+                result = wrapped(*args, **kwargs)
+                
+                if result and hasattr(result, "answer"):
+                    span.set_attribute("neo4j_graphrag.answer", result.answer)
+
+                    if hasattr(result, "retriever_result") and result.retriever_result:
+                        try:
+                            retriever_items = len(result.retriever_result.items)
+                            span.set_attribute("neo4j_graphrag.context_items", retriever_items)
+                        except Exception:
+                            pass
+                
+                span.set_status(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
+
+
+def patch_retriever_search(operation_name: str, version: str, tracer: Tracer):
+    
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS.get("NEO4J_GRAPHRAG", "neo4j_graphrag")
+        extra_attributes = baggage.get_baggage(LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY)
+        
+        # Basic attributes
+        span_attributes = {
+            "langtrace.sdk.name": "langtrace-python-sdk",
+            "langtrace.service.name": service_provider,
+            "langtrace.service.type": "framework",
+            "langtrace.service.version": version,
+            "langtrace.version": v(LANGTRACE_SDK_NAME),
+            "neo4j.retriever.operation": operation_name,
+            "neo4j.retriever.type": instance.__class__.__name__,
+            **(extra_attributes if extra_attributes is not None else {}),
+        }
+
+        query_text = kwargs.get("query_text", args[0] if len(args) > 0 else None)
+        if query_text:
+            span_attributes["neo4j.retriever.query_text"] = query_text
+
+        if hasattr(instance, "__class__") and hasattr(instance.__class__, "__name__"):
+            retriever_type = instance.__class__.__name__
+
+            if retriever_type == "VectorRetriever" and hasattr(instance, "index_name"):
+                span_attributes["neo4j.vector_retriever.index_name"] = instance.index_name
+
+            if retriever_type == "KnowledgeGraphRetriever" and hasattr(instance, "cypher_query"):
+                span_attributes["neo4j.kg_retriever.cypher_query"] = instance.cypher_query
+
+        for param in ["top_k", "similarity_threshold"]:
+            if param in kwargs:
+                span_attributes[f"neo4j.retriever.{param}"] = kwargs[param]
+            elif hasattr(instance, param):
+                span_attributes[f"neo4j.retriever.{param}"] = getattr(instance, param)
+
+        attributes = FrameworkSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=f"neo4j.retriever.{operation_name}",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            try:
+                set_span_attributes(span, attributes)
+
+                result = wrapped(*args, **kwargs)
+
+                if result:
+                    if hasattr(result, "items") and isinstance(result.items, list):
+                        span.set_attribute("neo4j.retriever.items_count", len(result.items))
+
+                        try:
+                            item_ids = [item.id for item in result.items[:5] if hasattr(item, "id")]
+                            if item_ids:
+                                span.set_attribute("neo4j.retriever.item_ids", json.dumps(item_ids))
+                        except Exception:
+                            pass
+                
+                span.set_status(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

langtrace_python_sdk/langtrace.py

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

langtrace_python_sdk/version.py

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

{langtrace_python_sdk-3.8.7.dist-info → langtrace_python_sdk-3.8.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langtrace-python-sdk
-Version: 3.8.7
+Version: 3.8.9
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>
@@ -23,7 +23,7 @@ Requires-Dist: sentry-sdk>=2.14.0
 Requires-Dist: setuptools
 Requires-Dist: sqlalchemy
 Requires-Dist: tiktoken>=0.1.1
-Requires-Dist: trace-attributes==7.2.0
+Requires-Dist: trace-attributes==7.2.1
 Requires-Dist: transformers>=4.11.3
 Requires-Dist: ujson>=5.10.0
 Provides-Extra: dev

{langtrace_python_sdk-3.8.7.dist-info → langtrace_python_sdk-3.8.9.dist-info}/RECORD

@@ -83,6 +83,11 @@ examples/mistral_example/complete.py,sha256=Ydf5iOCGM2tySSd5vzGvNh3Qz39aTpxcCItB
 examples/mistral_example/complete_async.py,sha256=-SVOLNLNi5JL3o4obBsCXKOt8A6b61j3otbrrYPOOtM,586
 examples/mistral_example/embeddings.py,sha256=LKq_k3Y-TS2SCkyFKw2tLiYreVRHqlY6z0Gfh1y_7u0,468
 examples/mongo_vector_search_example/main.py,sha256=6bHOxR3LfAZwmGmv74jgNCciATS5h5e-Tb5z0cD0z3Q,1542
+examples/neo4j_example/__init__.py,sha256=YG7-_cQnXt-fJQtH2bJu_uSNyVCmYyUevU4H-GML_Sk,200
+examples/neo4j_example/basic.py,sha256=xw5fNi5y-U5wzrVIsnvbKBZDURUzoYriawu6nyEXi3o,780
+examples/neo4j_graphrag_example/__init__.py,sha256=YCXutcM3otTKEff1S3qTciFBxOMrQw7mQ4WGfJ7-mU8,230
+examples/neo4j_graphrag_example/basic.py,sha256=uYf4O8bT3GkU5Qa-tluEr9sYTm5gjdzHAKf-C8bi5GY,1727
+examples/neo4j_graphrag_example/data/abramov.pdf,sha256=Fm93bGdz74PJel-Dw_wZeuqqedloTTZiOakIWLzYDPA,46933
 examples/ollama_example/__init__.py,sha256=qOx0jGCPuSpRCPiqtDVm7F0z8hIZ8C75hDZ_C8Apz-s,399
 examples/ollama_example/basic.py,sha256=EPbsigOF4xBDBgLgAD0EzPo737ycVm7aXZr7F5Xt-A4,1062
 examples/openai_example/__init__.py,sha256=6faH7wTegSozKmS89sd1Tgv8AcEH0GfKkC7YaBWA8tg,849
@@ -115,8 +120,8 @@ examples/vertexai_example/main.py,sha256=gndId5X5ksD-ycxnAWMdEqIDbLc3kz5Vt8vm4YP
 examples/weaviate_example/__init__.py,sha256=8JMDBsRSEV10HfTd-YC7xb4txBjD3la56snk-Bbg2Kw,618
 examples/weaviate_example/query_text.py,sha256=wPHQTc_58kPoKTZMygVjTj-2ZcdrIuaausJfMxNQnQc,127162
 langtrace_python_sdk/__init__.py,sha256=VZM6i71NR7pBQK6XvJWRelknuTYUhqwqE7PlicKa5Wg,1166
-langtrace_python_sdk/langtrace.py,sha256=T-DsDrwWaL4gAUK1lkTRRpmvoO7F2WtO5hQZdyrVAxE,13791
-langtrace_python_sdk/version.py,sha256=zMoy0NhWOOCYaBDKkUvjVGd7u0NZBhM_7SModsbbTxQ,22
+langtrace_python_sdk/langtrace.py,sha256=P7aHGunNP8YQqd2glsVLXhSe4QDeAarfG2eCV9hhAm4,13944
+langtrace_python_sdk/version.py,sha256=KpxarCD6mNaS0PTJWWJOOJiShAjznuFT62LLRdWsuYI,22
 langtrace_python_sdk/constants/__init__.py,sha256=3CNYkWMdd1DrkGqzLUgNZXjdAlM6UFMlf_F-odAToyc,146
 langtrace_python_sdk/constants/exporter/langtrace_exporter.py,sha256=EVCrouYCpY98f0KSaKr4PzNxPULTZZO6dSA_crEOyJU,106
 langtrace_python_sdk/constants/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -131,6 +136,7 @@ langtrace_python_sdk/constants/instrumentation/groq.py,sha256=VFXmIl4aqGY_fS0PAm
 langtrace_python_sdk/constants/instrumentation/litellm.py,sha256=bMAlpY2scFe6Lql0Nl7euGNSO9QEV5Uzne12hnw3mSE,449
 langtrace_python_sdk/constants/instrumentation/milvus.py,sha256=lnl7RtV2Wc66mEmtVMjV5eWalt-caH4-sQEHpN1LmEE,1029
 langtrace_python_sdk/constants/instrumentation/mistral.py,sha256=9PlmcC5P5_BHJ-zsX1xekht6rSm7arTin58HAfdYvLk,730
+langtrace_python_sdk/constants/instrumentation/neo4j.py,sha256=-coUUCiroN2shr_C-HcWPYPmEcjIz1CQpaXtZo2uVFU,1002
 langtrace_python_sdk/constants/instrumentation/ollama.py,sha256=H_-S0xjqRsi5qSp7mAlK7Y9NlQ3BqOkG6ASogqqgdJY,212
 langtrace_python_sdk/constants/instrumentation/openai.py,sha256=uEOH5UXapU2DSf2AdgXTRhhJEHGWXUNFkUGD5QafflM,1164
 langtrace_python_sdk/constants/instrumentation/pinecone.py,sha256=0TityERbGWaHGSN8-vyYZtYCjVj8fQOKae8lng0O0Bk,478
@@ -141,7 +147,7 @@ langtrace_python_sdk/constants/instrumentation/weaviate.py,sha256=gtv-JBxvNGClEM
 langtrace_python_sdk/extensions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/extensions/langtrace_exporter.py,sha256=ckd8dMmY6h2oxE04p1JFLwUB5PSJX_Cy4eDFEM6aj4Y,6605
 langtrace_python_sdk/extensions/langtrace_filesystem.py,sha256=34fZutG28EJ66l67OvTGsydAH3ZpXgikdE7hVLqBpG4,7863
-langtrace_python_sdk/instrumentation/__init__.py,sha256=DC96oELorZedDf5zooQ6HGi-dmieXwUephgeB_LzfaU,2559
+langtrace_python_sdk/instrumentation/__init__.py,sha256=sUVSZtidQev87AspZmWK4M8J51mBnJXXNDwfBsl_Dcs,2692
 langtrace_python_sdk/instrumentation/agno/__init__.py,sha256=95fn4oA-CHB0mxc6KnVB20KSbXGl_ZZr9n99EEaXzrY,91
 langtrace_python_sdk/instrumentation/agno/instrumentation.py,sha256=XUnfvqpp13IgdF03xGKasq7kGjeaN1sXLIwCf-Nt_Nc,2667
 langtrace_python_sdk/instrumentation/agno/patch.py,sha256=qCUxCkzU9cYu_d8BzLgj_Ss97qib07tRVYpYDiNnNMs,16876
@@ -218,6 +224,12 @@ langtrace_python_sdk/instrumentation/milvus/patch.py,sha256=0yY5aQz0x7hpQZ8U-0qf
 langtrace_python_sdk/instrumentation/mistral/__init__.py,sha256=mkGALBQvq0jSfwDl6TU09SFwnVs6O4zkUi-yVmd3SNg,90
 langtrace_python_sdk/instrumentation/mistral/instrumentation.py,sha256=qtCkHCSOaiicUChbmTID4lcK1rbeW8oRSbpda2ogbgM,2328
 langtrace_python_sdk/instrumentation/mistral/patch.py,sha256=5EzqMeIEcMUH8P-l9-ijpRNWSc6mYwe9Gz0VY4PLrS0,6559
+langtrace_python_sdk/instrumentation/neo4j/__init__.py,sha256=6pG7SKOSKENaFEsstMJ4ZUH9wUP53JdQKdjHzJ4Lbko,86
+langtrace_python_sdk/instrumentation/neo4j/instrumentation.py,sha256=KECuFv7OuVBgh_NTZl4_W18GAGbKnaeQqRZM8NPZgjU,1668
+langtrace_python_sdk/instrumentation/neo4j/patch.py,sha256=wPAxslFdsHdYU5k1nHzWAqaHxQA-HhxgfVnb_U5qaZ0,7922
+langtrace_python_sdk/instrumentation/neo4j_graphrag/__init__.py,sha256=K-Qpfh-kn2EPxnYZ_oubLZNqAUO9Cx7mnMrJu_ms3UU,102
+langtrace_python_sdk/instrumentation/neo4j_graphrag/instrumentation.py,sha256=FLnddPSozuu5puqYD_FnXLMFet8YToVdWP-cY5QZrvM,2216
+langtrace_python_sdk/instrumentation/neo4j_graphrag/patch.py,sha256=4c7NKcTJeYsbi5V4aSsuPzTiFoBqtlcXmYx2bMlE7qg,9713
 langtrace_python_sdk/instrumentation/ollama/__init__.py,sha256=g2zJsXnDHinXPzTc-WxDeTtHmr9gmAj3K6l_00kP8c8,82
 langtrace_python_sdk/instrumentation/ollama/instrumentation.py,sha256=jdsvkqUJAAUNLVPtAkn_rG26HXetVQXWtjn4a6eWZro,2029
 langtrace_python_sdk/instrumentation/ollama/patch.py,sha256=w99r9wCCVDdJnZQEezYE2lW_iNFEtrldt9vq3ISAsag,5375
@@ -300,8 +312,8 @@ tests/pinecone/cassettes/test_query.yaml,sha256=b5v9G3ssUy00oG63PlFUR3JErF2Js-5A
 tests/pinecone/cassettes/test_upsert.yaml,sha256=neWmQ1v3d03V8WoLl8FoFeeCYImb8pxlJBWnFd_lITU,38607
 tests/qdrant/conftest.py,sha256=9n0uHxxIjWk9fbYc4bx-uP8lSAgLBVx-cV9UjnsyCHM,381
 tests/qdrant/test_qdrant.py,sha256=pzjAjVY2kmsmGfrI2Gs2xrolfuaNHz7l1fqGQCjp5_o,3353
-langtrace_python_sdk-3.8.7.dist-info/METADATA,sha256=au8c2l_qdpKK8lgICNgotmeqr6K6mD4r_nhOS5X9pJo,15844
-langtrace_python_sdk-3.8.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langtrace_python_sdk-3.8.7.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
-langtrace_python_sdk-3.8.7.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-langtrace_python_sdk-3.8.7.dist-info/RECORD,,
+langtrace_python_sdk-3.8.9.dist-info/METADATA,sha256=k-mlsrxBFSyNhSlXR2UkkEpPjIT8YzLHtCG3R82NeL8,15844
+langtrace_python_sdk-3.8.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langtrace_python_sdk-3.8.9.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
+langtrace_python_sdk-3.8.9.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+langtrace_python_sdk-3.8.9.dist-info/RECORD,,