Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions integrations/weights_and_biases_weave/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,23 +7,22 @@ name = "weave-haystack"
dynamic = ["version"]
description = "An integration of Weights & Biases Weave tracer with Haystack"
readme = "README.md"
requires-python = ">=3.9"
requires-python = ">=3.10"
license = "Apache-2.0"
keywords = []
authors = [{ name = "deepset GmbH", email = "[email protected]" }]
classifiers = [
"License :: OSI Approved :: Apache Software License",
"Development Status :: 4 - Beta",
"Programming Language :: Python",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
]
dependencies = ["haystack-ai", "weave>=0.52.7"]
dependencies = ["haystack-ai>=2.22.0", "weave>=0.52.7"]

[project.urls]
Source = "https://github.com/deepset-ai/haystack-core-integrations"
Expand Down Expand Up @@ -74,7 +73,6 @@ check_untyped_defs = true
disallow_incomplete_defs = true

[tool.ruff]
target-version = "py39"
line-length = 120

[tool.ruff.lint]
Expand Down Expand Up @@ -121,10 +119,6 @@ ignore = [
"PLR0913",
"PLR0915",
]
unfixable = [
# Don't touch unused imports
"F401",
]

[tool.ruff.lint.isort]
known-first-party = ["haystack_integrations"]
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, Optional
from typing import Any

from haystack import component, default_from_dict, default_to_dict, logging, tracing

Expand Down Expand Up @@ -71,7 +71,7 @@ class WeaveConnector:

"""

def __init__(self, pipeline_name: str, weave_init_kwargs: Optional[dict[str, Any]] = None) -> None:
def __init__(self, pipeline_name: str, weave_init_kwargs: dict[str, Any] | None = None) -> None:
"""
Initialize WeaveConnector.

Expand All @@ -80,7 +80,7 @@ def __init__(self, pipeline_name: str, weave_init_kwargs: Optional[dict[str, Any
"""
self.pipeline_name = pipeline_name
self.weave_init_kwargs = weave_init_kwargs or {}
self.tracer: Optional[WeaveTracer] = None
self.tracer: WeaveTracer | None = None

def warm_up(self) -> None:
"""Initialize the WeaveTracer."""
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import contextlib
import os
from collections.abc import Iterator
from typing import Any, Optional, Union
from typing import Any

from haystack import logging
from haystack.tracing import Span, Tracer
Expand All @@ -22,9 +22,7 @@ class WeaveSpan(Span):
that describe the operation.
"""

def __init__(
self, call: Optional[Call] = None, parent: Optional[Call] = None, operation: Optional[str] = None
) -> None:
def __init__(self, call: Call | None = None, parent: Call | None = None, operation: str | None = None) -> None:
self._call = call
self._parent = parent
self._operation = operation
Expand Down Expand Up @@ -86,12 +84,10 @@ def __init__(self, project_name: str, **weave_init_kwargs: Any) -> None:
)

self._client = weave.init(project_name, **weave_init_kwargs)
self._current_span: Optional[WeaveSpan] = None
self._current_span: WeaveSpan | None = None

@staticmethod
def create_call(
attributes: dict, client: WeaveClient, parent_span: Union[WeaveSpan, None], operation_name: str
) -> Call:
def create_call(attributes: dict, client: WeaveClient, parent_span: WeaveSpan | None, operation_name: str) -> Call:
comp_name = attributes.pop("haystack.component.name", "")
comp_type = attributes.pop("haystack.component.type", "")
comp_input = attributes.pop("haystack.component.input", {})
Expand All @@ -104,18 +100,16 @@ def create_call(
)
return call

def current_span(self) -> Optional[Span]:
def current_span(self) -> Span | None:
"""Get the current active span."""
return self._current_span

@staticmethod
def _create_component_span(parent_span: Optional[WeaveSpan], operation_name: str) -> WeaveSpan:
def _create_component_span(parent_span: WeaveSpan | None, operation_name: str) -> WeaveSpan:
"""Create a span for component runs with deferred call creation."""
return WeaveSpan(parent=parent_span.raw_span() if parent_span else None, operation=operation_name)

def _create_regular_span(
self, operation_name: str, tags: Optional[dict], parent_span: Optional[WeaveSpan]
) -> WeaveSpan:
def _create_regular_span(self, operation_name: str, tags: dict | None, parent_span: WeaveSpan | None) -> WeaveSpan:
"""Create a span for regular operations with immediate call creation."""
call = self._client.create_call(
op=operation_name,
Expand All @@ -124,7 +118,7 @@ def _create_regular_span(
)
return WeaveSpan(call=call)

def _finish_component_call(self, span: WeaveSpan, parent_span: Optional[WeaveSpan], operation_name: str) -> None:
def _finish_component_call(self, span: WeaveSpan, parent_span: WeaveSpan | None, operation_name: str) -> None:
"""Create and finish call for component runs."""
attributes = span.get_attributes()
call = self.create_call(
Expand Down Expand Up @@ -155,8 +149,8 @@ def _finish_regular_call(self, span: WeaveSpan) -> None:
def trace( # type: ignore[override]
self,
operation_name: str,
tags: Optional[dict[str, Any]] = None,
parent_span: Optional[WeaveSpan] = None,
tags: dict[str, Any] | None = None,
parent_span: WeaveSpan | None = None,
) -> Iterator[WeaveSpan]:
"""
A context manager that creates and manages spans for tracking operations in Weights & Biases Weave.
Expand Down