code review

docs/quickstart/quickstart2_code.py

@@ -77,7 +77,7 @@ def from_country(self, country: Annotated[str, country_similarity]) -> sqlalchem
 
 
 async def main():
-    dbally.event_handlers.append(CLIEventHandler())
+    dbally.event_handlers = [CLIEventHandler()]
     await country_similarity.update()
 
     llm = LiteLLM(model_name="gpt-3.5-turbo")

docs/quickstart/quickstart3_code.py

@@ -126,7 +126,7 @@ def display_results(result: ExecutionResult):
 
 
 async def main():
-    dbally.event_handlers.append(CLIEventHandler())
+    dbally.event_handlers = [CLIEventHandler()]
     await country_similarity.update()
 
     llm = LiteLLM(model_name="gpt-3.5-turbo")

docs/quickstart/quickstart_code.py

@@ -58,7 +58,7 @@ def from_country(self, country: str) -> sqlalchemy.ColumnElement:
 
 async def main():
     llm = LiteLLM(model_name="gpt-3.5-turbo")
-    dbally.event_handlers.append(CLIEventHandler())
+    dbally.event_handlers = [CLIEventHandler()]
 
     collection = dbally.create_collection("recruitment", llm)
     collection.add(CandidateView, lambda: CandidateView(engine))

src/dbally/__init__.py

@@ -1,8 +1,7 @@
 """ dbally """
 
-from typing import List
+from typing import TYPE_CHECKING, List
 
-from dbally.collection.collection import Collection, create_collection
 from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
 from dbally.collection.results import ExecutionResult
 from dbally.views import decorators
@@ -12,6 +11,7 @@
 from dbally.views.structured import BaseStructuredView
 
 from .__version__ import __version__
