diff --git a/benchmark/dbally_benchmark/e2e_benchmark.py b/benchmark/dbally_benchmark/e2e_benchmark.py index 06d84287..46edce2f 100644 --- a/benchmark/dbally_benchmark/e2e_benchmark.py +++ b/benchmark/dbally_benchmark/e2e_benchmark.py @@ -23,7 +23,7 @@ import dbally from dbally.collection import Collection from dbally.iql_generator.iql_prompt_template import default_iql_template -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient from dbally.utils.errors import NoViewFoundError, UnsupportedQueryError from dbally.view_selection.view_selector_prompt_template import default_view_selector_template diff --git a/benchmark/dbally_benchmark/iql_benchmark.py b/benchmark/dbally_benchmark/iql_benchmark.py index f99411e5..1fd8bd50 100644 --- a/benchmark/dbally_benchmark/iql_benchmark.py +++ b/benchmark/dbally_benchmark/iql_benchmark.py @@ -22,7 +22,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.iql_generator.iql_generator import IQLGenerator from dbally.iql_generator.iql_prompt_template import default_iql_template -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient from dbally.utils.errors import UnsupportedQueryError from dbally.views.structured import BaseStructuredView diff --git a/benchmark/dbally_benchmark/text2sql_benchmark.py b/benchmark/dbally_benchmark/text2sql_benchmark.py index 99eebe6a..459fb132 100644 --- a/benchmark/dbally_benchmark/text2sql_benchmark.py +++ b/benchmark/dbally_benchmark/text2sql_benchmark.py @@ -21,8 +21,8 @@ from sqlalchemy import create_engine from dbally.audit.event_tracker import EventTracker -from dbally.llm_client.base import LLMClient -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.base import LLMClient +from dbally.llms.litellm import OpenAIClient def _load_db_schema(db_name: str, encoding: Optional[str] = None) -> str: diff --git a/docs/how-to/custom_views_code.py b/docs/how-to/custom_views_code.py index 5ea4fb8e..104d5281 100644 --- a/docs/how-to/custom_views_code.py +++ b/docs/how-to/custom_views_code.py @@ -10,7 +10,7 @@ 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.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient @dataclass class Candidate: diff --git a/docs/how-to/pandas_views_code.py b/docs/how-to/pandas_views_code.py index fe17f232..0118c7ea 100644 --- a/docs/how-to/pandas_views_code.py +++ b/docs/how-to/pandas_views_code.py @@ -8,7 +8,7 @@ from dbally import decorators, DataFrameBaseView from dbally.audit.event_handlers.cli_event_handler import CLIEventHandler -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient class CandidateView(DataFrameBaseView): diff --git a/docs/quickstart/quickstart2_code.py b/docs/quickstart/quickstart2_code.py index 9fa8cd4a..5a7f8ec7 100644 --- a/docs/quickstart/quickstart2_code.py +++ b/docs/quickstart/quickstart2_code.py @@ -12,7 +12,7 @@ from dbally.audit.event_handlers.cli_event_handler import CLIEventHandler from dbally.similarity import SimpleSqlAlchemyFetcher, FaissStore, SimilarityIndex from dbally.embedding_client.openai import OpenAiEmbeddingClient -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient engine = create_engine('sqlite:///candidates.db') diff --git a/docs/quickstart/quickstart3_code.py b/docs/quickstart/quickstart3_code.py index 1d474e8b..ce44ce71 100644 --- a/docs/quickstart/quickstart3_code.py +++ b/docs/quickstart/quickstart3_code.py @@ -12,7 +12,7 @@ from dbally import decorators, SqlAlchemyBaseView, DataFrameBaseView, ExecutionResult from dbally.similarity import SimpleSqlAlchemyFetcher, FaissStore, SimilarityIndex from dbally.embedding_client.openai import OpenAiEmbeddingClient -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient engine = create_engine('sqlite:///candidates.db') diff --git a/docs/quickstart/quickstart_code.py b/docs/quickstart/quickstart_code.py index 77d9d87e..1f8e3efa 100644 --- a/docs/quickstart/quickstart_code.py +++ b/docs/quickstart/quickstart_code.py @@ -9,7 +9,7 @@ from dbally import decorators, SqlAlchemyBaseView from dbally.audit.event_handlers.cli_event_handler import CLIEventHandler -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient engine = create_engine('sqlite:///candidates.db') diff --git a/examples/recruiting.py b/examples/recruiting.py index b5290344..ffccbcdd 100644 --- a/examples/recruiting.py +++ b/examples/recruiting.py @@ -8,7 +8,7 @@ import dbally from dbally.audit.event_handlers.cli_event_handler import CLIEventHandler from dbally.audit.event_tracker import EventTracker -from dbally.llm_client.openai_client import OpenAIClient +from dbally.llms.litellm import OpenAIClient from dbally.prompts import PromptTemplate TEXT2SQL_PROMPT_TEMPLATE = PromptTemplate( diff --git a/src/dbally/_main.py b/src/dbally/_main.py index 316bec01..456d1d66 100644 --- a/src/dbally/_main.py +++ b/src/dbally/_main.py @@ -2,7 +2,7 @@ from .audit.event_handlers.base import EventHandler from .collection import Collection -from .llm_client.base import LLMClient +from .llms.base import LLMClient from .nl_responder.nl_responder import NLResponder from .view_selection.base import ViewSelector from .view_selection.llm_view_selector import LLMViewSelector diff --git a/src/dbally/collection.py b/src/dbally/collection.py index ee2271ba..78939086 100644 --- a/src/dbally/collection.py +++ b/src/dbally/collection.py @@ -8,7 +8,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.data_models.audit import RequestEnd, RequestStart from dbally.data_models.execution_result import ExecutionResult -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.nl_responder.nl_responder import NLResponder from dbally.similarity.index import AbstractSimilarityIndex from dbally.utils.errors import NoViewFoundError diff --git a/src/dbally/iql_generator/iql_generator.py b/src/dbally/iql_generator/iql_generator.py index 5bd310a0..ebdafa03 100644 --- a/src/dbally/iql_generator/iql_generator.py +++ b/src/dbally/iql_generator/iql_generator.py @@ -3,7 +3,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.iql_generator.iql_prompt_template import IQLPromptTemplate, default_iql_template -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.views.exposed_functions import ExposedFunction diff --git a/src/dbally/llm_client/__init__.py b/src/dbally/llms/__init__.py similarity index 100% rename from src/dbally/llm_client/__init__.py rename to src/dbally/llms/__init__.py diff --git a/src/dbally/llm_client/base.py b/src/dbally/llms/base.py similarity index 100% rename from src/dbally/llm_client/base.py rename to src/dbally/llms/base.py diff --git a/src/dbally/llm_client/openai_client.py b/src/dbally/llms/litellm.py similarity index 98% rename from src/dbally/llm_client/openai_client.py rename to src/dbally/llms/litellm.py index 4458698b..c64b4e44 100644 --- a/src/dbally/llm_client/openai_client.py +++ b/src/dbally/llms/litellm.py @@ -5,7 +5,7 @@ from openai import NotGiven as OpenAINotGiven from dbally.data_models.audit import LLMEvent -from dbally.llm_client.base import LLMClient +from dbally.llms.base import LLMClient from dbally.prompts import ChatFormat from .._types import NOT_GIVEN, NotGiven diff --git a/src/dbally/nl_responder/nl_responder.py b/src/dbally/nl_responder/nl_responder.py index a0d093b4..bfeb2e26 100644 --- a/src/dbally/nl_responder/nl_responder.py +++ b/src/dbally/nl_responder/nl_responder.py @@ -5,7 +5,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.data_models.execution_result import ViewExecutionResult -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.nl_responder.nl_responder_prompt_template import NLResponderPromptTemplate, default_nl_responder_template from dbally.nl_responder.query_explainer_prompt_template import ( QueryExplainerPromptTemplate, diff --git a/src/dbally/view_selection/base.py b/src/dbally/view_selection/base.py index 94495e39..f06bacc8 100644 --- a/src/dbally/view_selection/base.py +++ b/src/dbally/view_selection/base.py @@ -2,7 +2,7 @@ from typing import Dict, Optional from dbally.audit.event_tracker import EventTracker -from dbally.llm_client.base import LLMOptions +from dbally.llms.base import LLMOptions class ViewSelector(abc.ABC): diff --git a/src/dbally/view_selection/llm_view_selector.py b/src/dbally/view_selection/llm_view_selector.py index 477f763f..1254fa9d 100644 --- a/src/dbally/view_selection/llm_view_selector.py +++ b/src/dbally/view_selection/llm_view_selector.py @@ -3,7 +3,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.iql_generator.iql_prompt_template import IQLPromptTemplate -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.view_selection.base import ViewSelector from dbally.view_selection.view_selector_prompt_template import default_view_selector_template diff --git a/src/dbally/view_selection/random_view_selector.py b/src/dbally/view_selection/random_view_selector.py index 8b2b1173..7b93d6f7 100644 --- a/src/dbally/view_selection/random_view_selector.py +++ b/src/dbally/view_selection/random_view_selector.py @@ -2,7 +2,7 @@ from typing import Dict, Optional from dbally.audit.event_tracker import EventTracker -from dbally.llm_client.base import LLMOptions +from dbally.llms.base import LLMOptions from dbally.view_selection.base import ViewSelector diff --git a/src/dbally/views/base.py b/src/dbally/views/base.py index a0846153..dbab6288 100644 --- a/src/dbally/views/base.py +++ b/src/dbally/views/base.py @@ -3,7 +3,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.data_models.execution_result import ViewExecutionResult -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions class BaseView(metaclass=abc.ABCMeta): diff --git a/src/dbally/views/freeform/text2sql/_autodiscovery.py b/src/dbally/views/freeform/text2sql/_autodiscovery.py index b579876d..63ad6ded 100644 --- a/src/dbally/views/freeform/text2sql/_autodiscovery.py +++ b/src/dbally/views/freeform/text2sql/_autodiscovery.py @@ -4,7 +4,7 @@ from sqlalchemy.sql.ddl import CreateTable from typing_extensions import Self -from dbally.llm_client.base import LLMClient +from dbally.llms.base import LLMClient from dbally.prompts import PromptTemplate from ._config import Text2SQLConfig, Text2SQLTableConfig diff --git a/src/dbally/views/freeform/text2sql/_view.py b/src/dbally/views/freeform/text2sql/_view.py index 3efe0e66..3cd8301e 100644 --- a/src/dbally/views/freeform/text2sql/_view.py +++ b/src/dbally/views/freeform/text2sql/_view.py @@ -5,7 +5,7 @@ from dbally.audit.event_tracker import EventTracker from dbally.data_models.execution_result import ViewExecutionResult -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.prompts import PromptTemplate from dbally.views.base import BaseView diff --git a/src/dbally/views/structured.py b/src/dbally/views/structured.py index 4cead890..db6c36d7 100644 --- a/src/dbally/views/structured.py +++ b/src/dbally/views/structured.py @@ -5,7 +5,7 @@ from dbally.data_models.execution_result import ViewExecutionResult from dbally.iql import IQLError, IQLQuery from dbally.iql_generator.iql_generator import IQLGenerator -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.views.exposed_functions import ExposedFunction from .base import BaseView diff --git a/tests/unit/mocks.py b/tests/unit/mocks.py index 2e8ceafa..e8583f52 100644 --- a/tests/unit/mocks.py +++ b/tests/unit/mocks.py @@ -12,7 +12,7 @@ from dbally.iql import IQLQuery from dbally.iql_generator.iql_generator import IQLGenerator from dbally.iql_generator.iql_prompt_template import IQLPromptTemplate, default_iql_template -from dbally.llm_client.base import LLMClient, LLMOptions +from dbally.llms.base import LLMClient, LLMOptions from dbally.similarity.index import AbstractSimilarityIndex from dbally.view_selection.base import ViewSelector from dbally.views.structured import BaseStructuredView, ExposedFunction, ViewExecutionResult diff --git a/tests/unit/test_view_selector.py b/tests/unit/test_view_selector.py index 09a7f1d3..62a5a78a 100644 --- a/tests/unit/test_view_selector.py +++ b/tests/unit/test_view_selector.py @@ -7,7 +7,7 @@ import dbally from dbally.audit.event_tracker import EventTracker -from dbally.llm_client.base import LLMClient +from dbally.llms.base import LLMClient from dbally.view_selection.llm_view_selector import LLMViewSelector from .mocks import MockLLMClient