From 4de179a076414fa4f9af2efa1fd34d44ea3f59a8 Mon Sep 17 00:00:00 2001
From: Dannon Baker <dannon.baker@gmail.com>
Date: Mon, 18 Nov 2024 21:22:45 -0500
Subject: [PATCH] Refactor typing annotations in ChatManager and ChatAPI to use
 Union for older python.

---
 lib/galaxy/managers/chat.py           | 25 ++++---------------------
 lib/galaxy/webapps/galaxy/api/chat.py | 24 ++++++++++++++++++------
 2 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/lib/galaxy/managers/chat.py b/lib/galaxy/managers/chat.py
index f3447e9aca72..ecf651602684 100644
--- a/lib/galaxy/managers/chat.py
+++ b/lib/galaxy/managers/chat.py
@@ -1,6 +1,5 @@
-from typing import Optional
+from typing import Union
 
-from fastapi import Path
 from sqlalchemy import (
     and_,
     select,
@@ -9,7 +8,6 @@
     MultipleResultsFound,
     NoResultFound,
 )
-from typing_extensions import Annotated
 
 from galaxy.exceptions import (
     InconsistentDatabase,
@@ -19,30 +17,15 @@
 from galaxy.managers.context import ProvidesUserContext
 from galaxy.model import ChatExchange
 from galaxy.model.base import transaction
-from galaxy.schema.fields import DecodedDatabaseIdField
 from galaxy.util import unicodify
 
-JobIdPathParam = Optional[
-    Annotated[
-        DecodedDatabaseIdField,
-        Path(title="Job ID", description="The Job ID the chat exchange is linked to."),
-    ]
-]
-
-MessageIdPathParam = Optional[
-    Annotated[
-        DecodedDatabaseIdField,
-        Path(title="Job ID", description="The ChatMessage ID."),
-    ]
-]
-
 
 class ChatManager:
     """
     Business logic for chat exchanges.
     """
 
-    def create(self, trans: ProvidesUserContext, job_id: JobIdPathParam, message: str) -> ChatExchange:
+    def create(self, trans: ProvidesUserContext, job_id: int, message: str) -> ChatExchange:
         """
         Create a new chat exchange in the DB.  Currently these are *only* job-based chat exchanges, will need to generalize down the road.
         :param  job_id:      id of the job to associate the response with
@@ -59,7 +42,7 @@ def create(self, trans: ProvidesUserContext, job_id: JobIdPathParam, message: st
             trans.sa_session.commit()
         return chat_exchange
 
-    def get(self, trans: ProvidesUserContext, job_id: JobIdPathParam) -> ChatExchange | None:
+    def get(self, trans: ProvidesUserContext, job_id: int) -> Union[ChatExchange, None]:
         """
         Returns the chat response from the DB based on the given job id.
         :param  job_id:      id of the job to load a response for from the DB
@@ -84,7 +67,7 @@ def get(self, trans: ProvidesUserContext, job_id: JobIdPathParam) -> ChatExchang
             raise InternalServerError(f"Error loading from the database.{unicodify(e)}")
         return chat_response
 
-    def set_feedback_for_job(self, trans: ProvidesUserContext, job_id: JobIdPathParam, feedback: int) -> ChatExchange:
+    def set_feedback_for_job(self, trans: ProvidesUserContext, job_id: int, feedback: int) -> ChatExchange:
         """
         Set the feedback for a chat response.
         :param  message_id:      id of the job to associate the feedback with
diff --git a/lib/galaxy/webapps/galaxy/api/chat.py b/lib/galaxy/webapps/galaxy/api/chat.py
index 5723fbf4815c..de32b6a452f4 100644
--- a/lib/galaxy/webapps/galaxy/api/chat.py
+++ b/lib/galaxy/webapps/galaxy/api/chat.py
@@ -3,15 +3,20 @@
 """
 
 import logging
+from typing import (
+    Optional,
+    Union,
+)
+
+from fastapi import Path
+from typing_extensions import Annotated
 
 from galaxy.config import GalaxyAppConfiguration
 from galaxy.exceptions import ConfigurationError
-from galaxy.managers.chat import (
-    ChatManager,
-    JobIdPathParam,
-)
+from galaxy.managers.chat import ChatManager
 from galaxy.managers.context import ProvidesUserContext
 from galaxy.managers.jobs import JobManager
+from galaxy.schema.fields import DecodedDatabaseIdField
 from galaxy.schema.schema import ChatPayload
 from galaxy.webapps.galaxy.api import (
     depends,
@@ -32,6 +37,13 @@
 Please only say that something went wrong when configuing the ai prompt in your response.
 """
 
+JobIdPathParam = Optional[
+    Annotated[
+        DecodedDatabaseIdField,
+        Path(title="Job ID", description="The Job ID the chat exchange is linked to."),
+    ]
+]
+
 
 @router.cbv
 class ChatAPI:
@@ -54,7 +66,7 @@ def query(
             # If there's an existing response for this job, just return that one for now.
             # TODO: Support regenerating the response as a new message, and
             # asking follow-up questions.
-            existing_response = self.chat_manager.get(trans, job_id)
+            existing_response = self.chat_manager.get(trans, job.id)
             if existing_response and existing_response.messages[0]:
                 return existing_response.messages[0].message
 
@@ -75,7 +87,7 @@ def feedback(
         job_id: JobIdPathParam,
         feedback: int,
         trans: ProvidesUserContext = DependsOnTrans,
-    ) -> int | None:
+    ) -> Union[int, None]:
         """Provide feedback on the chatbot response."""
         job = self.job_manager.get_accessible_job(trans, job_id)
         chat_response = self.chat_manager.set_feedback_for_job(trans, job.id, feedback)