Skip to content

Commit 4461df1

Browse files
refactor(api): add SQLAlchemy 2.x Mapped type hints to Message model (#27709)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent f7f6b4a commit 4461df1

File tree

1 file changed

+24
-17
lines changed

1 file changed

+24
-17
lines changed

api/models/model.py

Lines changed: 24 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import uuid
44
from collections.abc import Mapping
55
from datetime import datetime
6+
from decimal import Decimal
67
from enum import StrEnum, auto
78
from typing import TYPE_CHECKING, Any, Literal, Optional, cast
89

@@ -914,34 +915,40 @@ class Message(Base):
914915
)
915916

916917
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
917-
app_id = mapped_column(StringUUID, nullable=False)
918-
model_provider = mapped_column(String(255), nullable=True)
919-
model_id = mapped_column(String(255), nullable=True)
920-
override_model_configs = mapped_column(sa.Text)
921-
conversation_id = mapped_column(StringUUID, sa.ForeignKey("conversations.id"), nullable=False)
918+
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
919+
model_provider: Mapped[str | None] = mapped_column(String(255), nullable=True)
920+
model_id: Mapped[str | None] = mapped_column(String(255), nullable=True)
921+
override_model_configs: Mapped[str | None] = mapped_column(sa.Text)
922+
conversation_id: Mapped[str] = mapped_column(StringUUID, sa.ForeignKey("conversations.id"), nullable=False)
922923
_inputs: Mapped[dict[str, Any]] = mapped_column("inputs", sa.JSON)
923924
query: Mapped[str] = mapped_column(sa.Text, nullable=False)
924-
message = mapped_column(sa.JSON, nullable=False)
925+
message: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False)
925926
message_tokens: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0"))
926-
message_unit_price = mapped_column(sa.Numeric(10, 4), nullable=False)
927-
message_price_unit = mapped_column(sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001"))
927+
message_unit_price: Mapped[Decimal] = mapped_column(sa.Numeric(10, 4), nullable=False)
928+
message_price_unit: Mapped[Decimal] = mapped_column(
929+
sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001")
930+
)
928931
answer: Mapped[str] = mapped_column(sa.Text, nullable=False)
929932
answer_tokens: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0"))
930-
answer_unit_price = mapped_column(sa.Numeric(10, 4), nullable=False)
931-
answer_price_unit = mapped_column(sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001"))
932-
parent_message_id = mapped_column(StringUUID, nullable=True)
933-
provider_response_latency = mapped_column(sa.Float, nullable=False, server_default=sa.text("0"))
934-
total_price = mapped_column(sa.Numeric(10, 7))
933+
answer_unit_price: Mapped[Decimal] = mapped_column(sa.Numeric(10, 4), nullable=False)
934+
answer_price_unit: Mapped[Decimal] = mapped_column(
935+
sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001")
936+
)
937+
parent_message_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
938+
provider_response_latency: Mapped[float] = mapped_column(sa.Float, nullable=False, server_default=sa.text("0"))
939+
total_price: Mapped[Decimal | None] = mapped_column(sa.Numeric(10, 7))
935940
currency: Mapped[str] = mapped_column(String(255), nullable=False)
936-
status = mapped_column(String(255), nullable=False, server_default=sa.text("'normal'::character varying"))
937-
error = mapped_column(sa.Text)
938-
message_metadata = mapped_column(sa.Text)
941+
status: Mapped[str] = mapped_column(
942+
String(255), nullable=False, server_default=sa.text("'normal'::character varying")
943+
)
944+
error: Mapped[str | None] = mapped_column(sa.Text)
945+
message_metadata: Mapped[str | None] = mapped_column(sa.Text)
939946
invoke_from: Mapped[str | None] = mapped_column(String(255), nullable=True)
940947
from_source: Mapped[str] = mapped_column(String(255), nullable=False)
941948
from_end_user_id: Mapped[str | None] = mapped_column(StringUUID)
942949
from_account_id: Mapped[str | None] = mapped_column(StringUUID)
943950
created_at: Mapped[datetime] = mapped_column(sa.DateTime, server_default=func.current_timestamp())
944-
updated_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
951+
updated_at: Mapped[datetime] = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
945952
agent_based: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
946953
workflow_run_id: Mapped[str | None] = mapped_column(StringUUID)
947954
app_mode: Mapped[str | None] = mapped_column(String(255), nullable=True)

0 commit comments

Comments
 (0)