refactor: move ExecutionResult to collection (#46)

Author: micpst

benchmark/dbally_benchmark/e2e_benchmark.py

@@ -21,7 +21,8 @@
 from sqlalchemy import create_engine
 
 import dbally
-from dbally.collection import Collection, NoViewFoundError
+from dbally.collection import Collection
+from dbally.collection.exceptions import NoViewFoundError
 from dbally.iql_generator.iql_prompt_template import UnsupportedQueryError, default_iql_template
 from dbally.llms.litellm import LiteLLM
 from dbally.view_selection.view_selector_prompt_template import default_view_selector_template

docs/concepts/collections.md

@@ -26,6 +26,6 @@ my_collection.ask("Find me Italian recipes for soups")
 In this scenario, the LLM first determines the most suitable view to address the query, and then that view is used to pull the relevant data.
 
 !!! info
-    The result of a query is an [`ExecutionResult`][dbally.data_models.execution_result.ExecutionResult] object, which contains the data fetched by the view. It contains a `results` attribute that holds the actual data, structured as a list of dictionaries. The exact structure of these dictionaries depends on the view that was used to fetch the data, which can be obtained by looking at the `view_name` attribute of the `ExecutionResult` object.
+    The result of a query is an [`ExecutionResult`][dbally.collection.results.ExecutionResult] object, which contains the data fetched by the view. It contains a `results` attribute that holds the actual data, structured as a list of dictionaries. The exact structure of these dictionaries depends on the view that was used to fetch the data, which can be obtained by looking at the `view_name` attribute of the `ExecutionResult` object.
 
 It's possible for projects to feature several collections, each potentially housing a different set of views. Moreover, a single view can be associated with multiple collections, offering versatile usage across various contexts.

docs/how-to/views/custom.md

@@ -154,7 +154,7 @@ import abc
 from typing import Callable, Any, Iterable
 
 from dbally.iql import IQLQuery
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 
 @abc.abstractmethod
 def get_data(self) -> Iterable:

docs/how-to/views/custom_views_code.py

@@ -8,7 +8,7 @@
 from dbally import decorators, MethodsBaseView
 from dbally.audit.event_handlers.cli_event_handler import CLIEventHandler
 from dbally.iql import IQLQuery, syntax
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.llms.litellm import LiteLLM
 
 @dataclass

docs/quickstart/quickstart3.md

@@ -8,7 +8,7 @@ The guide illustrates how to use multiple views to handle queries requiring diff
 
 You can register multiple views with a collection. When you run a query, the AI model decides which view to use based on the query. This allows for handling diffrent kinds of queries with different views. Those views can be based on data from the same source (e.g., different tables in the same database), or from different sources (e.g. a database and an API).
 
-Upon selecting the view, the AI model uses it to extract the relevant data from its data source. The query result is an [`ExecutionResult`][dbally.data_models.execution_result.ExecutionResult] object. It contains the data extracted by the view, along with other metadata including the name of the view that fetched the data.
+Upon selecting the view, the AI model uses it to extract the relevant data from its data source. The query result is an [`ExecutionResult`][dbally.collection.results.ExecutionResult] object. It contains the data extracted by the view, along with other metadata including the name of the view that fetched the data.
 
 ## Defining the JobView
 

docs/reference/collection.md

@@ -5,8 +5,10 @@
 
 ::: dbally.create_collection
 
-::: dbally.Collection
+::: dbally.collection.Collection
 
-::: dbally.data_models.execution_result.ExecutionResult
+::: dbally.collection.results.ExecutionResult
 
-::: dbally.collection.IndexUpdateError
+::: dbally.collection.exceptions.IndexUpdateError
+
+::: dbally.collection.exceptions.NoViewFoundError

docs/reference/views/index.md

@@ -6,4 +6,4 @@
 
 ::: dbally.views.base.BaseView
 
-::: dbally.data_models.execution_result.ViewExecutionResult
+::: dbally.collection.results.ViewExecutionResult

src/dbally/__init__.py

@@ -1,6 +1,8 @@
 """ dbally """
 
-from dbally.data_models.execution_result import ExecutionResult
+from dbally.collection.collection import Collection
+from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
+from dbally.collection.results import ExecutionResult
 from dbally.views import decorators
 from dbally.views.methods_base import MethodsBaseView
 from dbally.views.pandas_base import DataFrameBaseView
@@ -10,7 +12,6 @@
 from .__version__ import __version__
 from ._main import create_collection
 from ._types import NOT_GIVEN, NotGiven
-from .collection import Collection
 from .embeddings.exceptions import (
     EmbeddingConnectionError,
     EmbeddingError,
@@ -39,6 +40,8 @@
     "LLMConnectionError",
     "LLMResponseError",
     "LLMStatusError",
+    "NoViewFoundError",
+    "IndexUpdateError",
     "NotGiven",
     "NOT_GIVEN",
 ]

src/dbally/collection/__init__.py

@@ -0,0 +1,11 @@
+from dbally.collection.collection import Collection
+from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
+from dbally.collection.results import ExecutionResult, ViewExecutionResult
+
+__all__ = [
+    "Collection",
+    "ExecutionResult",
+    "ViewExecutionResult",
+    "NoViewFoundError",
+    "IndexUpdateError",
+]

src/dbally/collection.py renamed to src/dbally/collection/collection.py

@@ -7,9 +7,9 @@
 
 from dbally.audit.event_handlers.base import EventHandler
 from dbally.audit.event_tracker import EventTracker
+from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
+from dbally.collection.results import ExecutionResult
 from dbally.data_models.audit import RequestEnd, RequestStart
-from dbally.data_models.execution_result import ExecutionResult
-from dbally.exceptions import DbAllyError
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.nl_responder.nl_responder import NLResponder
@@ -18,29 +18,6 @@
 from dbally.views.base import BaseView, IndexLocation
 
 
-class NoViewFoundError(DbAllyError):
-    """
-    Error raised when there is no view with the given name.
-    """
-
-
-class IndexUpdateError(DbAllyError):
-    """
-    Exception for when updating any of the Collection's similarity indexes fails.
-
-    Provides a dictionary mapping failed indexes to their
-    respective exceptions as the `failed_indexes` attribute.
-    """
-
-    def __init__(self, message: str, failed_indexes: Dict[AbstractSimilarityIndex, Exception]) -> None:
-        """
-        Args:
-            failed_indexes: Dictionary mapping failed indexes to their respective exceptions.
-        """
-        self.failed_indexes = failed_indexes
-        super().__init__(message)
-
-
 class Collection:
     """
     Collection is a container for a set of views that can be used by db-ally to answer user questions.
@@ -158,7 +135,7 @@ def get(self, name: str) -> BaseView:
         """
 
         if name not in self._views:
-            raise NoViewFoundError
+            raise NoViewFoundError(name)
 
         return self._builders[name]()
 
@@ -298,5 +275,4 @@ async def update_similarity_indexes(self) -> None:
         }
         if failed_indexes:
             failed_locations = [loc for index in failed_indexes for loc in indexes[index]]
-            descriptions = ", ".join(".".join(name for name in location) for location in failed_locations)
-            raise IndexUpdateError(f"Failed to update similarity indexes for {descriptions}", failed_indexes)
+            raise IndexUpdateError(failed_indexes, failed_locations)