openinference-instrumentation-haystack 0.1.0__tar.gz

openinference_instrumentation_haystack-0.1.0/.gitignore

@@ -0,0 +1,2 @@
+# vendored virtual environments
+.venv

openinference_instrumentation_haystack-0.1.0/LICENSE

@@ -0,0 +1,201 @@
+ Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+    "License" shall mean the terms and conditions for use, reproduction,
+    and distribution as defined by Sections 1 through 9 of this document.
+
+    "Licensor" shall mean the copyright owner or entity authorized by
+    the copyright owner that is granting the License.
+
+    "Legal Entity" shall mean the union of the acting entity and all
+    other entities that control, are controlled by, or are under common
+    control with that entity. For the purposes of this definition,
+    "control" means (i) the power, direct or indirect, to cause the
+    direction or management of such entity, whether by contract or
+    otherwise, or (ii) ownership of fifty percent (50%) or more of the
+    outstanding shares, or (iii) beneficial ownership of such entity.
+
+    "You" (or "Your") shall mean an individual or Legal Entity
+    exercising permissions granted by this License.
+
+    "Source" form shall mean the preferred form for making modifications,
+    including but not limited to software source code, documentation
+    source, and configuration files.
+
+    "Object" form shall mean any form resulting from mechanical
+    transformation or translation of a Source form, including but
+    not limited to compiled object code, generated documentation,
+    and conversions to other media types.
+
+    "Work" shall mean the work of authorship, whether in Source or
+    Object form, made available under the License, as indicated by a
+    copyright notice that is included in or attached to the work
+    (an example is provided in the Appendix below).
+
+    "Derivative Works" shall mean any work, whether in Source or Object
+    form, that is based on (or derived from) the Work and for which the
+    editorial revisions, annotations, elaborations, or other modifications
+    represent, as a whole, an original work of authorship. For the purposes
+    of this License, Derivative Works shall not include works that remain
+    separable from, or merely link (or bind by name) to the interfaces of,
+    the Work and Derivative Works thereof.
+
+    "Contribution" shall mean any work of authorship, including
+    the original version of the Work and any modifications or additions
+    to that Work or Derivative Works thereof, that is intentionally
+    submitted to Licensor for inclusion in the Work by the copyright owner
+    or by an individual or Legal Entity authorized to submit on behalf of
+    the copyright owner. For the purposes of this definition, "submitted"
+    means any form of electronic, verbal, or written communication sent
+    to the Licensor or its representatives, including but not limited to
+    communication on electronic mailing lists, source code control systems,
+    and issue tracking systems that are managed by, or on behalf of, the
+    Licensor for the purpose of discussing and improving the Work, but
+    excluding communication that is conspicuously marked or otherwise
+    designated in writing by the copyright owner as "Not a Contribution."
+
+    "Contributor" shall mean Licensor and any individual or Legal Entity
+    on behalf of whom a Contribution has been received by Licensor and
+    subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+   this License, each Contributor hereby grants to You a perpetual,
+   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+   copyright license to reproduce, prepare Derivative Works of,
+   publicly display, publicly perform, sublicense, and distribute the
+   Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+   this License, each Contributor hereby grants to You a perpetual,
+   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+   (except as stated in this section) patent license to make, have made,
+   use, offer to sell, sell, import, and otherwise transfer the Work,
+   where such license applies only to those patent claims licensable
+   by such Contributor that are necessarily infringed by their
+   Contribution(s) alone or by combination of their Contribution(s)
+   with the Work to which such Contribution(s) was submitted. If You
+   institute patent litigation against any entity (including a
+   cross-claim or counterclaim in a lawsuit) alleging that the Work
+   or a Contribution incorporated within the Work constitutes direct
+   or contributory patent infringement, then any patent licenses
+   granted to You under this License for that Work shall terminate
+   as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+   Work or Derivative Works thereof in any medium, with or without
+   modifications, and in Source or Object form, provided that You
+   meet the following conditions:
+
+    (a) You must give any other recipients of the Work or
+    Derivative Works a copy of this License; and
+
+    (b) You must cause any modified files to carry prominent notices
+    stating that You changed the files; and
+
+    (c) You must retain, in the Source form of any Derivative Works
+    that You distribute, all copyright, patent, trademark, and
+    attribution notices from the Source form of the Work,
+    excluding those notices that do not pertain to any part of
+    the Derivative Works; and
+
+    (d) If the Work includes a "NOTICE" text file as part of its
+    distribution, then any Derivative Works that You distribute must
+    include a readable copy of the attribution notices contained
+    within such NOTICE file, excluding those notices that do not
+    pertain to any part of the Derivative Works, in at least one
+    of the following places: within a NOTICE text file distributed
+    as part of the Derivative Works; within the Source form or
+    documentation, if provided along with the Derivative Works; or,
+    within a display generated by the Derivative Works, if and
+    wherever such third-party notices normally appear. The contents
+    of the NOTICE file are for informational purposes only and
+    do not modify the License. You may add Your own attribution
+    notices within Derivative Works that You distribute, alongside
+    or as an addendum to the NOTICE text from the Work, provided
+    that such additional attribution notices cannot be construed
+    as modifying the License.
+
+    You may add Your own copyright statement to Your modifications and
+    may provide additional or different license terms and conditions
+    for use, reproduction, or distribution of Your modifications, or
+    for any such Derivative Works as a whole, provided Your use,
+    reproduction, and distribution of the Work otherwise complies with
+    the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+   any Contribution intentionally submitted for inclusion in the Work
+   by You to the Licensor shall be under the terms and conditions of
+   this License, without any additional terms or conditions.
+   Notwithstanding the above, nothing herein shall supersede or modify
+   the terms of any separate license agreement you may have executed
+   with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+   names, trademarks, service marks, or product names of the Licensor,
+   except as required for reasonable and customary use in describing the
+   origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+   agreed to in writing, Licensor provides the Work (and each
+   Contributor provides its Contributions) on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+   implied, including, without limitation, any warranties or conditions
+   of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+   PARTICULAR PURPOSE. You are solely responsible for determining the
+   appropriateness of using or redistributing the Work and assume any
+   risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+   whether in tort (including negligence), contract, or otherwise,
+   unless required by applicable law (such as deliberate and grossly
+   negligent acts) or agreed to in writing, shall any Contributor be
+   liable to You for damages, including any direct, indirect, special,
+   incidental, or consequential damages of any character arising as a
+   result of this License or out of the use or inability to use the
+   Work (including but not limited to damages for loss of goodwill,
+   work stoppage, computer failure or malfunction, or any and all
+   other commercial damages or losses), even if such Contributor
+   has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+   the Work or Derivative Works thereof, You may choose to offer,
+   and charge a fee for, acceptance of support, warranty, indemnity,
+   or other liability obligations and/or rights consistent with this
+   License. However, in accepting such obligations, You may act only
+   on Your own behalf and on Your sole responsibility, not on behalf
+   of any other Contributor, and only if You agree to indemnify,
+   defend, and hold each Contributor harmless for any liability
+   incurred by, or claims asserted against, such Contributor by reason
+   of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+Copyright The OpenTelemetry Authors
+
+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.

