diff --git a/great_ai/__init__.py b/great_ai/__init__.py index 40c27e9..c56762d 100644 --- a/great_ai/__init__.py +++ b/great_ai/__init__.py @@ -3,7 +3,7 @@ __version__ = "0.0.12" from .context import configure -from .deploy import GreatAI, RouteConfig +from .deploy import GreatAI from .exceptions import ( ArgumentValidationError, MissingArgumentError, @@ -24,4 +24,4 @@ from .remote import ( call_remote_great_ai_async, ) from .tracing import add_ground_truth, delete_ground_truth, query_ground_truth -from .views import Trace +from .views import RouteConfig, Trace diff --git a/great_ai/context.py b/great_ai/context.py index 19010e1..38537d5 100644 --- a/great_ai/context.py +++ b/great_ai/context.py @@ -2,7 +2,7 @@ import os import random from logging import DEBUG, Logger from pathlib import Path -from typing import Any, Dict, Optional, Type, cast +from typing import Any, Dict, Optional, Type, Union, cast from pydantic import BaseModel @@ -17,9 +17,11 @@ from .constants import ( from .large_file import LargeFileBase, LargeFileLocal from .persistence import ParallelTinyDbDriver, TracingDatabaseDriver from .utilities import get_logger +from .views import RouteConfig class Context(BaseModel): + version: Union[int, str] tracing_database: TracingDatabaseDriver large_file_implementation: Type[LargeFileBase] is_production: bool @@ -27,6 +29,7 @@ class Context(BaseModel): should_log_exception_stack: bool prediction_cache_size: int dashboard_table_size: int + route_config: RouteConfig class Config: arbitrary_types_allowed = True @@ -54,6 +57,7 @@ def get_context() -> Context: def configure( *, + version: Union[int, str] = "0.0.1", log_level: int = DEBUG, seed: int = 42, tracing_database_factory: Optional[Type[TracingDatabaseDriver]] = None, @@ -62,6 +66,7 @@ def configure( prediction_cache_size: int = 512, disable_se4ml_banner: bool = False, dashboard_table_size: int = 20, + route_config: RouteConfig = RouteConfig(), ) -> None: global _context logger = get_logger("great_ai", level=log_level) @@ -89,6 +94,7 @@ def configure( logger.warning(message) _context = Context( + version=version, tracing_database=tracing_database, large_file_implementation=_initialize_large_file( large_file_implementation, logger=logger @@ -100,6 +106,7 @@ def configure( else should_log_exception_stack, prediction_cache_size=prediction_cache_size, dashboard_table_size=dashboard_table_size, + route_config=route_config, ) logger.info("Settings: configured ✅") diff --git a/great_ai/deploy/__init__.py b/great_ai/deploy/__init__.py index 63e52be..49f9bc1 100644 --- a/great_ai/deploy/__init__.py +++ b/great_ai/deploy/__init__.py @@ -1,2 +1 @@ from .great_ai import GreatAI -from .routes import RouteConfig diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index 334d702..a3059f7 100644 --- a/great_ai/deploy/great_ai.py +++ b/great_ai/deploy/great_ai.py @@ -7,7 +7,6 @@ from typing import ( Generic, List, Optional, - Protocol, Sequence, TypeVar, Union, @@ -32,7 +31,6 @@ from .routes.bootstrap_feedback_endpoints import bootstrap_feedback_endpoints from .routes.bootstrap_meta_endpoints import bootstrap_meta_endpoints from .routes.bootstrap_prediction_endpoint import bootstrap_prediction_endpoint from .routes.bootstrap_trace_endpoints import bootstrap_trace_endpoints -from .routes.route_config import RouteConfig T = TypeVar("T", bound=Union[Trace, Awaitable[Trace]]) V = TypeVar("V") @@ -42,27 +40,9 @@ class GreatAI(Generic[T, V]): __name__: str __doc__: str - class FactoryProtocol(Protocol): - @overload - def __call__( # type: ignore - self, - func: Callable[..., Awaitable[V]], - ) -> "GreatAI[Awaitable[Trace[V]], V]": - - ... - - @overload - def __call__( - self, - func: Callable[..., V], - ) -> "GreatAI[Trace[V], V]": - ... - def __init__( self, func: Callable[..., Union[V, Awaitable[V]]], - version: Union[str, int], - route_config: RouteConfig, ): func = automatically_decorate_parameters(func) get_function_metadata_store(func).is_finalised = True @@ -73,7 +53,7 @@ class GreatAI(Generic[T, V]): wraps(func)(self) self.__doc__ = f"GreatAI wrapper for interacting with the `{self.__name__}` function.\n\n{dedent(self.__doc__ or '')}" - self.version = str(version) + self.version = str(get_context().version) flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions) if flat_model_versions: self.version += f"+{flat_model_versions}" @@ -87,7 +67,7 @@ class GreatAI(Generic[T, V]): redoc_url=None, ) - self._bootstrap_rest_api(route_config) + self._bootstrap_rest_api() @overload @staticmethod @@ -105,36 +85,13 @@ class GreatAI(Generic[T, V]): ) -> "GreatAI[Trace[V], V]": ... - @overload @staticmethod def create( - func: None = ..., - *, - version: Union[str, int] = ..., - route_config: RouteConfig = ..., - ) -> FactoryProtocol: - ... - - @staticmethod - def create( - func: Optional[Callable] = None, - *, - version: Union[str, int] = "0.0.1", - route_config: RouteConfig = RouteConfig(), - ) -> Union[ - FactoryProtocol, "GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]" - ]: - def factory(_func): # type: ignore - return GreatAI[Trace[V], V]( - _func, - version=version, - route_config=route_config, - ) - - if func is None: - return cast(GreatAI.FactoryProtocol, factory) - else: - return factory(func) + func: Union[Callable[..., Awaitable[V]], Callable[..., V]], + ) -> Union["GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]"]: + return GreatAI[Trace[V], V]( + func, + ) def __call__(self, *args: Any, **kwargs: Any) -> T: return self._wrapped_func(*args, **kwargs) @@ -192,7 +149,11 @@ class GreatAI(Generic[T, V]): ), ) - def _bootstrap_rest_api(self, route_config: RouteConfig) -> None: + def _bootstrap_rest_api( + self, + ) -> None: + route_config = get_context().route_config + if route_config.prediction_endpoint_enabled: bootstrap_prediction_endpoint(self.app, self._wrapped_func) diff --git a/great_ai/deploy/routes/__init__.py b/great_ai/deploy/routes/__init__.py index 5ac4e86..693155e 100644 --- a/great_ai/deploy/routes/__init__.py +++ b/great_ai/deploy/routes/__init__.py @@ -4,4 +4,3 @@ from .bootstrap_feedback_endpoints import bootstrap_feedback_endpoints from .bootstrap_meta_endpoints import bootstrap_meta_endpoints from .bootstrap_prediction_endpoint import bootstrap_prediction_endpoint from .bootstrap_trace_endpoints import bootstrap_trace_endpoints -from .route_config import RouteConfig diff --git a/great_ai/views/__init__.py b/great_ai/views/__init__.py index 65ef6f8..cf82c4f 100644 --- a/great_ai/views/__init__.py +++ b/great_ai/views/__init__.py @@ -7,5 +7,6 @@ from .health_check_response import HealthCheckResponse from .model import Model from .operators import operators from .query import Query +from .route_config import RouteConfig from .sort_by import SortBy from .trace import Trace diff --git a/great_ai/deploy/routes/route_config.py b/great_ai/views/route_config.py similarity index 100% rename from great_ai/deploy/routes/route_config.py rename to great_ai/views/route_config.py diff --git a/tests/test_async_starters.py b/tests/test_async_starters.py index 9a3c5fc..16eb06d 100644 --- a/tests/test_async_starters.py +++ b/tests/test_async_starters.py @@ -20,20 +20,6 @@ async def test_create_trivial_cases() -> None: assert (await hello_world_2("andras")).output == "Hello andras!" - @GreatAI.create() - async def hello_world_3(name: str) -> str: - await sleep(0.5) - return f"Hello {name}!" - - assert (await hello_world_3("andras")).output == "Hello andras!" - - @GreatAI.create() - async def hello_world_4(name): # type: ignore - await sleep(0.5) - return f"Hello {name}!" - - assert (await hello_world_4("andras")).output == "Hello andras!" - @pytest.mark.asyncio async def test_with_parameter() -> None: diff --git a/tests/test_basic_starters.py b/tests/test_basic_starters.py index 47f6617..fc13da7 100644 --- a/tests/test_basic_starters.py +++ b/tests/test_basic_starters.py @@ -22,18 +22,6 @@ def test_create_trivial_cases() -> None: assert hello_world_2("andras").output == "Hello andras!" - @GreatAI.create() - def hello_world_3(name: str) -> str: - return f"Hello {name}!" - - assert hello_world_3("andras").output == "Hello andras!" - - @GreatAI.create() - def hello_world_4(name): # type: ignore - return f"Hello {name}!" - - assert hello_world_4("andras").output == "Hello andras!" - def test_create_with_other_decorator() -> None: @GreatAI.create @@ -44,7 +32,7 @@ def test_create_with_other_decorator() -> None: assert hello_world_1("andras").output == "Hello andras!" @lru_cache - @GreatAI.create() + @GreatAI.create def hello_world_2(name: str) -> str: return f"Hello {name}!"