[owl] Add image output columns plus image‑gen usage/billing & image token quotas (#872)

1) Gen Table image output columns
New ImageGenConfig enables image outputs for Gen Tables.
Each output cell stores exactly one image URI (e.g. s3://...).
Supports image_gen (text → image) and image_edit (image(s) + text → image).
Image generation itself is non‑streaming; SSE emits one chunk when the URI is ready.
2) Model types & LM routing
Added ModelType.IMAGE_GEN and ModelCapability.IMAGE_OUT.
Capability rules:
image_gen requires IMAGE_OUT
image_edit requires IMAGE_OUT + IMAGE
LM routing supports:
Standard image endpoints (aimage_generation / aimage_edit)
Gemini preview allowlist routed via acompletion(..., modalities=["image","text"]) to obtain correct usage
3) Usage & billing telemetry (image‑gen)
New ClickHouse table image_gen_usage for image‑gen token/cost breakdown:
text_input_token, text_output_token, image_input_token, image_output_token
Cost aggregation updates:
CostTable union now includes image‑gen cost categories:
image_text_input / image_text_output / image_image_input / image_image_output
LLM usage tables remain pure; image‑gen does not write to llm_usage.
4) Image token metrics & meters
Image token usage recorded to VictoriaMetrics via image_token_usage counter.
Metrics query path added via type=image and metricId=image in meters APIs.
5) Quota enforcement + org schema
Org fields:
image_tokens_quota_mtok
image_tokens_usage_mtok
Enforced in cloud (ELLM only; BYOK exempt):
Image tokens consume image quota
Image‑gen text tokens still consume LLM quota
API/Behavior Notes
Image breakdown fields (text_tokens, image_tokens) appear only when set.

Author: noobHappylife

clients/python/src/jamaibase/client.py

@@ -3386,7 +3386,7 @@ class _MeterClientAsync(_ClientAsync):
 
     async def get_usage_metrics(
         self,
-        type: Literal["llm", "embedding", "reranking"],
+        type: Literal["llm", "embedding", "reranking", "image"],
         from_: datetime,
         window_size: str,
         org_ids: list[str] | None = None,

clients/python/src/jamaibase/types/__init__.py

@@ -9,6 +9,7 @@
     EgressUsageData,
     EmbedUsageData,
     FileStorageUsageData,
+    ImageGenUsageData,
     LlmUsageData,
     RerankUsageData,
     UsageData,
@@ -155,6 +156,7 @@
     DiscriminatedGenConfig,
     EmbedGenConfig,
     GenConfigUpdateRequest,
+    ImageGenConfig,
     KnowledgeTableSchemaCreate,
     LLMGenConfig,
     MultiRowAddRequest,

clients/python/src/jamaibase/types/billing.py

@@ -50,6 +50,36 @@ class LlmUsageData(_BaseUsageData):
     )
 
 
