Skip to content

Commit

Permalink
Mount chroma integration under haystack_integrations.* (#193)
Browse files Browse the repository at this point in the history
* mount chroma integration under haystack_integrations.*

* sort imports

* elax relative imports linter

* fix mypy configuration for native namespaces

* specify what to put in the wheel
  • Loading branch information
masci authored Jan 17, 2024
1 parent 928c6fa commit 59b1017
Show file tree
Hide file tree
Showing 11 changed files with 30 additions and 28 deletions.
4 changes: 2 additions & 2 deletions integrations/chroma/example/example.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
from haystack.components.converters import TextFileToDocument
from haystack.components.writers import DocumentWriter

from chroma_haystack import ChromaDocumentStore
from chroma_haystack.retriever import ChromaQueryRetriever
from haystack_integrations.document_stores.chroma import ChromaDocumentStore
from haystack_integrations.components.retrievers.chroma import ChromaQueryRetriever

HERE = Path(__file__).resolve().parent
file_paths = [HERE / "data" / Path(name) for name in os.listdir("data")]
Expand Down
15 changes: 9 additions & 6 deletions integrations/chroma/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@ Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/m
Issues = "https://github.com/deepset-ai/haystack-core-integrations/issues"
Source = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/chroma"

[tool.hatch.build.targets.wheel]
packages = ["src/haystack_integrations"]

[tool.hatch.version]
source = "vcs"
tag-pattern = 'integrations\/chroma-v(?P<version>.*)'
Expand Down Expand Up @@ -69,7 +72,7 @@ dependencies = [
"numpy", # we need the stubs from the main package
]
[tool.hatch.envs.lint.scripts]
typing = "mypy --install-types --non-interactive {args:src/chroma_haystack tests}"
typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
style = [
"ruff {args:.}",
"black --check --diff {args:.}",
Expand Down Expand Up @@ -141,26 +144,25 @@ unfixable = [
exclude = ["example"]

[tool.ruff.isort]
known-first-party = ["chroma_haystack"]
known-first-party = ["src"]

[tool.ruff.flake8-tidy-imports]
ban-relative-imports = "all"
ban-relative-imports = "parents"

[tool.ruff.per-file-ignores]
# Tests can use magic values, assertions, and relative imports
"tests/**/*" = ["PLR2004", "S101", "TID252"]

[tool.coverage.run]
source_pkgs = ["chroma_haystack", "tests"]
source_pkgs = ["src", "tests"]
branch = true
parallel = true
omit = [
"src/chroma_haystack/__about__.py",
"example"
]

[tool.coverage.paths]
chroma_haystack = ["src/chroma_haystack", "*/chroma-haystack/src/chroma_haystack"]
chroma_haystack = ["src/haystack_integrations", "*/chroma-haystack/src/chroma_haystack"]
tests = ["tests"]

[tool.coverage.report]
Expand All @@ -181,6 +183,7 @@ markers = [
module = [
"chromadb.*",
"haystack.*",
"haystack_integrations.*",
"pytest.*"
]
ignore_missing_imports = true
6 changes: 0 additions & 6 deletions integrations/chroma/src/chroma_haystack/__init__.py

This file was deleted.

Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .retriever import ChromaEmbeddingRetriever, ChromaQueryRetriever, ChromaSingleQueryRetriever

__all__ = ["ChromaQueryRetriever", "ChromaEmbeddingRetriever", "ChromaSingleQueryRetriever"]
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,7 @@
from typing import Any, Dict, List, Optional

from haystack import Document, component, default_from_dict, default_to_dict

from chroma_haystack import ChromaDocumentStore
from haystack_integrations.document_stores.chroma import ChromaDocumentStore


@component
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .document_store import ChromaDocumentStore

__all__ = ["ChromaDocumentStore"]
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,10 @@
import numpy as np
from chromadb.api.types import GetResult, QueryResult, validate_where, validate_where_document
from haystack.dataclasses import Document
from haystack.document_stores.protocol import DuplicatePolicy
from haystack.document_stores.types import DuplicatePolicy

from chroma_haystack.errors import ChromaDocumentStoreFilterError
from chroma_haystack.utils import get_embedding_function
from .errors import ChromaDocumentStoreFilterError
from .utils import get_embedding_function

logger = logging.getLogger(__name__)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
Text2VecEmbeddingFunction,
)

from chroma_haystack.errors import ChromaDocumentStoreConfigError
from .errors import ChromaDocumentStoreConfigError

FUNCTION_REGISTRY = {
"default": DefaultEmbeddingFunction,
Expand Down
7 changes: 4 additions & 3 deletions integrations/chroma/tests/test_document_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,7 @@
DeleteDocumentsTest,
LegacyFilterDocumentsTest,
)

from chroma_haystack.document_store import ChromaDocumentStore
from haystack_integrations.document_stores.chroma import ChromaDocumentStore


class _TestEmbeddingFunction(EmbeddingFunction):
Expand All @@ -42,7 +41,9 @@ def document_store(self) -> ChromaDocumentStore:
This is the most basic requirement for the child class: provide
an instance of this document store so the base class can use it.
"""
with mock.patch("chroma_haystack.document_store.get_embedding_function") as get_func:
with mock.patch(
"haystack_integrations.document_stores.chroma.document_store.get_embedding_function"
) as get_func:
get_func.return_value = _TestEmbeddingFunction()
return ChromaDocumentStore(embedding_function="test_function", collection_name=str(uuid.uuid1()))

Expand Down
9 changes: 4 additions & 5 deletions integrations/chroma/tests/test_retriever.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,8 @@
#
# SPDX-License-Identifier: Apache-2.0
import pytest

from chroma_haystack.document_store import ChromaDocumentStore
from chroma_haystack.retriever import ChromaQueryRetriever
from haystack_integrations.components.retrievers.chroma import ChromaQueryRetriever
from haystack_integrations.document_stores.chroma import ChromaDocumentStore


@pytest.mark.integration
Expand All @@ -14,7 +13,7 @@ def test_retriever_to_json(request):
)
retriever = ChromaQueryRetriever(ds, filters={"foo": "bar"}, top_k=99)
assert retriever.to_dict() == {
"type": "chroma_haystack.retriever.ChromaQueryRetriever",
"type": "haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryRetriever",
"init_parameters": {
"filters": {"foo": "bar"},
"top_k": 99,
Expand All @@ -30,7 +29,7 @@ def test_retriever_to_json(request):
@pytest.mark.integration
def test_retriever_from_json(request):
data = {
"type": "chroma_haystack.retriever.ChromaQueryRetriever",
"type": "haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryRetriever",
"init_parameters": {
"filters": {"bar": "baz"},
"top_k": 42,
Expand Down

0 comments on commit 59b1017

Please sign in to comment.