wehub-resource-sync
b7f52be4c9
CI / Run CI (push) Has been cancelled
CI / check-backend (push) Has been cancelled
CI / check-frontend (push) Has been cancelled
CI / tests (push) Has been cancelled
CI / e2e-tests (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
276 lines
8.1 KiB
Python
276 lines
8.1 KiB
Python
import os
|
|
import os.path
|
|
import pickle
|
|
from typing import Dict, List, Optional
|
|
|
|
import chainlit as cl
|
|
import chainlit.data as cl_data
|
|
from chainlit.data.utils import queue_until_user_message
|
|
from chainlit.element import Element, ElementDict
|
|
from chainlit.socket import persist_user_session
|
|
from chainlit.step import StepDict
|
|
from chainlit.types import (
|
|
Feedback,
|
|
PageInfo,
|
|
PaginatedResponse,
|
|
Pagination,
|
|
ThreadDict,
|
|
ThreadFilter,
|
|
)
|
|
from chainlit.utils import utc_now
|
|
|
|
os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET" # nosec B105
|
|
|
|
now = utc_now()
|
|
|
|
thread_history = [
|
|
{
|
|
"id": "test1",
|
|
"name": "thread 1",
|
|
"createdAt": now,
|
|
"userId": "user1_id",
|
|
"userIdentifier": "user1",
|
|
"steps": [
|
|
{
|
|
"id": "test1",
|
|
"name": "test",
|
|
"createdAt": now,
|
|
"type": "user_message",
|
|
"output": "Message 1",
|
|
},
|
|
{
|
|
"id": "test2",
|
|
"name": "test",
|
|
"createdAt": now,
|
|
"type": "assistant_message",
|
|
"output": "Message 2",
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"id": "test2",
|
|
"createdAt": now,
|
|
"userId": "user1_id",
|
|
"userIdentifier": "user1",
|
|
"name": "thread 2",
|
|
"steps": [
|
|
{
|
|
"id": "test3",
|
|
"createdAt": now,
|
|
"name": "test",
|
|
"type": "user_message",
|
|
"output": "Message 3",
|
|
},
|
|
{
|
|
"id": "test4",
|
|
"createdAt": now,
|
|
"name": "test",
|
|
"type": "assistant_message",
|
|
"output": "Message 4",
|
|
},
|
|
],
|
|
},
|
|
] # type: List[ThreadDict]
|
|
deleted_thread_ids = [] # type: List[str]
|
|
ELEMENTS_STORAGE = []
|
|
|
|
THREAD_HISTORY_PICKLE_PATH = os.path.join(
|
|
os.path.dirname(__file__), "thread_history.pickle"
|
|
)
|
|
if THREAD_HISTORY_PICKLE_PATH and os.path.exists(THREAD_HISTORY_PICKLE_PATH):
|
|
with open(THREAD_HISTORY_PICKLE_PATH, "rb") as f:
|
|
thread_history = pickle.load(f)
|
|
|
|
|
|
async def save_thread_history():
|
|
# Force saving of thread history for reload when server restarts
|
|
await persist_user_session(
|
|
cl.context.session.thread_id, cl.context.session.to_persistable()
|
|
)
|
|
|
|
with open(THREAD_HISTORY_PICKLE_PATH, "wb") as out_file:
|
|
pickle.dump(thread_history, out_file)
|
|
|
|
|
|
class TestDataLayer(cl_data.BaseDataLayer):
|
|
async def get_user(self, identifier: str):
|
|
if identifier == "user1":
|
|
return cl.PersistedUser(id="user1_id", createdAt=now, identifier=identifier)
|
|
elif identifier == "user2":
|
|
return cl.PersistedUser(id="user2_id", createdAt=now, identifier=identifier)
|
|
return None
|
|
|
|
async def create_user(self, user: cl.User):
|
|
if user.identifier == "user1":
|
|
return cl.PersistedUser(
|
|
id="user1_id", createdAt=now, identifier=user.identifier
|
|
)
|
|
elif user.identifier == "user2":
|
|
return cl.PersistedUser(
|
|
id="user2_id", createdAt=now, identifier=user.identifier
|
|
)
|
|
return None
|
|
|
|
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,
|
|
):
|
|
thread = next((t for t in thread_history if t["id"] == thread_id), None)
|
|
if thread:
|
|
if name:
|
|
thread["name"] = name
|
|
if metadata:
|
|
thread["metadata"] = metadata
|
|
if tags:
|
|
thread["tags"] = tags
|
|
else:
|
|
thread_history.append(
|
|
{
|
|
"id": thread_id,
|
|
"name": name,
|
|
"metadata": metadata,
|
|
"tags": tags,
|
|
"createdAt": utc_now(),
|
|
"userId": user_id,
|
|
"userIdentifier": "user1"
|
|
if user_id == "user1_id"
|
|
else "user2"
|
|
if user_id == "user2_id"
|
|
else "unknown",
|
|
"steps": [],
|
|
}
|
|
)
|
|
|
|
@cl_data.queue_until_user_message()
|
|
async def create_step(self, step_dict: StepDict):
|
|
cl.user_session.set(
|
|
"create_step_counter", cl.user_session.get("create_step_counter") + 1
|
|
)
|
|
|
|
thread = next(
|
|
(t for t in thread_history if t["id"] == step_dict.get("threadId")), None
|
|
)
|
|
if thread:
|
|
thread["steps"].append(step_dict)
|
|
|
|
async def get_thread_author(self, thread_id: str):
|
|
thread = await self.get_thread(thread_id)
|
|
return thread["userIdentifier"] if thread else None
|
|
|
|
async def list_threads(
|
|
self, pagination: Pagination, filters: ThreadFilter
|
|
) -> PaginatedResponse[ThreadDict]:
|
|
return PaginatedResponse(
|
|
data=[t for t in thread_history if t["id"] not in deleted_thread_ids],
|
|
pageInfo=PageInfo(hasNextPage=False, startCursor=None, endCursor=None),
|
|
)
|
|
|
|
async def get_thread(self, thread_id: str):
|
|
thread = next((t for t in thread_history if t["id"] == thread_id), None)
|
|
if not thread:
|
|
return None
|
|
thread["steps"] = sorted(thread["steps"], key=lambda x: x["createdAt"])
|
|
return thread
|
|
|
|
async def delete_thread(self, thread_id: str):
|
|
deleted_thread_ids.append(thread_id)
|
|
|
|
async def delete_feedback(
|
|
self,
|
|
feedback_id: str,
|
|
) -> bool:
|
|
return True
|
|
|
|
async def upsert_feedback(
|
|
self,
|
|
feedback: Feedback,
|
|
) -> str:
|
|
return ""
|
|
|
|
@queue_until_user_message()
|
|
async def create_element(self, element: "Element"):
|
|
if element.url == "http://example.org/test.txt":
|
|
element.url = "http://example.com/test.txt"
|
|
|
|
ELEMENTS_STORAGE.append(element.to_dict())
|
|
|
|
async def get_element(
|
|
self, thread_id: str, element_id: str
|
|
) -> Optional["ElementDict"]:
|
|
return next((e for e in ELEMENTS_STORAGE if e["id"] == element_id), None)
|
|
|
|
@queue_until_user_message()
|
|
async def delete_element(self, element_id: str, thread_id: Optional[str] = None):
|
|
pass
|
|
|
|
@queue_until_user_message()
|
|
async def update_step(self, step_dict: "StepDict"):
|
|
pass
|
|
|
|
@queue_until_user_message()
|
|
async def delete_step(self, step_id: str):
|
|
pass
|
|
|
|
async def get_favorite_steps(self, user_id: str) -> List["StepDict"]:
|
|
return []
|
|
|
|
async def build_debug_url(self) -> str:
|
|
return ""
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
|
|
@cl.data_layer
|
|
def data_layer():
|
|
return TestDataLayer()
|
|
|
|
|
|
async def send_count():
|
|
create_step_counter = cl.user_session.get("create_step_counter")
|
|
await cl.Message(f"Create step counter: {create_step_counter}").send()
|
|
|
|
|
|
@cl.on_chat_start
|
|
async def main():
|
|
# Add step counter to session so that it is saved in thread metadata
|
|
cl.user_session.set("create_step_counter", 0)
|
|
await cl.Message("Hello, send me a message!").send()
|
|
await send_count()
|
|
|
|
|
|
@cl.on_message
|
|
async def handle_message():
|
|
# Wait for queue to be flushed
|
|
await cl.sleep(2)
|
|
await send_count()
|
|
async with cl.Step(type="tool", name="thinking") as step:
|
|
step.output = "Thinking..."
|
|
await cl.Message("Ok!").send()
|
|
await send_count()
|
|
|
|
await save_thread_history()
|
|
|
|
|
|
@cl.password_auth_callback
|
|
def auth_callback(username: str, password: str) -> Optional[cl.User]:
|
|
if (username, password) == ("user1", "user1"):
|
|
return cl.User(identifier="user1")
|
|
elif (username, password) == ("user2", "user2"):
|
|
return cl.User(identifier="user2")
|
|
else:
|
|
return None
|
|
|
|
|
|
@cl.on_chat_resume
|
|
async def on_chat_resume(thread: ThreadDict):
|
|
await cl.Message(f"Welcome back to {thread['name']}").send()
|
|
if "metadata" in thread:
|
|
await cl.Message(thread["metadata"], author="metadata", language="json").send()
|
|
if "tags" in thread:
|
|
await cl.Message(thread["tags"], author="tags", language="json").send()
|