+from ._main import create_collection
 from ._types import NOT_GIVEN, NotGiven
 from .embeddings.exceptions import (
     EmbeddingConnectionError,
@@ -22,15 +22,17 @@
 from .exceptions import DbAllyError
 from .llms.clients.exceptions import LLMConnectionError, LLMError, LLMResponseError, LLMStatusError
 
-event_handlers: List = []
+if TYPE_CHECKING:
+    from .audit import EventHandler
+
+event_handlers: List["EventHandler"] = []
 
 __all__ = [
     "__version__",
     "create_collection",
     "decorators",
     "event_handlers",
     "BaseStructuredView",
-    "Collection",
     "DataFrameBaseView",
     "DbAllyError",
     "ExecutionResult",

src/dbally/_main.py

@@ -0,0 +1,68 @@
+from typing import TYPE_CHECKING, List, Optional
+
+from dbally.audit import EventHandler
+from dbally.llms import LLM
+from dbally.nl_responder.nl_responder import NLResponder
+from dbally.view_selection import LLMViewSelector
+from dbally.view_selection.base import ViewSelector
+
+if TYPE_CHECKING:
+    from dbally.collection import Collection
+
+
+def create_collection(
+    name: str,
+    llm: LLM,
+    event_handlers: Optional[List[EventHandler]] = None,
+    view_selector: Optional[ViewSelector] = None,
+    nl_responder: Optional[NLResponder] = None,
+) -> "Collection":
+    """
+    Create a new [Collection](collection.md) that is a container for registering views and the\
+    main entrypoint to db-ally features.
+
+    Unlike instantiating a [Collection][dbally.Collection] directly, this function\
+    provides a set of default values for various dependencies like LLM client, view selector,\
+    IQL generator, and NL responder.
+
+    ##Example
+
+    ```python
+        from dbally import create_collection
+        from dbally.llms.litellm import LiteLLM
+
+        collection = create_collection("my_collection", llm=LiteLLM())
+    ```
+
+    Args:
+        name: Name of the collection is available for [Event handlers](event_handlers/index.md) and is\
+        used to distinguish different db-ally runs.
+        llm: LLM used by the collection to generate responses for natural language queries.
+        event_handlers: Event handlers used by the collection during query executions. Can be used to\
+        log events as [CLIEventHandler](event_handlers/cli_handler.md) or to validate system performance as\
+        [LangSmithEventHandler](event_handlers/langsmith_handler.md). If provided, this parameter overrides the
+        global dbally.event_handlers_list
+        view_selector: View selector used by the collection to select the best view for the given query.\
+        If None, a new instance of [LLMViewSelector][dbally.view_selection.llm_view_selector.LLMViewSelector]\
+        will be used.
+        nl_responder: NL responder used by the collection to respond to natural language queries. If None,\
+        a new instance of [NLResponder][dbally.nl_responder.nl_responder.NLResponder] will be used.
+
+    Returns:
+        a new instance of db-ally Collection
+
+    Raises:
+        ValueError: if default LLM client is not configured
+    """
+    from dbally.collection import Collection  # pylint: disable=import-outside-toplevel
+
+    view_selector = view_selector or LLMViewSelector(llm=llm)
+    nl_responder = nl_responder or NLResponder(llm=llm)
+
+    return Collection(
+        name,
+        nl_responder=nl_responder,
+        view_selector=view_selector,
+        llm=llm,
+        event_handlers=event_handlers,
+    )

src/dbally/collection/__init__.py

@@ -1,9 +1,8 @@
-from dbally.collection.collection import Collection, create_collection
+from dbally.collection.collection import Collection
 from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
 from dbally.collection.results import ExecutionResult, ViewExecutionResult
 
 __all__ = [
-    "create_collection",
     "Collection",
     "ExecutionResult",
     "ViewExecutionResult",

src/dbally/collection/collection.py

@@ -2,7 +2,6 @@
 import inspect
 import textwrap
 import time
-import warnings
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Type, TypeVar
 
@@ -16,7 +15,6 @@
 from dbally.llms.clients.base import LLMOptions
 from dbally.nl_responder.nl_responder import NLResponder
 from dbally.similarity.index import AbstractSimilarityIndex
-from dbally.view_selection import LLMViewSelector
 from dbally.view_selection.base import ViewSelector
 from dbally.views.base import BaseView, IndexLocation
 
@@ -63,13 +61,7 @@ def __init__(
         self._nl_responder = nl_responder
         self._llm = llm
 
-        if not event_handlers:
-            event_handlers = dbally.event_handlers
-        elif event_handlers != dbally.event_handlers:
-            # At this moment, there is no event tracker initialized to record an event
-            warnings.warn("Default event handler has been overwritten for {self.name}.")
-
-        self._event_handlers = event_handlers
+        self._event_handlers = event_handlers or dbally.event_handlers
 
     T = TypeVar("T", bound=BaseView)
 
@@ -276,59 +268,3 @@ async def update_similarity_indexes(self) -> None:
         if failed_indexes:
             failed_locations = [loc for index in failed_indexes for loc in indexes[index]]
             raise IndexUpdateError(failed_indexes, failed_locations)
-
-
-def create_collection(
-    name: str,
-    llm: LLM,
-    event_handlers: Optional[List[EventHandler]] = None,
-    view_selector: Optional[ViewSelector] = None,
-    nl_responder: Optional[NLResponder] = None,
-) -> Collection:
-    """
-    Create a new [Collection](collection.md) that is a container for registering views and the\
-    main entrypoint to db-ally features.
-
-    Unlike instantiating a [Collection][dbally.Collection] directly, this function\
-    provides a set of default values for various dependencies like LLM client, view selector,\
-    IQL generator, and NL responder.
-
-    ##Example
-
-    ```python
-        from dbally import create_collection
-        from dbally.llms.litellm import LiteLLM
-
-        collection = create_collection("my_collection", llm=LiteLLM())
-    ```
-
-    Args:
-        name: Name of the collection is available for [Event handlers](event_handlers/index.md) and is\
-        used to distinguish different db-ally runs.
-        llm: LLM used by the collection to generate responses for natural language queries.
-        event_handlers: Event handlers used by the collection during query executions. Can be used to\
-        log events as [CLIEventHandler](event_handlers/cli_handler.md) or to validate system performance as\
-        [LangSmithEventHandler](event_handlers/langsmith_handler.md). If provided, this parameter overrides the
-        global dbally.event_handlers_list
-        view_selector: View selector used by the collection to select the best view for the given query.\
-        If None, a new instance of [LLMViewSelector][dbally.view_selection.llm_view_selector.LLMViewSelector]\
-        will be used.
-        nl_responder: NL responder used by the collection to respond to natural language queries. If None,\
-        a new instance of [NLResponder][dbally.nl_responder.nl_responder.NLResponder] will be used.
-
-    Returns:
-        a new instance of db-ally Collection
-
-    Raises:
-        ValueError: if default LLM client is not configured
-    """
-    view_selector = view_selector or LLMViewSelector(llm=llm)
-    nl_responder = nl_responder or NLResponder(llm=llm)
-
-    return Collection(
-        name,
-        nl_responder=nl_responder,
-        view_selector=view_selector,
-        llm=llm,
-        event_handlers=event_handlers,
-    )

tests/unit/test_collection.py

@@ -6,7 +6,8 @@
 import pytest
 from typing_extensions import Annotated
 
-from dbally.collection import Collection, create_collection
+import dbally
+from dbally.collection import Collection
 from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
 from dbally.collection.results import ViewExecutionResult
 from dbally.iql._exceptions import IQLError
@@ -126,7 +127,7 @@ def mock_collection() -> Collection:
     """
     Returns a collection with two mock views
     """
-    collection = create_collection(
+    collection = dbally.create_collection(
         "foo",
         llm=MockLLM(),
         view_selector=MockViewSelector("MockView1"),