Skip to content

Commit

Permalink
refactor: move ExecutionResult to collection (#46)
Browse files Browse the repository at this point in the history
  • Loading branch information
micpst committed Jun 14, 2024
1 parent 825b8fe commit eae24da
Show file tree
Hide file tree
Showing 23 changed files with 89 additions and 52 deletions.
3 changes: 2 additions & 1 deletion benchmark/dbally_benchmark/e2e_benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/concepts/collections.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
2 changes: 1 addition & 1 deletion docs/how-to/views/custom.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion docs/how-to/views/custom_views_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/quickstart/quickstart3.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 5 additions & 3 deletions docs/reference/collection.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion docs/reference/views/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,4 @@

::: dbally.views.base.BaseView

::: dbally.data_models.execution_result.ViewExecutionResult
::: dbally.collection.results.ViewExecutionResult
7 changes: 5 additions & 2 deletions src/dbally/__init__.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -39,6 +40,8 @@
"LLMConnectionError",
"LLMResponseError",
"LLMStatusError",
"NoViewFoundError",
"IndexUpdateError",
"NotGiven",
"NOT_GIVEN",
]
Expand Down
11 changes: 11 additions & 0 deletions src/dbally/collection/__init__.py
Original file line number Diff line number Diff line change
@@ -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",
]
32 changes: 4 additions & 28 deletions src/dbally/collection.py → src/dbally/collection/collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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]()

Expand Down Expand Up @@ -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)
42 changes: 42 additions & 0 deletions src/dbally/collection/exceptions.py
Original file line number Diff line number Diff line change
@@ -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
File renamed without changes.
2 changes: 1 addition & 1 deletion src/dbally/data_models/audit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
3 changes: 2 additions & 1 deletion src/dbally/gradio/gradio_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/dbally/nl_responder/nl_responder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/dbally/views/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/dbally/views/freeform/text2sql/view.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/dbally/views/pandas_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/dbally/views/sqlalchemy_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/dbally/views/structured.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions tests/unit/test_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_nl_responder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion tests/unit/views/test_methods_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit eae24da

Please sign in to comment.