openinference_instrumentation_haystack-0.1.0/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.3
+Name: openinference-instrumentation-haystack
+Version: 0.1.0
+Summary: OpenInference Haystack Instrumentation
+Project-URL: Homepage, https://github.com/Arize-ai/openinference/tree/main/python/instrumentation/openinference-instrumentation-haystack
+Author-email: OpenInference Authors <oss@arize.com>
+License-Expression: Apache-2.0
+License-File: LICENSE
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: <3.13,>=3.8
+Requires-Dist: openinference-instrumentation>=0.1.7
+Requires-Dist: openinference-semantic-conventions
+Requires-Dist: opentelemetry-api
+Requires-Dist: opentelemetry-instrumentation
+Requires-Dist: opentelemetry-semantic-conventions
+Requires-Dist: typing-extensions
+Requires-Dist: wrapt
+Provides-Extra: instruments
+Requires-Dist: haystack-ai>=2.0.0; extra == 'instruments'
+Provides-Extra: test
+Requires-Dist: haystack-ai>=2.0.0; extra == 'test'
+Requires-Dist: opentelemetry-sdk; extra == 'test'
+Requires-Dist: responses; extra == 'test'
+Description-Content-Type: text/markdown
+
+# OpenInference Haystack Instrumentation
+
+Python auto-instrumentation library for LLM applications implemented with [Haystack](https://haystack.deepset.ai/).
+
+Haystack [Pipelines](https://docs.haystack.deepset.ai/docs/pipelines) and [Components](https://docs.haystack.deepset.ai/docs/components) (ex. PromptBuilder, OpenAIGenerator, etc.) are fully OpenTelemetry-compatible and can be sent to an OpenTelemetry collector for monitoring, such as [`arize-phoenix`](https://github.com/Arize-ai/phoenix).
+
+## Installation
+
+```shell
+pip install openinference-instrumentation-haystack
+```
+
+## Quickstart
+
+This quickstart shows you how to instrument your Haystack-orchestrated LLM application
+
+Through your *terminal*, install required packages.
+
+```shell
+pip install openinference-instrumentation-haystack haystack arize-phoenix opentelemetry-sdk opentelemetry-exporter-otlp
+```
+
+You can install Phoenix and start it with the following terminal commands:
+```shell
+pip install arize-phoenix
+python -m phoenix.server.main serve
+````
+Start Phoenix in the background as a collector. By default, it listens on `http://localhost:6006`. You can visit the app via a browser at the same address. (Phoenix does not send data over the internet. It only operates locally on your machine.)
+
+
+Try the following in a *Python file*.
+
+Set up `HaystackInstrumentor` to trace your application and sends the traces to Phoenix at the endpoint defined below.
+
+```python
+from openinference.instrumentation.haystack import HaystackInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+import os
+
+# Set your OpenAI API key
+os.environ["OPENAI_API_KEY"] = "YOUR_KEY_HERE"
+
+# Set up the tracer, using Arize Phoenix as the endpoint
+endpoint = "http://127.0.0.1:6006/v1/traces"
+tracer_provider = trace_sdk.TracerProvider()
+trace_api.set_tracer_provider(tracer_provider)
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+
+# Instrument the Haystack application
+HaystackInstrumentor().instrument()
+```
+
+Set up a simple Pipeline with a template using `OpenAIGenerator`.
+```python
+from haystack import Pipeline
+from haystack.components.generators import OpenAIGenerator
+
+# Initialize the pipeline
+pipeline = Pipeline()
+
+# Initialize the OpenAI generator component
+llm = OpenAIGenerator(model="gpt-3.5-turbo")
+
+# Add the generator component to the pipeline
+pipeline.add_component("llm", llm)
+
+# Define the question
+question = "What is the location of the Hanging Gardens of Babylon?"
+
+# Run the pipeline with the question
+response = pipeline.run({"llm": {"prompt": question}})
+
+print(response)
+```
+Now, on the Phoenix UI on your browser, you should see the traces from your Haystack application. Specifically, you can see attributes from the execution of the OpenAIGenerator.
+

openinference_instrumentation_haystack-0.1.0/README.md

@@ -0,0 +1,79 @@
+# OpenInference Haystack Instrumentation
+
+Python auto-instrumentation library for LLM applications implemented with [Haystack](https://haystack.deepset.ai/).
+
+Haystack [Pipelines](https://docs.haystack.deepset.ai/docs/pipelines) and [Components](https://docs.haystack.deepset.ai/docs/components) (ex. PromptBuilder, OpenAIGenerator, etc.) are fully OpenTelemetry-compatible and can be sent to an OpenTelemetry collector for monitoring, such as [`arize-phoenix`](https://github.com/Arize-ai/phoenix).
+
+## Installation
+
+```shell
+pip install openinference-instrumentation-haystack
+```
+
+## Quickstart
+
+This quickstart shows you how to instrument your Haystack-orchestrated LLM application
+
+Through your *terminal*, install required packages.
+
+```shell
+pip install openinference-instrumentation-haystack haystack arize-phoenix opentelemetry-sdk opentelemetry-exporter-otlp
+```
+
+You can install Phoenix and start it with the following terminal commands:
+```shell
+pip install arize-phoenix
+python -m phoenix.server.main serve
+````
+Start Phoenix in the background as a collector. By default, it listens on `http://localhost:6006`. You can visit the app via a browser at the same address. (Phoenix does not send data over the internet. It only operates locally on your machine.)
+
+
+Try the following in a *Python file*.
+
+Set up `HaystackInstrumentor` to trace your application and sends the traces to Phoenix at the endpoint defined below.
+
+```python
+from openinference.instrumentation.haystack import HaystackInstrumentor
+from opentelemetry import trace as trace_api
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+import os
+
+# Set your OpenAI API key
+os.environ["OPENAI_API_KEY"] = "YOUR_KEY_HERE"
+
+# Set up the tracer, using Arize Phoenix as the endpoint
+endpoint = "http://127.0.0.1:6006/v1/traces"
+tracer_provider = trace_sdk.TracerProvider()
+trace_api.set_tracer_provider(tracer_provider)
+tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))
+
+# Instrument the Haystack application
+HaystackInstrumentor().instrument()
+```
+
+Set up a simple Pipeline with a template using `OpenAIGenerator`.
+```python
+from haystack import Pipeline
+from haystack.components.generators import OpenAIGenerator
+
+# Initialize the pipeline
+pipeline = Pipeline()
+
+# Initialize the OpenAI generator component
+llm = OpenAIGenerator(model="gpt-3.5-turbo")
+
+# Add the generator component to the pipeline
+pipeline.add_component("llm", llm)
+
+# Define the question
+question = "What is the location of the Hanging Gardens of Babylon?"
+
+# Run the pipeline with the question
+response = pipeline.run({"llm": {"prompt": question}})
+
+print(response)
+```
+Now, on the Phoenix UI on your browser, you should see the traces from your Haystack application. Specifically, you can see attributes from the execution of the OpenAIGenerator.
+

openinference_instrumentation_haystack-0.1.0/pyproject.toml

@@ -0,0 +1,91 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "openinference-instrumentation-haystack"
+dynamic = ["version"]
+description = "OpenInference Haystack Instrumentation"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.8, <3.13"
+authors = [
+    { name = "OpenInference Authors", email = "oss@arize.com" },
+]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+    "opentelemetry-api",
+    "opentelemetry-instrumentation",
+    "opentelemetry-semantic-conventions",
+    "openinference-instrumentation>=0.1.7",
+    "openinference-semantic-conventions",
+    "wrapt",
+    "typing-extensions",
+]
+
+
+[project.optional-dependencies]
+instruments = [
+    "haystack-ai >= 2.0.0",
+]
+test = [
+    "haystack-ai >= 2.0.0",
+    "opentelemetry-sdk",
+    "responses",
+]
+
+[project.urls]
+Homepage = "https://github.com/Arize-ai/openinference/tree/main/python/instrumentation/openinference-instrumentation-haystack"
+
+[tool.hatch.version]
+path = "src/openinference/instrumentation/haystack/version.py"
+
+[tool.hatch.build.targets.sdist]
+include = [
+    "/src",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/openinference"]
+
+[tool.mypy]
+strict = true
+explicit_package_bases = true
+exclude = [
+    "examples",
+    "dist",
+    "sdist",
+]
+
+[[tool.mypy.overrides]]
+ignore_missing_imports = true
+module = [
+    "haystack.*",
+    "datasets",
+    "openinference.instrumentation.haystack.*",
+    "wrapt",
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py38"
+
+[tool.ruff.lint.per-file-ignores]
+"*.ipynb" = ["E402", "E501"]
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+
+[tool.ruff.lint.isort]
+force-single-line = false

openinference_instrumentation_haystack-0.1.0/src/openinference/instrumentation/haystack/__init__.py

@@ -0,0 +1,53 @@
+import logging
+from importlib import import_module
+from typing import Any, Collection
+
+from openinference.instrumentation.haystack._wrappers import _ComponentWrapper, _PipelineWrapper
+from openinference.instrumentation.haystack.version import __version__
+from opentelemetry import trace as trace_api
+from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
+    BaseInstrumentor,
+)
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper
+
+logger = logging.getLogger(__name__)
+
+_instruments = ("haystack-ai >= 2.0.0",)
+
+
+class HaystackInstrumentor(BaseInstrumentor):  # type: ignore[misc]
+    """An instrumentor for the Haystack framework."""
+
+    __slots__ = ("_original_pipeline_run", "_original_pipeline_run_component")
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any) -> None:
+        if not (tracer_provider := kwargs.get("tracer_provider")):
+            tracer_provider = trace_api.get_tracer_provider()
+        tracer = get_tracer(__name__, __version__, tracer_provider)
+        pipeline = import_module("haystack.core.pipeline.pipeline")
+
+        self._original_pipeline_run = pipeline.Pipeline.run
+        wrap_function_wrapper(
+            module="haystack.core.pipeline.pipeline",
+            name="Pipeline.run",
+            wrapper=_PipelineWrapper(tracer=tracer),
+        )
+        self._original_pipeline_run_component = pipeline.Pipeline._run_component
+        wrap_function_wrapper(
+            module="haystack.core.pipeline.pipeline",
+            name="Pipeline._run_component",
+            wrapper=_ComponentWrapper(tracer=tracer),
+        )
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        haystack = import_module("haystack.core.pipeline.pipeline")
+
+        if self._original_pipeline_run is not None:
+            haystack.Pipeline.run = self._original_pipeline_run
+
+        if self._original_pipeline_run_component is not None:
+            haystack.Pipeline._run_component = self._original_pipeline_run_component

openinference_instrumentation_haystack-0.1.0/src/openinference/instrumentation/haystack/_wrappers.py

@@ -0,0 +1,254 @@
+from abc import ABC
+from enum import Enum
+from typing import Any, Callable, Iterator, List, Mapping, Tuple
+
+import opentelemetry.context as context_api
+from openinference.instrumentation import get_attributes_from_context, safe_json_dumps
+from openinference.semconv.trace import (
+    DocumentAttributes,
+    EmbeddingAttributes,
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+)
+from opentelemetry import trace as trace_api
+from opentelemetry.util.types import AttributeValue
+
+
+def _flatten(mapping: Mapping[str, Any]) -> Iterator[Tuple[str, AttributeValue]]:
+    for key, value in mapping.items():
+        if value is None:
+            continue
+        if isinstance(value, Mapping):
+            for sub_key, sub_value in _flatten(value):
+                yield f"{key}.{sub_key}", sub_value
+        elif isinstance(value, List) and any(isinstance(item, Mapping) for item in value):
+            for index, sub_mapping in enumerate(value):
+                for sub_key, sub_value in _flatten(sub_mapping):
+                    yield f"{key}.{index}.{sub_key}", sub_value
+        else:
+            if isinstance(value, Enum):
+                value = value.value
+            yield key, value
+
+
+class _WithTracer(ABC):
+    """
+    Base class for wrappers that need a tracer.
+    """
+
+    def __init__(self, tracer: trace_api.Tracer, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._tracer = tracer
+
+
+class _ComponentWrapper(_WithTracer):
+    """
+    Wrapper for the pipeline processing
+    Captures all calls to the pipeline
+    """
+
+    def __call__(
+        self,
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: Tuple[Any, ...],
+        kwargs: Mapping[str, Any],
+    ) -> Any:
+        component_type = args[0]
+        input_data = args[1]
+
+        # Diving into the instance to retrieve the Component instance
+        component = instance.graph.nodes._nodes[component_type]["instance"]
+
+        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        # Prepare invocation parameters by merging args and kwargs
+        invocation_parameters = {}
+        for arg in args:
+            if arg and isinstance(arg, dict):
+                invocation_parameters.update(arg)
+        invocation_parameters.update(kwargs)
+
+        span_name = component.__class__.__name__
+
+        with self._tracer.start_as_current_span(span_name, attributes={}) as span:
+            if component_type == "llm":
+                attributes = dict(
+                    _flatten(
+                        {
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND: LLM,
+                            SpanAttributes.INPUT_VALUE: safe_json_dumps(input_data),
+                            SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                            SpanAttributes.LLM_PROMPT_TEMPLATE: instance.graph.nodes._nodes[
+                                "prompt_builder"
+                            ]["instance"]._template_string,
+                        }
+                    )
+                )
+            elif component_type == "text_embedder":
+                attributes = dict(
+                    _flatten(
+                        {
+                            SpanAttributes.EMBEDDING_MODEL_NAME: component.model,
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND: EMBEDDING,
+                            SpanAttributes.INPUT_VALUE: safe_json_dumps(invocation_parameters),
+                            SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                        }
+                    )
+                )
+            elif component_type == "retriever":
+                attributes = dict(
+                    _flatten(
+                        {
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND: RETRIEVER,
+                            SpanAttributes.INPUT_VALUE: safe_json_dumps(
+                                {
+                                    "dimensions": len(input_data["query_embedding"]),
+                                    "embedding": input_data["query_embedding"],
+                                }
+                            ),
+                            SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                        }
+                    )
+                )
+
+            elif component_type == "prompt_builder":
+                attributes = dict(
+                    _flatten(
+                        {
+                            SpanAttributes.OPENINFERENCE_SPAN_KIND: CHAIN,
+                            SpanAttributes.LLM_PROMPT_TEMPLATE: instance.graph.nodes._nodes[
+                                "prompt_builder"
+                            ]["instance"]._template_string,
+                            SpanAttributes.INPUT_VALUE: safe_json_dumps(invocation_parameters),
+                            SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                            SpanAttributes.RETRIEVAL_DOCUMENTS: safe_json_dumps(
+                                invocation_parameters["documents"]
+                            ),
+                        }
+                    )
+                )
+
+            try:
+                response = wrapped(*args, **kwargs)
+            except Exception as exception:
+                span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, str(exception)))
+                span.record_exception(exception)
+                raise
+            span.set_status(trace_api.StatusCode.OK)
+
+            if component_type == "llm":
+                attributes.update(
+                    _flatten(
+                        {
+                            SpanAttributes.LLM_MODEL_NAME: response["meta"][0]["model"],
+                            SpanAttributes.LLM_TOKEN_COUNT_COMPLETION: response["meta"][0]["usage"][
+                                "completion_tokens"
+                            ],
+                            SpanAttributes.LLM_TOKEN_COUNT_PROMPT: response["meta"][0]["usage"][
+                                "prompt_tokens"
+                            ],
+                            SpanAttributes.LLM_TOKEN_COUNT_TOTAL: response["meta"][0]["usage"][
+                                "total_tokens"
+                            ],
+                            SpanAttributes.LLM_OUTPUT_MESSAGES: response["replies"],
+                        }
+                    )
+                )
+            elif component_type == "text_embedder":
+                emb_len = len(response["embedding"])
+                emb_vec_0 = f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.0."
+
+                attributes.update(
+                    _flatten(
+                        {
+                            f"{emb_vec_0}{EMBEDDING_VECTOR}": f"<{emb_len} dimensional vector>",
+                            f"{emb_vec_0}{EMBEDDING_TEXT}": invocation_parameters["text"],
+                        }
+                    )
+                )
+            elif component_type == "retriever":
+                attributes.update(
+                    _flatten(
+                        {
+                            SpanAttributes.OUTPUT_VALUE: safe_json_dumps(response["documents"]),
+                            SpanAttributes.OUTPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                        }
+                    )
+                )
+
+                for i, document in enumerate(response["documents"]):
+                    attributes[
+                        f"{SpanAttributes.RETRIEVAL_DOCUMENTS}.{i}.{DocumentAttributes.DOCUMENT_CONTENT}"
+                    ] = document.content
+                    attributes[
+                        f"{SpanAttributes.RETRIEVAL_DOCUMENTS}.{i}.{DocumentAttributes.DOCUMENT_ID}"
+                    ] = document.id
+                    attributes[
+                        f"{SpanAttributes.RETRIEVAL_DOCUMENTS}.{i}.{DocumentAttributes.DOCUMENT_SCORE}"
+                    ] = document.score
+                    attributes[
+                        f"{SpanAttributes.RETRIEVAL_DOCUMENTS}.{i}.{DocumentAttributes.DOCUMENT_METADATA}"
+                    ] = safe_json_dumps(document.meta)
+
+            span.set_attributes(attributes)
+
+        return response
+
+
+class _PipelineWrapper(_WithTracer):
+    """
+    Wrapper for the pipeline processing
+    Captures all calls to the pipeline
+    """
+
+    def __call__(
+        self,
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: Tuple[Any, ...],
+        kwargs: Mapping[str, Any],
+    ) -> Any:
+        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        # Prepare invocation parameters by merging args and kwargs
+        invocation_parameters = {}
+        for arg in args:
+            if arg and isinstance(arg, dict):
+                invocation_parameters.update(arg)
+        invocation_parameters.update(kwargs)
+
+        span_name = "Pipeline"
+        with self._tracer.start_as_current_span(
+            span_name,
+            attributes=dict(
+                _flatten(
+                    {
+                        SpanAttributes.OPENINFERENCE_SPAN_KIND: CHAIN,
+                        SpanAttributes.INPUT_VALUE: safe_json_dumps(invocation_parameters),
+                        SpanAttributes.INPUT_MIME_TYPE: OpenInferenceMimeTypeValues.JSON,
+                    }
+                )
+            ),
+        ) as span:
+            span.set_attributes(dict(get_attributes_from_context()))
+            try:
+                response = wrapped(*args, **kwargs)
+            except Exception as exception:
+                span.set_status(trace_api.Status(trace_api.StatusCode.ERROR, str(exception)))
+                span.record_exception(exception)
+                raise
+            span.set_status(trace_api.StatusCode.OK)
+
+        return response
+
+
+CHAIN = OpenInferenceSpanKindValues.CHAIN
+RETRIEVER = OpenInferenceSpanKindValues.RETRIEVER
+EMBEDDING = OpenInferenceSpanKindValues.EMBEDDING
+LLM = OpenInferenceSpanKindValues.LLM
+EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
+EMBEDDING_TEXT = EmbeddingAttributes.EMBEDDING_TEXT

openinference_instrumentation_haystack-0.1.0/src/openinference/instrumentation/haystack/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"