feat(styles): record the style each generation ran under

- WaitingPrompt.style_id records the style a request runs under, the drawn style for a collection, and record_image_statistic and record_text_statistic copy it to image_gen_stats.style_id and text_gen_stats.style_id, so style use can be grouped per style
- the columns are nullable UUIDs with no foreign key, so a deleted style leaves its statistics in place; sql_statements/5.1.13.txt adds them as catalogue-only changes
- integration tests check a queued text and image request records its style, the drawn style under a collection, and none when unstyled; unit tests check the text statistic copies it

docs/reference/style_contract.md

@@ -478,6 +478,7 @@ The span that validates a styled generate request carries `horde.style.id`, `hor
 | Worker context ceiling | `horde/database/functions.py` | `get_highest_text_worker_max_context_length` |
 | Text pricing | `horde/classes/kobold/waiting_prompt.py` | `TextWaitingPrompt.calculate_kudos` |
 | Style columns | `horde/classes/base/style.py` | `Style.parameter_policy`, `Style.template_fields` |
+| The style a generation ran under | `horde/classes/base/waiting_prompt.py`, `horde/classes/stable/genstats.py`, `horde/classes/kobold/genstats.py` | `WaitingPrompt.style_id`, `ImageGenerationStatistic.style_id`, `TextGenerationStatistic.style_id` |
 | The owner-only `shared_key` on a served style | `horde/apis/v2/styles.py`, `horde/classes/base/style.py` | `SingleStyleTemplateGet.get_existing_style`, `Style.get_details` |
 | Style application log line, counter and span attributes | `horde/apis/v2/base.py` | `GenerateTemplate.report_style_application` |
 | Style and collection write log lines | `horde/apis/v2/styles.py` | `describe_style`, `describe_collection`, `STYLE_PATCH_FIELDS` |
@@ -527,6 +528,7 @@ The span that validates a styled generate request carries `horde.style.id`, `hor
 | Context fit outcomes are counted | `tests/integration/test_text_style_application.py`, `TestContextFitTelemetry` |
 | Style writes are logged | `tests/integration/test_text_styles.py`, `TestTextStyleWriteLogs` |
 | Author credits and forgiven surcharges are counted | `tests/unit/test_kudos_ledger.py`, `TestStyleRewardOutcomeCounters` |
+| A queued request records its style, and its generation statistic copies it | `tests/integration/test_text_style_application.py`, `TestTextStyleAttribution`; `tests/integration/test_image_style_application.py`, `TestImageStyleAttribution`; `tests/unit/test_text_genstats.py`, `TestStatisticStyle` |
 | API refusals are counted by return code and route | `tests/integration/test_malformed_requests.py`, `test_a_refusal_is_counted_by_return_code_and_route` |
 | A write to a style is served by its single-style reads straight away | `tests/integration/test_text_styles.py`, `TestTextStyleReadCache`; `tests/integration/test_image_styles.py`, `TestImageStyleReadCache` |
 
