From b47626b96ad261b83ff5e7cf4bb5998f5b023a0e Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Sat, 23 Nov 2024 11:47:34 -0800
Subject: [PATCH] typing

---
 backend/ee/danswer/db/query_history.py | 5 +++--
 backend/ee/danswer/db/usage_export.py  | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/backend/ee/danswer/db/query_history.py b/backend/ee/danswer/db/query_history.py
index b1eac1debd1..8fb77f0a2dd 100644
--- a/backend/ee/danswer/db/query_history.py
+++ b/backend/ee/danswer/db/query_history.py
@@ -8,6 +8,7 @@
 from sqlalchemy.orm import contains_eager
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm import Session
+from sqlalchemy.sql.expression import UnaryExpression
 
 from danswer.db.models import ChatMessage
 from danswer.db.models import ChatSession
@@ -22,8 +23,8 @@ def fetch_chat_sessions_eagerly_by_time(
     limit: int | None = 500,
     initial_time: datetime.datetime | None = None,
 ) -> list[ChatSession]:
-    time_order = desc(ChatSession.time_created)
-    message_order = asc(ChatMessage.id)
+    time_order: UnaryExpression = desc(ChatSession.time_created)
+    message_order: UnaryExpression = asc(ChatMessage.id)
 
     filters: list[ColumnElement | BinaryExpression] = [
         ChatSession.time_created.between(start, end)
diff --git a/backend/ee/danswer/db/usage_export.py b/backend/ee/danswer/db/usage_export.py
index 9a00b11f038..074e1ae7d6d 100644
--- a/backend/ee/danswer/db/usage_export.py
+++ b/backend/ee/danswer/db/usage_export.py
@@ -2,6 +2,7 @@
 from collections.abc import Generator
 from datetime import datetime
 from typing import IO
+from typing import Optional
 
 from fastapi_users_db_sqlalchemy import UUID_ID
 from sqlalchemy.orm import Session
@@ -21,7 +22,7 @@ def get_empty_chat_messages_entries__paginated(
     period: tuple[datetime, datetime],
     limit: int | None = 500,
     initial_time: datetime | None = None,
-) -> tuple[int, list[ChatMessageSkeleton]]:
+) -> tuple[Optional[datetime], list[ChatMessageSkeleton]]:
     chat_sessions = fetch_chat_sessions_eagerly_by_time(
         start=period[0],
         end=period[1],
@@ -63,7 +64,7 @@ def get_all_empty_chat_message_entries(
     db_session: Session,
     period: tuple[datetime, datetime],
 ) -> Generator[list[ChatMessageSkeleton], None, None]:
-    initial_time = period[0]
+    initial_time: Optional[datetime] = period[0]
     ind = 0
     while True:
         ind += 1