Skip to content

Commit

Permalink
Refactor typing annotations in ChatManager and ChatAPI to use Union f…
Browse files Browse the repository at this point in the history
…or older python.
  • Loading branch information
dannon committed Nov 19, 2024
1 parent 8aadf1a commit 4de179a
Showing 2 changed files with 22 additions and 27 deletions.
25 changes: 4 additions & 21 deletions lib/galaxy/managers/chat.py
Original file line number Diff line number Diff line change
@@ -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
24 changes: 18 additions & 6 deletions lib/galaxy/webapps/galaxy/api/chat.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 4de179a

Please sign in to comment.