@@ -560,7 +562,13 @@ The span that validates a styled generate request carries `horde.style.id`, `hor
 ## Schema
 
 `sql_statements/5.1.13.txt` adds `styles.parameter_policy` and `styles.template_fields`, both nullable
-`JSONB` with no default, along with column comments. Files at the `sql_statements/` level are run by
-hand with psql rather than by the application; run it in autocommit with `-v ON_ERROR_STOP=1`. Adding a
-nullable column with no default is a catalogue-only change in PostgreSQL, so the migration takes only a
-brief lock and can be applied before the code that reads the columns is deployed.
+`JSONB` with no default, along with column comments. It also adds a nullable `style_id` UUID to
+`waiting_prompts`, `image_gen_stats` and `text_gen_stats`. A queued request records the style it runs
+under, the style drawn from a collection when the request specified one, and the generation statistic
+copies it on submit, so style use can be grouped per style. The column has no foreign key, so a deleted
+style leaves its statistics in place, and no index.
+
+Files at the `sql_statements/` level are run by hand with psql rather than by the application. Run it in
+autocommit with `-v ON_ERROR_STOP=1`. Adding a nullable column with no default is a catalogue-only change
+in PostgreSQL, so the migration takes only a brief lock on each table and can be applied before the code
+that reads the columns is deployed.

horde/apis/v2/kobold.py

@@ -161,6 +161,7 @@ class TextAsyncGenerate(GenerateTemplate):
             safe_ip=True,
             client_agent=self.args["Client-Agent"],
             sharedkey_id=self.sharedkey.id if self.sharedkey else None,
+            style_id=self.existing_style.id if self.existing_style is not None else None,
             proxied_account=self.args["proxied_account"],
             webhook=self.args.webhook,
         )

horde/apis/v2/stable.py

@@ -333,6 +333,7 @@ class ImageAsyncGenerate(GenerateTemplate):
             shared=shared,
             client_agent=self.args["Client-Agent"],
             sharedkey_id=self.sharedkey.id if self.sharedkey else None,
+            style_id=self.existing_style.id if self.existing_style is not None else None,
             proxied_account=self.args["proxied_account"],
             disable_batching=self.args["disable_batching"],
             webhook=self.args.webhook,

horde/classes/base/waiting_prompt.py

@@ -199,6 +199,9 @@ class WaitingPrompt(db.Model):
         nullable=True,
     )
     sharedkey = db.relationship("UserSharedKey", back_populates="waiting_prompts")
+    # The style the request runs under, drawn from a collection when the request specified one. Copied to
+    # the generation statistics on submit. No foreign key, so deleting a style leaves running requests alone.
+    style_id = db.Column(uuid_column_type(), nullable=True)
     proxied_account = db.Column(db.String(255), nullable=True)
 
     tricked_workers: Mapped[list[WPTrickedWorkers]] = relationship(

horde/classes/kobold/genstats.py

@@ -5,12 +5,15 @@
 from datetime import datetime
 
 from sqlalchemy import Enum
+from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.exc import SQLAlchemyError
 
 from horde.enums import ImageGenState
-from horde.flask import db
+from horde.flask import SQLITE_MODE, db
 from horde.logger import logger
 
+uuid_column_type = lambda: UUID(as_uuid=True) if not SQLITE_MODE else db.String(36)  # noqa E731
+
 
 def record_text_statistic(procgen):
     """Store a finished text generation's statistic row, best-effort.
@@ -37,6 +40,7 @@ def record_text_statistic(procgen):
         prompt_length=len(procgen.wp.prompt),
         bridge_agent=procgen.worker.bridge_agent,
         client_agent=procgen.wp.client_agent,
+        style_id=procgen.wp.style_id,
         state=state,
     )
     try:
@@ -64,6 +68,9 @@ class TextGenerationStatistic(db.Model):
     client_agent = db.Column(db.Text, default="unknown:0:unknown", nullable=False, index=True)
     bridge_agent = db.Column(db.Text, default="unknown:0:unknown", nullable=False, index=True)
     state = db.Column(Enum(ImageGenState), default=ImageGenState.OK, nullable=False, index=True)
+    # The style the generation ran under, drawn from a collection when the request specified one. No foreign
+    # key: the statistic outlives a deleted style, and an id is what per-style analysis groups by.
+    style_id = db.Column(uuid_column_type(), nullable=True)
 
 
 class CompiledTextGensStatsTotals(db.Model):

horde/classes/stable/genstats.py

@@ -5,11 +5,14 @@
 from datetime import datetime
 
 from sqlalchemy import Enum
+from sqlalchemy.dialects.postgresql import UUID
 
 from horde.enums import ImageGenState
-from horde.flask import db
+from horde.flask import SQLITE_MODE, db
 from horde.logger import logger
 
+uuid_column_type = lambda: UUID(as_uuid=True) if not SQLITE_MODE else db.String(36)  # noqa E731
+
 
 class ImageGenerationStatisticPP(db.Model):
     __tablename__ = "image_gen_stats_post_processors"
@@ -80,6 +83,9 @@ class ImageGenerationStatistic(db.Model):
     state = db.Column(Enum(ImageGenState), default=ImageGenState.OK, nullable=False, index=True)
     client_agent = db.Column(db.Text, default="unknown:0:unknown", nullable=False, index=True)
     bridge_agent = db.Column(db.Text, default="unknown:0:unknown", nullable=False, index=True)
+    # The style the generation ran under, drawn from a collection when the request specified one. No foreign
+    # key: the statistic outlives a deleted style, and an id is what per-style analysis groups by.
+    style_id = db.Column(uuid_column_type(), nullable=True)
     post_processors = db.relationship(
         "ImageGenerationStatisticPP",
         back_populates="imgstat",
@@ -130,6 +136,7 @@ def record_image_statistic(procgen):
         nsfw=procgen.wp.nsfw,
         bridge_agent=procgen.worker.bridge_agent,
         client_agent=procgen.wp.client_agent,
+        style_id=procgen.wp.style_id,
         state=state,
     )
     db.session.add(statistic)

sql_statements/5.1.13.txt

@@ -20,3 +20,22 @@ COMMENT ON COLUMN styles.parameter_policy IS
   'Which generation params a request under this style may set, and the ceilings it may not exceed. Null means none.';
 COMMENT ON COLUMN styles.template_fields IS
   'The placeholders this style''s prompt accepts from the request, beyond {p}. Null means none.';
+
+-- The style a generation ran under, recorded on the waiting prompt and copied to
+-- the generation statistic on submit, so style use can be analysed per style. A
+-- collection request records the style drawn from it. Nullable with no default
+-- and no foreign key, so these are catalogue-only changes as well, and a deleted
+-- style leaves its statistics in place.
+ALTER TABLE waiting_prompts
+  ADD COLUMN IF NOT EXISTS style_id UUID NULL;
+ALTER TABLE image_gen_stats
+  ADD COLUMN IF NOT EXISTS style_id UUID NULL;
+ALTER TABLE text_gen_stats
+  ADD COLUMN IF NOT EXISTS style_id UUID NULL;
+
+COMMENT ON COLUMN waiting_prompts.style_id IS
+  'The style the request runs under, drawn from a collection when the request specified one. Null when unstyled.';
+COMMENT ON COLUMN image_gen_stats.style_id IS
+  'The style the generation ran under. Null when unstyled.';
+COMMENT ON COLUMN text_gen_stats.style_id IS
+  'The style the generation ran under. Null when unstyled.';

tests/integration/test_image_style_application.py

@@ -908,3 +908,47 @@ class TestImageStyleCollection:
 
             details = client.get(f"/api/v2/collections/{collection_id}", headers=request_headers).get_json()
             assert details["use_count"] == 1
+
+
+def waiting_prompt_style_id(app: Flask, request_id: str) -> str | None:
+    """Return the style a queued request's waiting prompt records.
+
+    Args:
+        app: The Flask app, for the database session.
+        request_id: The queued request's id.
+
+    Returns:
+        The recorded style id as a string, or None when the request runs under no style.
+    """
+    from horde.classes.base.waiting_prompt import WaitingPrompt
+    from horde.flask import db
+
+    with app.app_context():
+        style_id = db.session.query(WaitingPrompt.style_id).filter(WaitingPrompt.id == request_id).scalar()
+        return str(style_id) if style_id is not None else None
+
+
+class TestImageStyleAttribution:
+    """A queued request records the style it runs under, for the generation statistics."""
+
+    def test_a_request_under_a_style_records_it(self, app: Flask, client: FlaskClient, request_headers: dict[str, str]) -> None:
+        with created_style(client, request_headers, style_body("image attributed")) as style_id:
+            with queued_request(client, request_headers, style=style_id) as request_id:
+                assert waiting_prompt_style_id(app, request_id) == style_id
+
+    def test_a_request_under_a_collection_records_the_drawn_style(
+        self,
+        app: Flask,
+        client: FlaskClient,
+        request_headers: dict[str, str],
+    ) -> None:
+        with (
+            created_style(client, request_headers, style_body("image attributed member")) as style_id,
+            created_collection(client, request_headers, name="image attributed collection", style_ids=[style_id]) as collection_id,
+        ):
+            with queued_request(client, request_headers, style=collection_id) as request_id:
+                assert waiting_prompt_style_id(app, request_id) == style_id
+
+    def test_an_unstyled_request_records_no_style(self, app: Flask, client: FlaskClient, request_headers: dict[str, str]) -> None:
+        with queued_request(client, request_headers) as request_id:
+            assert waiting_prompt_style_id(app, request_id) is None

tests/integration/test_text_style_application.py

@@ -1598,3 +1598,47 @@ class TestContextFitTelemetry:
         assert context_fit_counter.attributes == [
             {"horde.context_fit.mode": "reject", "horde.context_fit.outcome": "fits"},
         ]
+
+
+def waiting_prompt_style_id(app: Flask, request_id: str) -> str | None:
+    """Return the style a queued request's waiting prompt records.
+
+    Args:
+        app: The Flask app, for the database session.
+        request_id: The queued request's id.
+
+    Returns:
+        The recorded style id as a string, or None when the request runs under no style.
+    """
+    from horde.classes.base.waiting_prompt import WaitingPrompt
+    from horde.flask import db
+
+    with app.app_context():
+        style_id = db.session.query(WaitingPrompt.style_id).filter(WaitingPrompt.id == request_id).scalar()
+        return str(style_id) if style_id is not None else None
+
+
+class TestTextStyleAttribution:
+    """A queued request records the style it runs under, for the generation statistics."""
+
+    def test_a_request_under_a_style_records_it(self, app: Flask, client: FlaskClient, request_headers: dict[str, str]) -> None:
+        with created_style(client, request_headers, style_body("text attributed")) as style_id:
+            with queued_request(client, request_headers, style=style_id) as request_id:
+                assert waiting_prompt_style_id(app, request_id) == style_id
+
+    def test_a_request_under_a_collection_records_the_drawn_style(
+        self,
+        app: Flask,
+        client: FlaskClient,
+        request_headers: dict[str, str],
+    ) -> None:
+        with (
+            created_style(client, request_headers, style_body("text attributed member")) as style_id,
+            created_collection(client, request_headers, name="text attributed collection", style_ids=[style_id]) as collection_id,
+        ):
+            with queued_request(client, request_headers, style=collection_id) as request_id:
+                assert waiting_prompt_style_id(app, request_id) == style_id
+
+    def test_an_unstyled_request_records_no_style(self, app: Flask, client: FlaskClient, request_headers: dict[str, str]) -> None:
+        with queued_request(client, request_headers) as request_id:
+            assert waiting_prompt_style_id(app, request_id) is None

tests/unit/test_text_genstats.py

@@ -10,6 +10,7 @@ failed insert is rolled back and logged without failing the worker's submit.
 
 from __future__ import annotations
 
+import uuid
 from typing import Any
 
 import pytest
@@ -27,13 +28,14 @@ pytestmark = pytest.mark.unit
 _MODEL = "elinas/chronos-70b-v2"
 
 
-def _make_text_wp(user: Any) -> TextWaitingPrompt:
+def _make_text_wp(user: Any, *, style_id: uuid.UUID | None = None) -> TextWaitingPrompt:
     wp = TextWaitingPrompt(
         [],
         [_MODEL],
         prompt="a unit-test text prompt",
         user_id=user.id,
         params={"n": 1, "max_length": 80, "max_context_length": 1024},
+        style_id=style_id,
     )
     db.session.commit()
     return wp
@@ -99,3 +101,30 @@ class TestStatisticWriteIsBestEffort:
 
         assert errors == []
         assert db.session.query(TextGenerationStatistic).count() == rows_before + 1
+
+
+class TestStatisticStyle:
+    """The statistic records the style its request ran under."""
+
+    @staticmethod
+    def _newest_statistic_style(procgen: TextProcessingGeneration) -> uuid.UUID | None:
+        record_text_statistic(procgen)
+        newest = db.session.query(TextGenerationStatistic).order_by(TextGenerationStatistic.id.desc()).first()
+        return newest.style_id
+
+    def test_a_styled_request_records_its_style(self, fake_redis, make_user) -> None:
+        user = make_user()
+        style_id = uuid.uuid4()
+        wp = _make_text_wp(user, style_id=style_id)
+        procgen = TextProcessingGeneration(wp_id=wp.id, worker_id=_make_text_worker(user).id, model=_MODEL)
+        db.session.commit()
+
+        assert self._newest_statistic_style(procgen) == style_id
+
+    def test_an_unstyled_request_records_no_style(self, fake_redis, make_user) -> None:
+        user = make_user()
+        wp = _make_text_wp(user)
+        procgen = TextProcessingGeneration(wp_id=wp.id, worker_id=_make_text_worker(user).id, model=_MODEL)
+        db.session.commit()
+
+        assert self._newest_statistic_style(procgen) is None