Remove options from GreatAI.create

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

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 ✅")

great_ai/deploy/__init__.py

@@ -1,2 +1 @@
 from .great_ai import GreatAI
-from .routes import RouteConfig

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)
 

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

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

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:

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}!"