feat: associate generation calls with devices

This commit is contained in:
QiuSW
2026-07-21 09:40:40 +08:00
parent 69f4517958
commit 11c0d63ac8
16 changed files with 398 additions and 19 deletions
+10 -2
View File
@@ -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,