refactor: move ExecutionResult to collection (#46)

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 → 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)

src/dbally/collection/exceptions.py

@@ -0,0 +1,42 @@
+from typing import Dict, List
+
+from dbally.exceptions import DbAllyError
+from dbally.similarity.index import AbstractSimilarityIndex
+from dbally.views.base import IndexLocation
+
+
+class NoViewFoundError(DbAllyError):
+    """
+    Error raised when there is no view with the given name.
+    """
+
+    def __init__(self, view_name: str) -> None:
+        """
+        Args:
+            view_name: Name of the view that was not found.
+        """
+        super().__init__(f"No view found with name '{view_name}'.")
+        self.view_name = view_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,
+        failed_indexes: Dict[AbstractSimilarityIndex, Exception],
+        failed_locations: List[IndexLocation],
+    ) -> None:
+        """
+        Args:
+            failed_indexes: Dictionary mapping failed indexes to their respective exceptions.
+            failed_locations: List of locations of failed indexes.
+        """
+        description = ", ".join(".".join(name for name in location) for location in failed_locations)
+        super().__init__(f"Failed to update similarity indexes for {description}.")
+        self.failed_indexes = failed_indexes

src/dbally/data_models/audit.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from typing import Optional, Union
 
-from dbally.data_models.execution_result import ExecutionResult
+from dbally.collection.results import ExecutionResult
 from dbally.prompts import ChatFormat
 
 

src/dbally/gradio/gradio_interface.py

@@ -6,7 +6,8 @@
 
 from dbally import BaseStructuredView
 from dbally.audit import CLIEventHandler
-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
 from dbally.prompts import PromptTemplateError
 

src/dbally/nl_responder/nl_responder.py

@@ -4,7 +4,7 @@
 import pandas as pd
 
 from dbally.audit.event_tracker import EventTracker
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.nl_responder.nl_responder_prompt_template import NLResponderPromptTemplate, default_nl_responder_template

src/dbally/views/base.py

@@ -2,7 +2,7 @@
 from typing import Dict, List, Optional, Tuple
 
 from dbally.audit.event_tracker import EventTracker
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.similarity import AbstractSimilarityIndex

src/dbally/views/freeform/text2sql/view.py

@@ -7,7 +7,7 @@
 from sqlalchemy import ColumnClause, Engine, MetaData, Table, text
 
 from dbally.audit.event_tracker import EventTracker
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.llms.base import LLM
 from dbally.llms.clients.base import LLMOptions
 from dbally.prompts import PromptTemplate

src/dbally/views/pandas_base.py

@@ -3,7 +3,7 @@
 
 import pandas as pd
 
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.iql import IQLQuery, syntax
 from dbally.views.methods_base import MethodsBaseView
 

src/dbally/views/sqlalchemy_base.py

@@ -3,7 +3,7 @@
 
 import sqlalchemy
 
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.iql import IQLQuery, syntax
 from dbally.views.methods_base import MethodsBaseView
 

src/dbally/views/structured.py

@@ -3,7 +3,7 @@
 from typing import Dict, List, Optional
 
 from dbally.audit.event_tracker import EventTracker
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.iql import IQLError, IQLQuery
 from dbally.iql_generator.iql_generator import IQLGenerator
 from dbally.llms.base import LLM

tests/unit/test_collection.py

@@ -7,8 +7,9 @@
 from typing_extensions import Annotated
 
 from dbally._main import create_collection
-from dbally.collection import Collection, IndexUpdateError, NoViewFoundError
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection import Collection
+from dbally.collection.exceptions import IndexUpdateError, NoViewFoundError
+from dbally.collection.results import ViewExecutionResult
 from dbally.iql._exceptions import IQLError
 from dbally.views.exposed_functions import ExposedFunction, MethodParamWithTyping
 from dbally.views.structured import BaseStructuredView
@@ -438,7 +439,7 @@ async def test_update_similarity_indexes_error(
         await collection.update_similarity_indexes()
     assert (
         str(e.value) == "Failed to update similarity indexes for MockViewWithSimilarity.test_filter.dog, "
-        "MockViewWithSimilarity2.test_filter.monkey"
+        "MockViewWithSimilarity2.test_filter.monkey."
     )
     assert e.value.failed_indexes == {
         foo_index: foo_exception,

tests/unit/test_nl_responder.py

@@ -3,7 +3,7 @@
 import pytest
 
 from dbally.audit.event_tracker import EventTracker
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.nl_responder.nl_responder import NLResponder
 from tests.unit.mocks import MockLLM
 

tests/unit/views/test_methods_base.py

@@ -3,7 +3,7 @@
 
 from typing import List, Literal, Tuple
 
-from dbally.data_models.execution_result import ViewExecutionResult
+from dbally.collection.results import ViewExecutionResult
 from dbally.iql import IQLQuery
 from dbally.views.decorators import view_filter
 from dbally.views.exposed_functions import MethodParamWithTyping