+class ImageGenUsageData(_BaseUsageData):
+    model: str = Field(
+        description="Model used.",
+    )
+    text_input_token: int = Field(
+        description="Number of text input tokens used.",
+    )
+    text_output_token: int = Field(
+        description="Number of text output tokens used.",
+    )
+    image_input_token: int = Field(
+        description="Number of image input tokens used.",
+    )
+    image_output_token: int = Field(
+        description="Number of image output tokens used.",
+    )
+    text_input_cost: float = Field(
+        description="Cost in USD per million text input tokens.",
+    )
+    text_output_cost: float = Field(
+        description="Cost in USD per million text output tokens.",
+    )
+    image_input_cost: float = Field(
+        description="Cost in USD per million image input tokens.",
+    )
+    image_output_cost: float = Field(
+        description="Cost in USD per million image output tokens.",
+    )
+
+
 class EmbedUsageData(_BaseUsageData):
     model: str = Field(
         description="Model used.",
@@ -94,6 +124,7 @@ class DBStorageUsageData(_BaseUsageData):
 
 class UsageData(BaseModel):
     llm_usage: list[LlmUsageData] = []
+    image_gen_usage: list[ImageGenUsageData] = []
     embed_usage: list[EmbedUsageData] = []
     rerank_usage: list[RerankUsageData] = []
     egress_usage: list[EgressUsageData] = []
@@ -105,6 +136,7 @@ def as_list_by_type(self) -> dict[str, list[list]]:
         """Returns a dictionary of lists, where each key is a usage type and the value is a list of lists."""
         return {
             "llm_usage": [usage.as_list() for usage in self.llm_usage],
+            "image_gen_usage": [usage.as_list() for usage in self.image_gen_usage],
             "embed_usage": [usage.as_list() for usage in self.embed_usage],
             "rerank_usage": [usage.as_list() for usage in self.rerank_usage],
             "egress_usage": [usage.as_list() for usage in self.egress_usage],
@@ -117,6 +149,7 @@ def total_usage_events(self) -> int:
         """Returns the total number of usage events across all types."""
         return (
             len(self.llm_usage)
+            + len(self.image_gen_usage)
             + len(self.embed_usage)
             + len(self.rerank_usage)
             + len(self.egress_usage)
@@ -128,6 +161,7 @@ def __add__(self, other: "UsageData") -> "UsageData":
         """Overload the + operator to combine two UsageData objects."""
         combined = UsageData()
         combined.llm_usage = self.llm_usage + other.llm_usage
+        combined.image_gen_usage = self.image_gen_usage + other.image_gen_usage
         combined.embed_usage = self.embed_usage + other.embed_usage
         combined.rerank_usage = self.rerank_usage + other.rerank_usage
         combined.egress_usage = self.egress_usage + other.egress_usage

clients/python/src/jamaibase/types/db.py

@@ -168,6 +168,9 @@ class Products(BaseModel):
     llm_tokens: Product = Field(
         description="LLM token quota to this plan or tier.",
     )
+    image_tokens: Product = Field(
+        description="Image token quota to this plan or tier.",
+    )
     embedding_tokens: Product = Field(
         description="Embedding token quota to this plan or tier.",
     )
@@ -188,6 +191,7 @@ class Products(BaseModel):
     def null(cls):
         return cls(
             llm_tokens=Product.null("ELLM tokens", "Million Tokens"),
+            image_tokens=Product.null("Image tokens", "Million Tokens"),
             embedding_tokens=Product.null("Embedding tokens", "Million Tokens"),
             reranker_searches=Product.null("Reranker searches", "Thousand Searches"),
             db_storage=Product.null("Database storage", "GiB"),
@@ -199,6 +203,7 @@ def null(cls):
     def unlimited(cls, unit_cost: float = 0.0):
         return cls(
             llm_tokens=Product.unlimited("ELLM tokens", "Million Tokens", unit_cost),
+            image_tokens=Product.unlimited("Image tokens", "Million Tokens", unit_cost),
             embedding_tokens=Product.unlimited("Embedding tokens", "Million Tokens", unit_cost),
             reranker_searches=Product.unlimited(
                 "Reranker searches", "Thousand Searches", unit_cost
@@ -213,6 +218,7 @@ def unlimited(cls, unit_cost: float = 0.0):
     credit=("credit",),
     credit_grant=("credit_grant",),
     llm_tokens=("llm_tokens_quota_mtok", "llm_tokens_usage_mtok"),
+    image_tokens=("image_tokens_quota_mtok", "image_tokens_usage_mtok"),
     embedding_tokens=(
         "embedding_tokens_quota_mtok",
         "embedding_tokens_usage_mtok",
@@ -228,6 +234,7 @@ class ProductType(StrEnum):
     CREDIT = "credit"
     CREDIT_GRANT = "credit_grant"
     LLM_TOKENS = "llm_tokens"
+    IMAGE_TOKENS = "image_tokens"
     EMBEDDING_TOKENS = "embedding_tokens"
     RERANKER_SEARCHES = "reranker_searches"
     DB_STORAGE = "db_storage"
@@ -308,6 +315,12 @@ def free(
                     tiers=[],
                     unit="Million Tokens",
                 ),
+                image_tokens=Product(
+                    name="Image tokens",
+                    included=PriceTier(unit_cost=0.5, up_to=0.75),
+                    tiers=[],
+                    unit="Million Tokens",
+                ),
                 embedding_tokens=Product(
                     name="Embedding tokens",
                     included=PriceTier(unit_cost=0.5, up_to=0.75),
@@ -514,6 +527,7 @@ def status(self) -> str:
 class ModelType(StrEnum):
     COMPLETION = "completion"
     LLM = "llm"
+    IMAGE_GEN = "image_gen"
     EMBED = "embed"
     RERANK = "rerank"
 
@@ -527,6 +541,7 @@ class ModelCapability(StrEnum):
     CHAT = "chat"
     TOOL = "tool"
     IMAGE = "image"  # TODO: Maybe change to "image_in" & "image_out"
+    IMAGE_OUT = "image_out"
     AUDIO = "audio"
     EMBED = "embed"
     RERANK = "rerank"
@@ -548,7 +563,7 @@ class ModelInfo(_BaseModel):
     )
     type: _ModelType = Field(
         "",
-        description="Model type. Can be completion, llm, embed, or rerank.",
+        description="Model type. Can be completion, llm, image_gen, embed, or rerank.",
         examples=[ModelType.LLM],
     )
     name: SanitisedNonEmptyStr = Field(
@@ -645,6 +660,15 @@ class ModelConfigUpdate(ModelInfo):
         -1.0,
         description="Cost in USD per million (mega) output / completion token.",
     )
+    # --- Image generation models --- #
+    image_input_cost_per_mtoken: float = Field(
+        -1.0,
+        description="Cost in USD per million (mega) image input tokens.",
+    )
+    image_output_cost_per_mtoken: float = Field(
+        -1.0,
+        description="Cost in USD per million (mega) image output tokens.",
+    )
     # --- Embedding models --- #
     embedding_size: PositiveNonZeroInt | None = Field(
         None,
@@ -703,6 +727,14 @@ def check_chat_cost_per_mtoken(self) -> Self:
             self.llm_output_cost_per_mtoken = 0.600
         return self
 
+    @model_validator(mode="after")
+    def check_image_cost_per_mtoken(self) -> Self:
+        if self.image_input_cost_per_mtoken < 0:
+            self.image_input_cost_per_mtoken = 0.0
+        if self.image_output_cost_per_mtoken < 0:
+            self.image_output_cost_per_mtoken = 0.0
+        return self
+
     @model_validator(mode="after")
     def check_embed_cost_per_mtoken(self) -> Self:
         # OpenAI text-embedding-3-small pricing (2024-09-09)
@@ -729,7 +761,7 @@ class ModelConfigCreate(ModelConfigUpdate):
         ),
     )
     type: _ModelType = Field(
-        description="Model type. Can be completion, llm, embed, or rerank.",
+        description="Model type. Can be completion, llm, image_gen, embed, or rerank.",
     )
     name: SanitisedNonEmptyStr = Field(
         max_length=255,
@@ -1087,6 +1119,12 @@ class Organization_(OrganizationCreate, _TableBase):
     llm_tokens_usage_mtok: float = Field(
         description="LLM token usage in millions of tokens.",
     )
+    image_tokens_quota_mtok: float | None = Field(
+        description="Image token quota in millions of tokens.",
+    )
+    image_tokens_usage_mtok: float = Field(
+        description="Image token usage in millions of tokens.",
+    )
     embedding_tokens_quota_mtok: float | None = Field(
         description="Embedding token quota in millions of tokens.",
     )

clients/python/src/jamaibase/types/gen_table.py

@@ -25,6 +25,7 @@
     ChatRequestBase,
     References,
 )
+from jamaibase.utils.exceptions import BadInputError
 from jamaibase.utils.types import StrEnum
 
 
@@ -130,6 +131,34 @@ def compat(cls, data: dict[str, Any] | BaseModel) -> dict[str, Any]:
         return data
 
 
+class ImageGenConfig(BaseModel):
+    object: Literal["gen_config.image"] = Field(
+        "gen_config.image",
+        description='The object type, which is always "gen_config.image".',
+        examples=["gen_config.image"],
+    )
+    model: str = Field(
+        "",
+        description='ID of the model to use. Defaults to "".',
+    )
+    prompt: str = Field(
+        "",
+        description="Prompt for the image generation/edit model.",
+    )
+    size: Literal["auto", "1024x1024", "1536x1024", "1024x1536"] | None = Field(
+        None,
+        description="Image size/aspect ratio hint. Defaults to None (provider default).",
+    )
+    quality: Literal["low", "medium", "high", "auto"] | None = Field(
+        None,
+        description="Image quality hint. Defaults to None (provider default).",
+    )
+    style: str | None = Field(
+        None,
+        description="Image style hint. Generation-only; ignored for edits.",
+    )
+
+
 class EmbedGenConfig(BaseModel):
     object: Literal["gen_config.embed"] = Field(
         "gen_config.embed",
@@ -179,6 +208,8 @@ def _gen_config_discriminator(x: Any) -> str | None:
     if isinstance(x, dict):
         if "object" in x:
             return x["object"]
+        if any(k in x for k in ("size", "quality", "style")):
+            raise BadInputError('ImageGenConfig requires explicit `object="gen_config.image"`.')
         if "embedding_model" in x:
             return "gen_config.embed"
         if "source_column" in x:
@@ -196,6 +227,7 @@ def _gen_config_discriminator(x: Any) -> str | None:
         Annotated[PythonGenConfig, Tag("gen_config.python")],
         Annotated[LLMGenConfig, Tag("gen_config.llm")],
         Annotated[LLMGenConfig, Tag("gen_config.chat")],
+        Annotated[ImageGenConfig, Tag("gen_config.image")],
         Annotated[EmbedGenConfig, Tag("gen_config.embed")],
     ],
     Discriminator(_gen_config_discriminator),

clients/python/src/jamaibase/types/lm.py

@@ -9,6 +9,7 @@
     ConfigDict,
     Field,
     field_validator,
+    model_serializer,
     model_validator,
 )
 
@@ -376,6 +377,14 @@ def _none_to_zero(v: int | None) -> int:
 
 
 class PromptUsageDetails(BaseModel):
+    text_tokens: int | None = Field(
+        None,
+        description="Text tokens present in the prompt.",
+    )
+    image_tokens: int | None = Field(
+        None,
+        description="Image tokens present in the prompt.",
+    )
     cached_tokens: ZeroIfNoneInt = Field(
         0,
         description="Cached tokens present in the prompt.",
@@ -385,8 +394,25 @@ class PromptUsageDetails(BaseModel):
         description="Audio input tokens present in the prompt or generated by the model.",
     )
 
+    @model_serializer(mode="wrap")
+    def _omit_optional_breakdown(self, handler):
+        data = handler(self)
+        if data.get("text_tokens") is None:
+            data.pop("text_tokens", None)
+        if data.get("image_tokens") is None:
+            data.pop("image_tokens", None)
+        return data
+
 
 class CompletionUsageDetails(BaseModel):
+    text_tokens: int | None = Field(
+        None,
+        description="Text tokens present in the completion.",
+    )
+    image_tokens: int | None = Field(
+        None,
+        description="Image tokens present in the completion.",
+    )
     audio_tokens: ZeroIfNoneInt = Field(
         0,
         description="Audio input tokens present in the prompt or generated by the model.",
@@ -404,6 +430,15 @@ class CompletionUsageDetails(BaseModel):
         description="When using Predicted Outputs, the number of tokens in the prediction that did not appear in the completion.",
     )
 
+    @model_serializer(mode="wrap")
+    def _omit_optional_breakdown(self, handler):
+        data = handler(self)
+        if data.get("text_tokens") is None:
+            data.pop("text_tokens", None)
+        if data.get("image_tokens") is None:
+            data.pop("image_tokens", None)
+        return data
+
 
 class ToolUsageDetails(BaseModel):
     web_search_calls: ZeroIfNoneInt = Field(

docker/ch_configs/create_ch_prom_db.sh

@@ -19,6 +19,28 @@ ENGINE=MergeTree
 PARTITION BY toYYYYMM(timestamp)
 ORDER BY (org_id, timestamp, model)"
 
+clickhouse-client --query="CREATE TABLE IF NOT EXISTS jamaibase_owl.image_gen_usage
+(
+    \`id\` UUID,
+    \`org_id\` String,
+    \`proj_id\` String,
+    \`user_id\` String,
+    \`timestamp\` DateTime64(6, 'UTC'),
+    \`model\` String,
+    \`text_input_token\` UInt32,
+    \`text_output_token\` UInt32,
+    \`image_input_token\` UInt32,
+    \`image_output_token\` UInt32,
+    \`text_input_cost\` Decimal128(12),
+    \`text_output_cost\` Decimal128(12),
+    \`image_input_cost\` Decimal128(12),
+    \`image_output_cost\` Decimal128(12),
+    \`cost\` Decimal128(12)
+)
+ENGINE=MergeTree
+PARTITION BY toYYYYMM(timestamp)
+ORDER BY (org_id, timestamp, model)"
+
 clickhouse-client --query="CREATE TABLE IF NOT EXISTS jamaibase_owl.embed_usage
 (
     \`id\` UUID,
@@ -179,4 +201,4 @@ clickhouse-client --query="ALTER TABLE jamaibase_owl.llm_usage MODIFY COLUMN out
 clickhouse-client --query="ALTER TABLE jamaibase_owl.embed_usage MODIFY COLUMN cost Decimal128(12)"
 clickhouse-client --query="ALTER TABLE jamaibase_owl.rerank_usage MODIFY COLUMN cost Decimal128(12)"
 clickhouse-client --query="ALTER TABLE jamaibase_owl.egress_usage MODIFY COLUMN cost Decimal128(12)"
-clickhouse-client --query="ALTER TABLE jamaibase_owl.egress_usage MODIFY COLUMN amount_gib Decimal128(12)"
+clickhouse-client --query="ALTER TABLE jamaibase_owl.egress_usage MODIFY COLUMN amount_gib Decimal128(12)"