import json
import os
import uuid
from typing import Optional, List, Dict
import requests
from chainlit import PersistedUser
from chainlit.data import BaseDataLayer
from chainlit.types import PageInfo, ThreadFilter, ThreadDict, Pagination, PaginatedResponse
from literalai.helper import utc_now
fastgpt_base_url = os.getenv("FASTGPT_BASE_URL")
share_id = os.getenv("FASTGPT_SHARE_ID")
now = utc_now()
user_cur_threads = []
thread_user_dict = {}
def change_type(user_type: str):
if user_type == 'AI':
return 'assistant_message'
if user_type == 'Human':
return 'user_message'
def get_app_info():
with requests.get(f"{fastgpt_base_url}/api/core/chat/outLink/init?chatId=&shareId={share_id}&outLinkUid=123456") as resp:
app = {}
if resp.status_code == 200:
res = json.loads(resp.content)
app = res.get('data').get('app')
appId = res.get('data').get('appId')
app['id'] = appId
return app
app_info = get_app_info()
app_id = app_info.get('id')
app_name = app_info.get('name')
welcome_text = app_info.get('chatConfig').get('welcomeText')
def getHistories(user_id):
histories = []
if user_id:
with requests.post(f"{fastgpt_base_url}/api/core/chat/getHistories", data={"shareId": share_id, "outLinkUid": user_id}) as resp:
if resp.status_code == 200:
res = json.loads(resp.content)
data = res["data"]
print(data)
histories = [{"id": item["chatId"], "name": item["title"], "createdAt": item["updateTime"], "userId": user_id, "userIdentifier": user_id} for item in data]
if user_cur_threads:
thread = next((t for t in user_cur_threads if t["userId"] == user_id), None)
if thread: # 确保 thread 不为 None
thread_id = thread.get("id")
if histories:
# 检查 thread 的 ID 是否已存在于 threads 中
if not any(t.get("id") == thread_id for t in histories):
histories.insert(0, thread)
else:
# 如果 threads 是空列表，则直接插入 thread
histories.insert(0, thread)
for item in histories:
thread_user_dict[item.get('id')] = item.get('userId')
return histories
class FastgptDataLayer(BaseDataLayer):
async def get_user(self, identifier: str):
print('get_user', identifier)
return PersistedUser(id=identifier, createdAt=now, identifier=identifier)
async def update_thread(self, thread_id: str, name: Optional[str] = None, user_id: Optional[str] = None, metadata: Optional[Dict] = None, tags: Optional[List[str]] = None,):
print('---------update_thread----------', thread_id)
thread = next((t for t in user_cur_threads if t["userId"] == user_id), None)
if thread:
if thread_id:
thread["id"] = thread_id
if name:
thread["name"] = name
if user_id:
thread["userId"] = user_id
thread["userIdentifier"] = user_id
if metadata:
thread["metadata"] = metadata
if tags:
thread["tags"] = tags
thread["createdAt"] = utc_now()
else:
print('---------update_thread----------thread_id ', thread_id, name)
user_cur_threads.append({"id": thread_id, "name": name, "metadata": metadata, "tags": tags, "createdAt": utc_now(), "userId": user_id, "userIdentifier": user_id,})
async def get_thread_author(self, thread_id: str):
print('get_thread_author')
return thread_user_dict.get(thread_id, None)
async def list_threads(self, pagination: Pagination, filters: ThreadFilter) -> PaginatedResponse[ThreadDict]:
threads = []
if filters:
threads = getHistories(filters.userId)
search = filters.search if filters.search else ""
filtered_threads = [thread for thread in threads if search in thread.get('name', '')]
start = 0
if pagination.cursor:
for i, thread in enumerate(filtered_threads):
if thread["id"] == pagination.cursor: # Find the start index using pagination.cursor
start = i + 1
break
end = start + pagination.first
paginated_threads = filtered_threads[start:end] or []
has_next_page = len(paginated_threads) > end
start_cursor = paginated_threads[0]["id"] if paginated_threads else None
end_cursor = paginated_threads[-1]["id"] if paginated_threads else None
return PaginatedResponse(pageInfo=PageInfo(hasNextPage=has_next_page, startCursor=start_cursor, endCursor=end_cursor,), data=paginated_threads,)
async def get_thread(self, thread_id: str):
print('get_thread', thread_id)
user_id = thread_user_dict.get(thread_id, None)
thread = None
if user_id:
params = {'chatId': thread_id, 'shareId': share_id, 'outLinkUid': user_id,}
with requests.get(f"{fastgpt_base_url}/api/core/chat/outLink/init", params=params,) as resp:
if resp.status_code == 200:
res = json.loads(resp.content)
data = res["data"]
if data:
history = data['history']
files = []
texts = []
for item in history:
for entry in item['value']:
if entry.get('type') == 'text':
text = {"id": item["_id"], "threadId": thread_id, "name": item["obj"], "type": change_type(item["obj"]), "input": None, "createdAt": utc_now(), "output": entry.get('text').get('content'),}
texts.append(text)
if entry.get('type') == 'file':
file = {"id": str(uuid.UUID), "threadId": thread_id, "forId": item["_id"], "name": entry.get('file').get('name'), "type": entry.get('file').get('type'), "url": entry.get('file').get('url'), "display": "inline", "size": "medium"}
files.append(file)
thread = {"id": thread_id, "name": data.get("title", ''), "createdAt": utc_now(), "userId": "admin", "userIdentifier": "admin", "metadata": {"appId": data["appId"]}, "steps": texts, "elements": files,}
return thread
async def delete_thread(self, thread_id: str):
print('delete_thread')
thread = next((t for t in user_cur_threads if t["id"] == thread_id), None)
user_id = thread_user_dict.get(thread_id, None)
if thread:
user_cur_threads.remove(thread)
if user_id:
params = {'appId': app_id, 'chatId': thread_id, 'shareId': share_id, 'outLinkUid': user_id,}
requests.get(f"{fastgpt_base_url}/api/core/chat/delHistory", params=params)