feat: associate generation calls with devices
This commit is contained in:
+10
-2
@@ -73,6 +73,7 @@ class GenerationInput:
|
||||
api_key: Any
|
||||
operation_type: str
|
||||
prompt: str
|
||||
client_device: Any | None = None
|
||||
alias: str | None = None
|
||||
resolution: str = "1K"
|
||||
parameters: Mapping[str, Any] = field(default_factory=dict)
|
||||
@@ -86,6 +87,7 @@ class GenerationInput:
|
||||
class PreparedGeneration:
|
||||
user: Any
|
||||
api_key: Any
|
||||
client_device: Any | None
|
||||
operation_type: str
|
||||
prompt: str
|
||||
alias: str
|
||||
@@ -139,11 +141,12 @@ IMAGE_URL_ALLOWED_SCHEMES = {"http", "https"}
|
||||
IMAGE_URL_CHUNK_SIZE = 64 * 1024
|
||||
|
||||
|
||||
def generate_title_response(*, user, api_key, request_data: Mapping[str, Any]) -> dict:
|
||||
def generate_title_response(*, user, api_key, client_device=None, request_data: Mapping[str, Any]) -> dict:
|
||||
result = run_synchronous_generation(
|
||||
GenerationInput(
|
||||
user=user,
|
||||
api_key=api_key,
|
||||
client_device=client_device,
|
||||
operation_type=CallRecord.OperationType.TITLE,
|
||||
prompt=request_data["prompt"],
|
||||
alias=request_data.get("model") or None,
|
||||
@@ -160,6 +163,7 @@ def generate_image_response(
|
||||
*,
|
||||
user,
|
||||
api_key,
|
||||
client_device=None,
|
||||
request_data: Mapping[str, Any],
|
||||
image_url_builder: ImageUrlBuilder | None = None,
|
||||
) -> dict:
|
||||
@@ -167,6 +171,7 @@ def generate_image_response(
|
||||
GenerationInput(
|
||||
user=user,
|
||||
api_key=api_key,
|
||||
client_device=client_device,
|
||||
operation_type=CallRecord.OperationType.IMAGE,
|
||||
prompt=request_data["prompt"],
|
||||
alias=request_data.get("model") or None,
|
||||
@@ -182,11 +187,12 @@ def generate_image_response(
|
||||
return result.as_response_data()
|
||||
|
||||
|
||||
def analyze_images_response(*, user, api_key, request_data: Mapping[str, Any]) -> dict:
|
||||
def analyze_images_response(*, user, api_key, client_device=None, request_data: Mapping[str, Any]) -> dict:
|
||||
result = run_synchronous_generation(
|
||||
GenerationInput(
|
||||
user=user,
|
||||
api_key=api_key,
|
||||
client_device=client_device,
|
||||
operation_type=CallRecord.OperationType.VISION,
|
||||
prompt=request_data["prompt"],
|
||||
alias=request_data.get("model") or None,
|
||||
@@ -257,6 +263,7 @@ def prepare_generation(generation_input: GenerationInput) -> PreparedGeneration:
|
||||
return PreparedGeneration(
|
||||
user=generation_input.user,
|
||||
api_key=generation_input.api_key,
|
||||
client_device=generation_input.client_device,
|
||||
operation_type=operation_type,
|
||||
prompt=prompt,
|
||||
alias=model_alias.alias,
|
||||
@@ -276,6 +283,7 @@ def precharge_generation(prepared: PreparedGeneration) -> PrechargedGeneration:
|
||||
charge = precharge_or_raise(
|
||||
user=prepared.user,
|
||||
api_key=prepared.api_key,
|
||||
client_device=prepared.client_device,
|
||||
operation_type=prepared.operation_type,
|
||||
alias=prepared.alias,
|
||||
model_used=prepared.resolved_model.model,
|
||||
|
||||
Reference in New Issue
Block a user