From 785c1db237cc8bc770ab59a76b616b0cfb7be0ee Mon Sep 17 00:00:00 2001
From: Aaron <29749331+aarnphm@users.noreply.github.com>
Date: Fri, 11 Aug 2023 06:14:59 -0400
Subject: [PATCH] fix(client): include openllm.client into main module [skip
 ci]

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
---
 .github/workflows/build.yml                   |  1 -
 DEVELOPMENT.md                                |  1 -
 pyproject.toml                                |  7 ++--
 src/openllm/cli/entrypoint.py                 | 14 +++-----
 src/openllm/client.py                         | 34 -------------------
 .../client}/__init__.py                       |  8 ++---
 .../client}/runtimes/__init__.py              |  0
 .../client}/runtimes/base.py                  |  0
 .../client}/runtimes/grpc.py                  |  1 -
 .../client}/runtimes/http.py                  |  1 -
 tests/models/conftest.py                      |  2 +-
 11 files changed, 13 insertions(+), 56 deletions(-)
 delete mode 100644 src/openllm/client.py
 rename src/{openllm_client => openllm/client}/__init__.py (81%)
 rename src/{openllm_client => openllm/client}/runtimes/__init__.py (100%)
 rename src/{openllm_client => openllm/client}/runtimes/base.py (100%)
 rename src/{openllm_client => openllm/client}/runtimes/grpc.py (99%)
 rename src/{openllm_client => openllm/client}/runtimes/http.py (99%)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 338bde72..172a4ec9 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -13,7 +13,6 @@ on:
       - ".github/workflows/build.yaml"
       - "src/openllm/bundle/oci/Dockerfile"
       - "src/openllm/**"
-      - "src/openllm_client/**"
 env:
   LINES: 120
   COLUMNS: 120
diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md
index 0a4376fe..48d853e5 100644
--- a/DEVELOPMENT.md
+++ b/DEVELOPMENT.md
@@ -74,7 +74,6 @@ openllm/
 ├── examples                 # Usage demonstration scripts
 ├── src
 │   ├── openllm              # openllm core
-│   ├── openllm_client       # openllm python client
 │   └── openllm-node         # openllm nodejs library
 ├── tests                    # Automated Tests
 ├── tools                    # Utilities Script
diff --git a/pyproject.toml b/pyproject.toml
index 27dabc69..7ddb6a0b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -284,7 +284,7 @@ avoid-escape = false
 "src/openllm/models/**" = ["I001", "E", "D", "F"]
 "src/openllm/utils/__init__.py" = ["I001"]
 "src/openllm/utils/import_utils.py" = ["PLW0603"]
