Skip to content

Commit

Permalink
Fix typing issues
Browse files Browse the repository at this point in the history
  • Loading branch information
cpsievert committed Sep 18, 2024
1 parent d0d7719 commit 7688287
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 9 deletions.
12 changes: 6 additions & 6 deletions shiny/ui/_chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -540,7 +540,7 @@ async def _append_message(
self._current_stream_message = ""

msg = await self._transform_message(
msg, chunk=chunk, chunk_content=chunk_content
msg, chunk=chunk, chunk_content=str(chunk_content)
)
if msg is None:
return
Expand Down Expand Up @@ -776,11 +776,11 @@ async def _transform_message(
key = res["transform_key"]

if message["role"] == "user" and self._transform_user is not None:
content = await self._transform_user(message["content"])
content = await self._transform_user(str(message["content"]))

elif message["role"] == "assistant" and self._transform_assistant is not None:
content = await self._transform_assistant(
message["content"],
str(message["content"]),
chunk_content or "",
chunk == "end" or chunk is False,
)
Expand All @@ -790,7 +790,7 @@ async def _transform_message(
if content is None:
return None

res[key] = content
res[key] = content # type: ignore

return res

Expand Down Expand Up @@ -950,7 +950,7 @@ def user_input(self, transform: bool = False) -> str | None:
if msg is None:
return None
key = "content_server" if transform else "content_client"
return msg[key]
return str(msg[key])

def _user_input(self) -> str:
id = self.user_input_id
Expand Down Expand Up @@ -1097,7 +1097,7 @@ def as_transformed_message(message: ChatMessage) -> TransformedMessage:

return TransformedMessage(
content_client=message["content"],
content_server=message["content"],
content_server=str(message["content"]),
role=message["role"],
transform_key=transform_key,
pre_transform_key=pre_transform_key,
Expand Down
6 changes: 4 additions & 2 deletions shiny/ui/_chat_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,21 @@

from typing import Literal, TypedDict

from htmltools import HTML

Role = Literal["assistant", "user", "system"]


# TODO: content should probably be [{"type": "text", "content": "..."}, {"type": "image", ...}]
# in order to support multiple content types...
class ChatMessage(TypedDict):
content: str
content: str | HTML
role: Role


# A message once transformed have been applied
class TransformedMessage(TypedDict):
content_client: str
content_client: str | HTML
content_server: str
role: Role
transform_key: Literal["content_client", "content_server"]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@

# TODO: test with append_message_stream() as well
@chat.transform_assistant_response
def transform(content: str) -> str:
def transform(content: str) -> str | ui.HTML:
if content == "return HTML":
return ui.HTML(f"<b>Transformed response</b>: {content}")
else:
Expand Down

0 comments on commit 7688287

Please sign in to comment.