349 lines
10 KiB
Python
349 lines
10 KiB
Python
|
import asyncio
|
||
|
|
import base64
|
||
|
|
import mimetypes
|
||
|
|
import os
|
||
|
|
import uuid
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
|
||
|
|
|
||
|
|
import filetype
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from botbuilder.core import TurnContext
|
||
|
|
from botbuilder.schema import Activity
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from botbuilder.core import (
|
||
|
|
BotFrameworkAdapter,
|
||
|
|
BotFrameworkAdapterSettings,
|
||
|
|
MessageFactory,
|
||
|
|
TurnContext,
|
||
|
|
)
|
||
|
|
from botbuilder.schema import (
|
||
|
|
ActionTypes,
|
||
|
|
Activity,
|
||
|
|
ActivityTypes,
|
||
|
|
Attachment,
|
||
|
|
CardAction,
|
||
|
|
ChannelAccount,
|
||
|
|
HeroCard,
|
||
|
|
)
|
||
|
|
|
||
|
|
from chainlit.config import config
|
||
|
|
from chainlit.context import ChainlitContext, HTTPSession, context, context_var
|
||
|
|
from chainlit.data import get_data_layer
|
||
|
|
from chainlit.element import Element, ElementDict
|
||
|
|
from chainlit.emitter import BaseChainlitEmitter
|
||
|
|
from chainlit.logger import logger
|
||
|
|
from chainlit.message import Message, StepDict
|
||
|
|
from chainlit.types import Feedback
|
||
|
|
from chainlit.user import PersistedUser, User
|
||
|
|
from chainlit.user_session import user_session
|
||
|
|
|
||
|
|
|
||
|
|
class TeamsEmitter(BaseChainlitEmitter):
|
||
|
|
def __init__(self, session: HTTPSession, turn_context: TurnContext):
|
||
|
|
super().__init__(session)
|
||
|
|
self.turn_context = turn_context
|
||
|
|
|
||
|
|
async def send_element(self, element_dict: ElementDict):
|
||
|
|
if element_dict.get("display") != "inline":
|
||
|
|
return
|
||
|
|
|
||
|
|
persisted_file = self.session.files.get(element_dict.get("chainlitKey") or "")
|
||
|
|
attachment: Optional[Attachment] = None
|
||
|
|
mime: Optional[str] = None
|
||
|
|
|
||
|
|
element_name: str = element_dict.get("name", "Untitled")
|
||
|
|
|
||
|
|
if mime:
|
||
|
|
file_extension = mimetypes.guess_extension(mime)
|
||
|
|
if file_extension:
|
||
|
|
element_name += file_extension
|
||
|
|
|
||
|
|
if persisted_file:
|
||
|
|
mime = element_dict.get("mime")
|
||
|
|
with open(persisted_file["path"], "rb") as file:
|
||
|
|
dencoded_string = base64.b64encode(file.read()).decode()
|
||
|
|
content_url = f"data:{mime};base64,{dencoded_string}"
|
||
|
|
attachment = Attachment(
|
||
|
|
content_type=mime, content_url=content_url, name=element_name
|
||
|
|
)
|
||
|
|
|
||
|
|
elif url := element_dict.get("url"):
|
||
|
|
attachment = Attachment(
|
||
|
|
content_type=mime, content_url=url, name=element_name
|
||
|
|
)
|
||
|
|
|
||
|
|
if not attachment:
|
||
|
|
return
|
||
|
|
|
||
|
|
await self.turn_context.send_activity(Activity(attachments=[attachment]))
|
||
|
|
|
||
|
|
async def send_step(self, step_dict: StepDict):
|
||
|
|
if not step_dict["type"] == "assistant_message":
|
||
|
|
return
|
||
|
|
|
||
|
|
step_type = step_dict.get("type")
|
||
|
|
is_message = step_type in [
|
||
|
|
"user_message",
|
||
|
|
"assistant_message",
|
||
|
|
]
|
||
|
|
is_empty_output = not step_dict.get("output")
|
||
|
|
|
||
|
|
if is_empty_output or not is_message:
|
||
|
|
return
|
||
|
|
else:
|
||
|
|
reply = MessageFactory.text(step_dict["output"])
|
||
|
|
enable_feedback = get_data_layer()
|
||
|
|
if enable_feedback:
|
||
|
|
current_run = context.current_run
|
||
|
|
scorable_id = current_run.id if current_run else step_dict["id"]
|
||
|
|
like_button = CardAction(
|
||
|
|
type=ActionTypes.message_back,
|
||
|
|
title="👍",
|
||
|
|
text="like",
|
||
|
|
value={"feedback": "like", "step_id": scorable_id},
|
||
|
|
)
|
||
|
|
dislike_button = CardAction(
|
||
|
|
type=ActionTypes.message_back,
|
||
|
|
title="👎",
|
||
|
|
text="dislike",
|
||
|
|
value={"feedback": "dislike", "step_id": scorable_id},
|
||
|
|
)
|
||
|
|
card = HeroCard(buttons=[like_button, dislike_button])
|
||
|
|
attachment = Attachment(
|
||
|
|
content_type="application/vnd.microsoft.card.hero", content=card
|
||
|
|
)
|
||
|
|
reply.attachments = [attachment]
|
||
|
|
|
||
|
|
await self.turn_context.send_activity(reply)
|
||
|
|
|
||
|
|
async def update_step(self, step_dict: StepDict):
|
||
|
|
if not step_dict["type"] == "assistant_message":
|
||
|
|
return
|
||
|
|
|
||
|
|
await self.send_step(step_dict)
|
||
|
|
|
||
|
|
|
||
|
|
adapter_settings = BotFrameworkAdapterSettings(
|
||
|
|
app_id=os.environ.get("TEAMS_APP_ID"),
|
||
|
|
app_password=os.environ.get("TEAMS_APP_PASSWORD"),
|
||
|
|
)
|
||
|
|
adapter = BotFrameworkAdapter(adapter_settings)
|
||
|
|
|
||
|
|
|
||
|
|
def init_teams_context(
|
||
|
|
session: HTTPSession,
|
||
|
|
turn_context: TurnContext,
|
||
|
|
) -> ChainlitContext:
|
||
|
|
emitter = TeamsEmitter(session=session, turn_context=turn_context)
|
||
|
|
context = ChainlitContext(session=session, emitter=emitter)
|
||
|
|
context_var.set(context)
|
||
|
|
user_session.set("teams_turn_context", turn_context)
|
||
|
|
return context
|
||
|
|
|
||
|
|
|
||
|
|
users_by_teams_id: Dict[str, Union[User, PersistedUser]] = {}
|
||
|
|
|
||
|
|
USER_PREFIX = "teams_"
|
||
|
|
|
||
|
|
|
||
|
|
async def get_user(teams_user: ChannelAccount):
|
||
|
|
if teams_user.id in users_by_teams_id:
|
||
|
|
return users_by_teams_id[teams_user.id]
|
||
|
|
|
||
|
|
metadata = {
|
||
|
|
"name": teams_user.name,
|
||
|
|
"id": teams_user.id,
|
||
|
|
}
|
||
|
|
user = User(identifier=USER_PREFIX + str(teams_user.name), metadata=metadata)
|
||
|
|
|
||
|
|
users_by_teams_id[teams_user.id] = user
|
||
|
|
|
||
|
|
if data_layer := get_data_layer():
|
||
|
|
try:
|
||
|
|
persisted_user = await data_layer.create_user(user)
|
||
|
|
if persisted_user:
|
||
|
|
users_by_teams_id[teams_user.id] = persisted_user
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error creating user: {e}")
|
||
|
|
|
||
|
|
return users_by_teams_id[teams_user.id]
|
||
|
|
|
||
|
|
|
||
|
|
async def download_teams_file(url: str):
|
||
|
|
async with httpx.AsyncClient() as client:
|
||
|
|
response = await client.get(url)
|
||
|
|
if response.status_code == 200:
|
||
|
|
return response.content
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
async def download_teams_files(
|
||
|
|
session: HTTPSession, attachments: Optional[List[Attachment]] = None
|
||
|
|
):
|
||
|
|
if not attachments:
|
||
|
|
return []
|
||
|
|
|
||
|
|
attachments = [
|
||
|
|
attachment for attachment in attachments if isinstance(attachment.content, dict)
|
||
|
|
]
|
||
|
|
download_coros = [
|
||
|
|
download_teams_file(attachment.content.get("downloadUrl"))
|
||
|
|
for attachment in attachments
|
||
|
|
]
|
||
|
|
file_bytes_list = await asyncio.gather(*download_coros)
|
||
|
|
file_refs = []
|
||
|
|
for idx, file_bytes in enumerate(file_bytes_list):
|
||
|
|
if file_bytes:
|
||
|
|
name = attachments[idx].name
|
||
|
|
mime_type = filetype.guess_mime(file_bytes) or "application/octet-stream"
|
||
|
|
file_ref = await session.persist_file(
|
||
|
|
name=name, mime=mime_type, content=file_bytes
|
||
|
|
)
|
||
|
|
file_refs.append(file_ref)
|
||
|
|
|
||
|
|
files_dicts = [
|
||
|
|
session.files[file["id"]] for file in file_refs if file["id"] in session.files
|
||
|
|
]
|
||
|
|
|
||
|
|
elements = [
|
||
|
|
Element.from_dict(
|
||
|
|
{
|
||
|
|
"id": file["id"],
|
||
|
|
"name": file["name"],
|
||
|
|
"path": str(file["path"]),
|
||
|
|
"chainlitKey": file["id"],
|
||
|
|
"display": "inline",
|
||
|
|
"type": Element.infer_type_from_mime(file["type"]),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
for file in files_dicts
|
||
|
|
]
|
||
|
|
|
||
|
|
return elements
|
||
|
|
|
||
|
|
|
||
|
|
def clean_content(activity: Activity):
|
||
|
|
return activity.text.strip()
|
||
|
|
|
||
|
|
|
||
|
|
async def process_teams_message(
|
||
|
|
turn_context: TurnContext,
|
||
|
|
thread_name: str,
|
||
|
|
):
|
||
|
|
user = await get_user(turn_context.activity.from_property)
|
||
|
|
|
||
|
|
thread_id = str(
|
||
|
|
uuid.uuid5(
|
||
|
|
uuid.NAMESPACE_DNS,
|
||
|
|
str(
|
||
|
|
turn_context.activity.conversation.id
|
||
|
|
+ datetime.today().strftime("%Y-%m-%d")
|
||
|
|
),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
text = clean_content(turn_context.activity)
|
||
|
|
teams_files = turn_context.activity.attachments
|
||
|
|
|
||
|
|
session_id = str(uuid.uuid4())
|
||
|
|
|
||
|
|
session = HTTPSession(
|
||
|
|
id=session_id,
|
||
|
|
thread_id=thread_id,
|
||
|
|
user=user,
|
||
|
|
client_type="teams",
|
||
|
|
)
|
||
|
|
|
||
|
|
ctx = init_teams_context(
|
||
|
|
session=session,
|
||
|
|
turn_context=turn_context,
|
||
|
|
)
|
||
|
|
|
||
|
|
file_elements = await download_teams_files(session, teams_files)
|
||
|
|
|
||
|
|
if on_chat_start := config.code.on_chat_start:
|
||
|
|
await on_chat_start()
|
||
|
|
|
||
|
|
msg = Message(
|
||
|
|
content=text,
|
||
|
|
elements=file_elements,
|
||
|
|
type="user_message",
|
||
|
|
author=user.metadata.get("name"),
|
||
|
|
)
|
||
|
|
|
||
|
|
await msg.send()
|
||
|
|
|
||
|
|
if on_message := config.code.on_message:
|
||
|
|
await on_message(msg)
|
||
|
|
|
||
|
|
if on_chat_end := config.code.on_chat_end:
|
||
|
|
await on_chat_end()
|
||
|
|
|
||
|
|
if data_layer := get_data_layer():
|
||
|
|
if isinstance(user, PersistedUser):
|
||
|
|
try:
|
||
|
|
await data_layer.update_thread(
|
||
|
|
thread_id=thread_id,
|
||
|
|
name=thread_name,
|
||
|
|
metadata=ctx.session.to_persistable(),
|
||
|
|
user_id=user.id,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error updating thread: {e}")
|
||
|
|
|
||
|
|
await ctx.session.delete()
|
||
|
|
|
||
|
|
|
||
|
|
async def handle_message(turn_context: TurnContext):
|
||
|
|
if turn_context.activity.type == ActivityTypes.message:
|
||
|
|
if (
|
||
|
|
turn_context.activity.text == "like"
|
||
|
|
or turn_context.activity.text == "dislike"
|
||
|
|
):
|
||
|
|
feedback_value: Literal[0, 1] = (
|
||
|
|
0 if turn_context.activity.text == "dislike" else 1
|
||
|
|
)
|
||
|
|
step_id = turn_context.activity.value.get("step_id")
|
||
|
|
if data_layer := get_data_layer():
|
||
|
|
await data_layer.upsert_feedback(
|
||
|
|
Feedback(forId=step_id, value=feedback_value)
|
||
|
|
)
|
||
|
|
updated_text = "👍" if turn_context.activity.text == "like" else "👎"
|
||
|
|
# Update the existing message to remove the buttons
|
||
|
|
updated_message = Activity(
|
||
|
|
type=ActivityTypes.message,
|
||
|
|
id=turn_context.activity.reply_to_id,
|
||
|
|
text=updated_text,
|
||
|
|
attachments=[],
|
||
|
|
)
|
||
|
|
await turn_context.update_activity(updated_message)
|
||
|
|
else:
|
||
|
|
# Send typing activity
|
||
|
|
typing_activity = Activity(
|
||
|
|
type=ActivityTypes.typing,
|
||
|
|
from_property=turn_context.activity.recipient,
|
||
|
|
recipient=turn_context.activity.from_property,
|
||
|
|
conversation=turn_context.activity.conversation,
|
||
|
|
)
|
||
|
|
await turn_context.send_activity(typing_activity)
|
||
|
|
thread_name = f"{turn_context.activity.from_property.name} Teams DM {datetime.today().strftime('%Y-%m-%d')}"
|
||
|
|
await process_teams_message(turn_context, thread_name)
|
||
|
|
|
||
|
|
|
||
|
|
async def on_turn(turn_context: TurnContext):
|
||
|
|
await handle_message(turn_context)
|
||
|
|
|
||
|
|
|
||
|
|
# Create the main bot class
|
||
|
|
class TeamsBot:
|
||
|
|
async def on_turn(self, turn_context: TurnContext):
|
||
|
|
await on_turn(turn_context)
|
||
|
|
|
||
|
|
|
||
|
|
# Create the bot instance
|
||
|
|
bot = TeamsBot()
|