-"src/openllm_client/runtimes/*" = ["D107"]
+"src/openllm/client/runtimes/*" = ["D107"]
 "tests/**/*" = [
     "S101",
     "TID252",
@@ -395,7 +395,7 @@ exclude = [
     "examples",
     "tests",
 ]
-include = ["src/openllm", "src/openllm_client"]
+include = ["src/openllm"]
 pythonVersion = "3.12"
 reportMissingImports = "warning"
 reportMissingTypeStubs = false
@@ -411,7 +411,7 @@ typeCheckingMode = "strict"
 # TODO: Enable model for strict type checking
 exclude = ["src/openllm/playground/", "src/openllm/utils/dummy_*.py", "src/openllm/models"]
 local_partial_types = true
-modules = ["openllm", "openllm_client"]
+modules = ["openllm"]
 mypy_path = "typings"
 pretty = true
 python_version = "3.11"
@@ -500,7 +500,6 @@ include = [
     "src/openllm/_quantisation.py",
     "src/openllm/_generation.py",
     "src/openllm/_strategies.py",
-    "src/openllm/client.py",
     "src/openllm/exceptions.py",
     "src/openllm/testing.py",
 ]
diff --git a/src/openllm/cli/entrypoint.py b/src/openllm/cli/entrypoint.py
index c82b1dce..b9c0b0a3 100644
--- a/src/openllm/cli/entrypoint.py
+++ b/src/openllm/cli/entrypoint.py
@@ -69,11 +69,7 @@ from ._factory import (
   start_command_factory,
   workers_per_resource_option,
 )
-from .. import (
-  bundle,
-  client as openllm_client,
-  serialisation,
-)
+from .. import bundle, serialisation
 from ..exceptions import OpenLLMException
 from ..models.auto import (
   CONFIG_MAPPING,
@@ -109,10 +105,10 @@ from ..utils import (
 
 if t.TYPE_CHECKING:
   import torch
-  from openllm_client.runtimes.base import BaseClient
 
   from bentoml._internal.bento import BentoStore
   from bentoml._internal.container import DefaultBuilder
+  from openllm.client import BaseClient
 
   from .._schema import EmbeddingsOutput
   from .._types import DictStrAny, LiteralRuntime, P
@@ -713,7 +709,7 @@ def instruct_command(endpoint: str, timeout: int, agent: t.LiteralString, output
         --text "¡Este es un API muy agradable!"
   ```
   """
-  client = openllm_client.HTTPClient(endpoint, timeout=timeout)
+  client = openllm.client.HTTPClient(endpoint, timeout=timeout)
 
   try:
     client.call("metadata")
@@ -745,7 +741,7 @@ def embed_command(ctx: click.Context, text: tuple[str, ...], endpoint: str, time
   $ openllm embed --endpoint http://12.323.2.1:3000 "What is the meaning of life?" "How many stars are there in the sky?"
   ```
   """
-  client = t.cast("BaseClient[t.Any]", openllm_client.HTTPClient(endpoint, timeout=timeout) if server_type == "http" else openllm_client.GrpcClient(endpoint, timeout=timeout))
+  client = t.cast("BaseClient[t.Any]", openllm.client.HTTPClient(endpoint, timeout=timeout) if server_type == "http" else openllm.client.GrpcClient(endpoint, timeout=timeout))
   try:
     gen_embed = client.embed(text)
   except ValueError:
@@ -778,7 +774,7 @@ def query_command(ctx: click.Context, /, prompt: str, endpoint: str, timeout: in
   """
   _memoized = {k: orjson.loads(v[0]) for k, v in _memoized.items() if v}
   if server_type == "grpc": endpoint = re.sub(r"http://", "", endpoint)
-  client = t.cast("BaseClient[t.Any]", openllm_client.HTTPClient(endpoint, timeout=timeout) if server_type == "http" else openllm_client.GrpcClient(endpoint, timeout=timeout))
+  client = t.cast("BaseClient[t.Any]", openllm.client.HTTPClient(endpoint, timeout=timeout) if server_type == "http" else openllm.client.GrpcClient(endpoint, timeout=timeout))
   input_fg, generated_fg = "magenta", "cyan"
   if output != "porcelain":
     termui.echo("==Input==\n", fg="white")
diff --git a/src/openllm/client.py b/src/openllm/client.py
deleted file mode 100644
index 9b6fa40c..00000000
--- a/src/openllm/client.py
+++ /dev/null
@@ -1,34 +0,0 @@
-"""OpenLLM client.
-
-To start interact with the server, you can do the following:
-
->>> import openllm
->>> client = openllm.client.HTTPClient("http://localhost:3000")
->>> client.query("What is the meaning of life?")
-"""
-from __future__ import annotations
-import importlib
-import typing as t
-
-_import_structure: dict[str, list[str]] = {"runtimes.grpc": ["AsyncGrpcClient", "GrpcClient"], "runtimes.http": ["AsyncHTTPClient", "HTTPClient"], "runtimes.base": ["BaseClient", "BaseAsyncClient"]}
-
-if t.TYPE_CHECKING:
-  from openllm_client.runtimes import (
-    AsyncGrpcClient as AsyncGrpcClient,
-    AsyncHTTPClient as AsyncHTTPClient,
-    BaseAsyncClient as BaseAsyncClient,
-    BaseClient as BaseClient,
-    GrpcClient as GrpcClient,
-    HTTPClient as HTTPClient,
-  )
-
-_module = "openllm_client"
-
-def __dir__() -> list[str]: return sorted(__all__)
-def __getattr__(name: str) -> t.Any:
-  if name in _import_structure: return importlib.import_module(f".{name}", _module)
-  try: module = next(module for module, attrs in _import_structure.items() if name in attrs)
-  except StopIteration: raise AttributeError(f"module {_module} has no attribute {name}") from None
-  return getattr(importlib.import_module(f".{module}", _module), name)
-# NOTE: Make sure to always keep this line at the bottom of the file. The update will be managed via tools/update-init-import.py
-__all__=["AsyncGrpcClient","AsyncHTTPClient","BaseAsyncClient","BaseClient","GrpcClient","HTTPClient"]
diff --git a/src/openllm_client/__init__.py b/src/openllm/client/__init__.py
similarity index 81%
rename from src/openllm_client/__init__.py
rename to src/openllm/client/__init__.py
index 24e6ae7d..aeb4448b 100644
--- a/src/openllm_client/__init__.py
+++ b/src/openllm/client/__init__.py
@@ -6,11 +6,11 @@ OpenLLM server. It is used to send requests to the server, and receive responses
 """
 from __future__ import annotations
 
-from .runtimes.grpc import (
+from .runtimes import (
     AsyncGrpcClient as AsyncGrpcClient,
-    GrpcClient as GrpcClient,
-)
-from .runtimes.http import (
     AsyncHTTPClient as AsyncHTTPClient,
+    BaseAsyncClient as BaseAsyncClient,
+    BaseClient as BaseClient,
+    GrpcClient as GrpcClient,
     HTTPClient as HTTPClient,
 )
diff --git a/src/openllm_client/runtimes/__init__.py b/src/openllm/client/runtimes/__init__.py
similarity index 100%
rename from src/openllm_client/runtimes/__init__.py
rename to src/openllm/client/runtimes/__init__.py
diff --git a/src/openllm_client/runtimes/base.py b/src/openllm/client/runtimes/base.py
similarity index 100%
rename from src/openllm_client/runtimes/base.py
rename to src/openllm/client/runtimes/base.py
diff --git a/src/openllm_client/runtimes/grpc.py b/src/openllm/client/runtimes/grpc.py
similarity index 99%
rename from src/openllm_client/runtimes/grpc.py
rename to src/openllm/client/runtimes/grpc.py
index 69fb754c..44069cc4 100644
--- a/src/openllm_client/runtimes/grpc.py
+++ b/src/openllm/client/runtimes/grpc.py
@@ -1,4 +1,3 @@
-
 from __future__ import annotations
 import asyncio
 import logging
diff --git a/src/openllm_client/runtimes/http.py b/src/openllm/client/runtimes/http.py
similarity index 99%
rename from src/openllm_client/runtimes/http.py
rename to src/openllm/client/runtimes/http.py
index dd43e287..6125ccff 100644
--- a/src/openllm_client/runtimes/http.py
+++ b/src/openllm/client/runtimes/http.py
@@ -1,4 +1,3 @@
-
 from __future__ import annotations
 import logging
 import typing as t
diff --git a/tests/models/conftest.py b/tests/models/conftest.py
index 6214e16a..5f9c6dd8 100644
--- a/tests/models/conftest.py
+++ b/tests/models/conftest.py
@@ -38,12 +38,12 @@ logger = logging.getLogger(__name__)
 if t.TYPE_CHECKING:
   import subprocess
 
-  from openllm_client.runtimes.base import BaseAsyncClient
   from syrupy.assertion import SnapshotAssertion
   from syrupy.types import PropertyFilter, PropertyMatcher, SerializableData, SerializedData
 
   from openllm._configuration import GenerationConfig
   from openllm._types import DictStrAny, ListAny
+  from openllm.client import BaseAsyncClient
 
 else:
   DictStrAny = dict