chore: import upstream snapshot with attribution
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
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
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,740 @@
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
|
||||
import aiofiles
|
||||
import asyncpg # type: ignore
|
||||
|
||||
from chainlit.data.base import BaseDataLayer
|
||||
from chainlit.data.storage_clients.base import BaseStorageClient
|
||||
from chainlit.data.utils import queue_until_user_message
|
||||
from chainlit.element import ElementDict
|
||||
from chainlit.logger import logger
|
||||
from chainlit.step import StepDict
|
||||
from chainlit.types import (
|
||||
Feedback,
|
||||
FeedbackDict,
|
||||
PageInfo,
|
||||
PaginatedResponse,
|
||||
Pagination,
|
||||
ThreadDict,
|
||||
ThreadFilter,
|
||||
)
|
||||
from chainlit.user import PersistedUser, User
|
||||
|
||||
# Import for runtime usage (isinstance checks)
|
||||
try:
|
||||
from chainlit.data.storage_clients.gcs import GCSStorageClient
|
||||
except ImportError:
|
||||
GCSStorageClient = None # type: ignore[assignment,misc]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from chainlit.data.storage_clients.gcs import GCSStorageClient
|
||||
from chainlit.element import Element, ElementDict
|
||||
from chainlit.step import StepDict
|
||||
|
||||
ISO_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
|
||||
|
||||
|
||||
class ChainlitDataLayer(BaseDataLayer):
|
||||
def __init__(
|
||||
self,
|
||||
database_url: str,
|
||||
storage_client: Optional[BaseStorageClient] = None,
|
||||
show_logger: bool = False,
|
||||
):
|
||||
self.database_url = database_url
|
||||
self.pool: Optional[asyncpg.Pool] = None
|
||||
self.storage_client = storage_client
|
||||
self.show_logger = show_logger
|
||||
|
||||
async def connect(self):
|
||||
if not self.pool:
|
||||
self.pool = await asyncpg.create_pool(self.database_url)
|
||||
|
||||
async def get_current_timestamp(self) -> datetime:
|
||||
return datetime.now()
|
||||
|
||||
async def execute_query(
|
||||
self, query: str, params: Union[Dict, None] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
if not self.pool:
|
||||
await self.connect()
|
||||
|
||||
try:
|
||||
async with self.pool.acquire() as connection: # type: ignore
|
||||
try:
|
||||
if params:
|
||||
records = await connection.fetch(query, *params.values())
|
||||
else:
|
||||
records = await connection.fetch(query)
|
||||
return [dict(record) for record in records]
|
||||
except Exception as e:
|
||||
logger.error(f"Database error: {e!s}")
|
||||
raise
|
||||
except (
|
||||
asyncpg.exceptions.ConnectionDoesNotExistError,
|
||||
asyncpg.exceptions.InterfaceError,
|
||||
) as e:
|
||||
# Handle connection issues by cleaning up and rethrowing
|
||||
logger.error(f"Connection error: {e!s}")
|
||||
await self.cleanup()
|
||||
raise
|
||||
|
||||
async def get_user(self, identifier: str) -> Optional[PersistedUser]:
|
||||
query = """
|
||||
SELECT * FROM "User"
|
||||
WHERE identifier = $1
|
||||
"""
|
||||
result = await self.execute_query(query, {"identifier": identifier})
|
||||
if not result or len(result) == 0:
|
||||
return None
|
||||
row = result[0]
|
||||
|
||||
return PersistedUser(
|
||||
id=str(row.get("id")),
|
||||
identifier=str(row.get("identifier")),
|
||||
createdAt=row.get("createdAt").isoformat(), # type: ignore
|
||||
metadata=json.loads(row.get("metadata", "{}")),
|
||||
)
|
||||
|
||||
async def create_user(self, user: User) -> Optional[PersistedUser]:
|
||||
query = """
|
||||
INSERT INTO "User" (id, identifier, metadata, "createdAt", "updatedAt")
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (identifier) DO UPDATE
|
||||
SET metadata = $3
|
||||
RETURNING *
|
||||
"""
|
||||
now = await self.get_current_timestamp()
|
||||
params = {
|
||||
"id": str(uuid.uuid4()),
|
||||
"identifier": user.identifier,
|
||||
"metadata": json.dumps(user.metadata),
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
result = await self.execute_query(query, params)
|
||||
row = result[0]
|
||||
|
||||
return PersistedUser(
|
||||
id=str(row.get("id")),
|
||||
identifier=str(row.get("identifier")),
|
||||
createdAt=row.get("createdAt").isoformat(), # type: ignore
|
||||
metadata=json.loads(row.get("metadata", "{}")),
|
||||
)
|
||||
|
||||
async def delete_feedback(self, feedback_id: str) -> bool:
|
||||
query = """
|
||||
DELETE FROM "Feedback" WHERE id = $1
|
||||
"""
|
||||
await self.execute_query(query, {"feedback_id": feedback_id})
|
||||
return True
|
||||
|
||||
async def upsert_feedback(self, feedback: Feedback) -> str:
|
||||
query = """
|
||||
INSERT INTO "Feedback" (id, "stepId", name, value, comment)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (id) DO UPDATE
|
||||
SET value = $4, comment = $5
|
||||
RETURNING id
|
||||
"""
|
||||
feedback_id = feedback.id or str(uuid.uuid4())
|
||||
params = {
|
||||
"id": feedback_id,
|
||||
"step_id": feedback.forId,
|
||||
"name": "user_feedback",
|
||||
"value": float(feedback.value),
|
||||
"comment": feedback.comment,
|
||||
}
|
||||
results = await self.execute_query(query, params)
|
||||
return str(results[0]["id"])
|
||||
|
||||
@queue_until_user_message()
|
||||
async def create_element(self, element: "Element"):
|
||||
if not element.for_id:
|
||||
return
|
||||
|
||||
if element.thread_id:
|
||||
query = 'SELECT id FROM "Thread" WHERE id = $1'
|
||||
results = await self.execute_query(query, {"thread_id": element.thread_id})
|
||||
if not results:
|
||||
await self.update_thread(thread_id=element.thread_id)
|
||||
|
||||
if element.for_id:
|
||||
query = 'SELECT id FROM "Step" WHERE id = $1'
|
||||
results = await self.execute_query(query, {"step_id": element.for_id})
|
||||
if not results:
|
||||
await self.create_step(
|
||||
{
|
||||
"id": element.for_id,
|
||||
"metadata": {},
|
||||
"type": "run",
|
||||
"start_time": await self.get_current_timestamp(),
|
||||
"end_time": await self.get_current_timestamp(),
|
||||
}
|
||||
)
|
||||
|
||||
# Handle file uploads only if storage_client is configured
|
||||
path = None
|
||||
if self.storage_client:
|
||||
content: Optional[Union[bytes, str]] = None
|
||||
|
||||
if element.path:
|
||||
async with aiofiles.open(element.path, "rb") as f:
|
||||
content = await f.read()
|
||||
elif element.content:
|
||||
content = element.content
|
||||
elif not element.url:
|
||||
raise ValueError("Element url, path or content must be provided")
|
||||
|
||||
if content is not None:
|
||||
if element.thread_id:
|
||||
path = f"threads/{element.thread_id}/files/{element.id}"
|
||||
else:
|
||||
path = f"files/{element.id}"
|
||||
|
||||
content_disposition = (
|
||||
f'attachment; filename="{element.name}"'
|
||||
if not (
|
||||
GCSStorageClient is not None
|
||||
and isinstance(self.storage_client, GCSStorageClient)
|
||||
)
|
||||
else None
|
||||
)
|
||||
await self.storage_client.upload_file(
|
||||
object_key=path,
|
||||
data=content,
|
||||
mime=element.mime or "application/octet-stream",
|
||||
overwrite=True,
|
||||
content_disposition=content_disposition,
|
||||
)
|
||||
|
||||
else:
|
||||
# Log warning only if element has file content that needs uploading
|
||||
if element.path or element.url or element.content:
|
||||
logger.warning(
|
||||
"Data Layer: No storage client configured. "
|
||||
"File will not be uploaded."
|
||||
)
|
||||
|
||||
# Always persist element metadata to database
|
||||
query = """
|
||||
INSERT INTO "Element" (
|
||||
id, "threadId", "stepId", metadata, mime, name, "objectKey", url,
|
||||
"chainlitKey", display, size, language, page, props
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14
|
||||
)
|
||||
ON CONFLICT (id) DO UPDATE SET
|
||||
props = EXCLUDED.props
|
||||
"""
|
||||
params = {
|
||||
"id": element.id,
|
||||
"thread_id": element.thread_id,
|
||||
"step_id": element.for_id,
|
||||
"metadata": json.dumps(
|
||||
{
|
||||
"size": element.size,
|
||||
"language": element.language,
|
||||
"display": element.display,
|
||||
"type": element.type,
|
||||
"page": getattr(element, "page", None),
|
||||
}
|
||||
),
|
||||
"mime": element.mime,
|
||||
"name": element.name,
|
||||
"object_key": path,
|
||||
"url": element.url,
|
||||
"chainlit_key": element.chainlit_key,
|
||||
"display": element.display,
|
||||
"size": element.size,
|
||||
"language": element.language,
|
||||
"page": getattr(element, "page", None),
|
||||
"props": json.dumps(getattr(element, "props", {})),
|
||||
}
|
||||
await self.execute_query(query, params)
|
||||
|
||||
async def get_element(
|
||||
self, thread_id: str, element_id: str
|
||||
) -> Optional[ElementDict]:
|
||||
query = """
|
||||
SELECT * FROM "Element"
|
||||
WHERE id = $1 AND "threadId" = $2
|
||||
"""
|
||||
results = await self.execute_query(
|
||||
query, {"element_id": element_id, "thread_id": thread_id}
|
||||
)
|
||||
|
||||
if not results:
|
||||
return None
|
||||
|
||||
row = results[0]
|
||||
metadata = json.loads(row.get("metadata", "{}"))
|
||||
|
||||
return ElementDict(
|
||||
id=str(row["id"]),
|
||||
threadId=str(row["threadId"]),
|
||||
type=metadata.get("type", "file"),
|
||||
url=str(row["url"]),
|
||||
name=str(row["name"]),
|
||||
mime=str(row["mime"]),
|
||||
objectKey=str(row["objectKey"]),
|
||||
forId=str(row["stepId"]),
|
||||
chainlitKey=row.get("chainlitKey"),
|
||||
display=row["display"],
|
||||
size=row["size"],
|
||||
language=row["language"],
|
||||
page=row["page"],
|
||||
autoPlay=row.get("autoPlay"),
|
||||
playerConfig=row.get("playerConfig"),
|
||||
props=json.loads(row.get("props", "{}")),
|
||||
)
|
||||
|
||||
@queue_until_user_message()
|
||||
async def delete_element(self, element_id: str, thread_id: Optional[str] = None):
|
||||
query = """
|
||||
SELECT * FROM "Element"
|
||||
WHERE id = $1
|
||||
"""
|
||||
elements = await self.execute_query(query, {"id": element_id})
|
||||
|
||||
if self.storage_client is not None and len(elements) > 0:
|
||||
if elements[0]["objectKey"]:
|
||||
await self.storage_client.delete_file(
|
||||
object_key=elements[0]["objectKey"]
|
||||
)
|
||||
query = """
|
||||
DELETE FROM "Element"
|
||||
WHERE id = $1
|
||||
"""
|
||||
params = {"id": element_id}
|
||||
|
||||
if thread_id:
|
||||
query += ' AND "threadId" = $2'
|
||||
params["thread_id"] = thread_id
|
||||
|
||||
await self.execute_query(query, params)
|
||||
|
||||
@queue_until_user_message()
|
||||
async def create_step(self, step_dict: StepDict):
|
||||
if step_dict.get("threadId"):
|
||||
thread_query = 'SELECT id FROM "Thread" WHERE id = $1'
|
||||
thread_results = await self.execute_query(
|
||||
thread_query, {"thread_id": step_dict["threadId"]}
|
||||
)
|
||||
if not thread_results:
|
||||
await self.update_thread(thread_id=step_dict["threadId"])
|
||||
|
||||
if step_dict.get("parentId"):
|
||||
parent_query = 'SELECT id FROM "Step" WHERE id = $1'
|
||||
parent_results = await self.execute_query(
|
||||
parent_query, {"parent_id": step_dict["parentId"]}
|
||||
)
|
||||
if not parent_results:
|
||||
await self.create_step(
|
||||
{
|
||||
"id": step_dict["parentId"],
|
||||
"metadata": {},
|
||||
"type": "run",
|
||||
"createdAt": step_dict.get("createdAt"),
|
||||
}
|
||||
)
|
||||
|
||||
query = """
|
||||
INSERT INTO "Step" (
|
||||
id, "threadId", "parentId", input, metadata, name, output,
|
||||
type, "startTime", "endTime", "showInput", "isError"
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12
|
||||
)
|
||||
ON CONFLICT (id) DO UPDATE SET
|
||||
"parentId" = COALESCE(EXCLUDED."parentId", "Step"."parentId"),
|
||||
input = COALESCE(NULLIF(EXCLUDED.input, ''), "Step".input),
|
||||
metadata = CASE
|
||||
WHEN EXCLUDED.metadata <> '{}' THEN EXCLUDED.metadata
|
||||
ELSE "Step".metadata
|
||||
END,
|
||||
name = COALESCE(EXCLUDED.name, "Step".name),
|
||||
output = COALESCE(NULLIF(EXCLUDED.output, ''), "Step".output),
|
||||
type = CASE
|
||||
WHEN EXCLUDED.type = 'run' THEN "Step".type
|
||||
ELSE EXCLUDED.type
|
||||
END,
|
||||
"threadId" = COALESCE(EXCLUDED."threadId", "Step"."threadId"),
|
||||
"endTime" = COALESCE(EXCLUDED."endTime", "Step"."endTime"),
|
||||
"startTime" = LEAST(EXCLUDED."startTime", "Step"."startTime"),
|
||||
"showInput" = COALESCE(EXCLUDED."showInput", "Step"."showInput"),
|
||||
"isError" = COALESCE(EXCLUDED."isError", "Step"."isError")
|
||||
"""
|
||||
|
||||
timestamp = await self.get_current_timestamp()
|
||||
created_at = step_dict.get("createdAt")
|
||||
if created_at:
|
||||
timestamp = datetime.strptime(created_at, ISO_FORMAT)
|
||||
|
||||
params = {
|
||||
"id": step_dict["id"],
|
||||
"thread_id": step_dict.get("threadId"),
|
||||
"parent_id": step_dict.get("parentId"),
|
||||
"input": step_dict.get("input"),
|
||||
"metadata": json.dumps(step_dict.get("metadata", {})),
|
||||
"name": step_dict.get("name"),
|
||||
"output": step_dict.get("output"),
|
||||
"type": step_dict["type"],
|
||||
"start_time": timestamp,
|
||||
"end_time": timestamp,
|
||||
"show_input": str(step_dict.get("showInput", "json")),
|
||||
"is_error": step_dict.get("isError", False),
|
||||
}
|
||||
await self.execute_query(query, params)
|
||||
|
||||
@queue_until_user_message()
|
||||
async def update_step(self, step_dict: StepDict):
|
||||
await self.create_step(step_dict)
|
||||
|
||||
@queue_until_user_message()
|
||||
async def delete_step(self, step_id: str):
|
||||
# Delete associated elements and feedbacks first
|
||||
await self.execute_query(
|
||||
'DELETE FROM "Element" WHERE "stepId" = $1', {"step_id": step_id}
|
||||
)
|
||||
await self.execute_query(
|
||||
'DELETE FROM "Feedback" WHERE "stepId" = $1', {"step_id": step_id}
|
||||
)
|
||||
# Delete the step
|
||||
await self.execute_query(
|
||||
'DELETE FROM "Step" WHERE id = $1', {"step_id": step_id}
|
||||
)
|
||||
|
||||
async def get_step(self, step_id: str) -> Optional[StepDict]:
|
||||
# Get step and related feedback
|
||||
query = """
|
||||
SELECT s.*,
|
||||
f.id feedback_id,
|
||||
f.value feedback_value,
|
||||
f."comment" feedback_comment
|
||||
FROM "Step" s left join "Feedback" f on s.id = f."stepId"
|
||||
WHERE s.id = $1
|
||||
"""
|
||||
result = await self.execute_query(query, {"step_id": step_id})
|
||||
if not result:
|
||||
return None
|
||||
return self._convert_step_row_to_dict(result[0])
|
||||
|
||||
async def get_thread_author(self, thread_id: str) -> str:
|
||||
query = """
|
||||
SELECT u.identifier
|
||||
FROM "Thread" t
|
||||
JOIN "User" u ON t."userId" = u.id
|
||||
WHERE t.id = $1
|
||||
"""
|
||||
results = await self.execute_query(query, {"thread_id": thread_id})
|
||||
if not results:
|
||||
raise ValueError(f"Thread {thread_id} not found")
|
||||
return results[0]["identifier"]
|
||||
|
||||
async def delete_thread(self, thread_id: str):
|
||||
elements_query = """
|
||||
SELECT * FROM "Element"
|
||||
WHERE "threadId" = $1
|
||||
"""
|
||||
elements_results = await self.execute_query(
|
||||
elements_query, {"thread_id": thread_id}
|
||||
)
|
||||
|
||||
if self.storage_client is not None:
|
||||
for elem in elements_results:
|
||||
if elem["objectKey"]:
|
||||
await self.storage_client.delete_file(object_key=elem["objectKey"])
|
||||
|
||||
await self.execute_query(
|
||||
'DELETE FROM "Thread" WHERE id = $1', {"thread_id": thread_id}
|
||||
)
|
||||
|
||||
async def list_threads(
|
||||
self, pagination: Pagination, filters: ThreadFilter
|
||||
) -> PaginatedResponse[ThreadDict]:
|
||||
query = """
|
||||
SELECT
|
||||
t.*,
|
||||
u.identifier as user_identifier,
|
||||
(SELECT COUNT(*) FROM "Thread" WHERE "userId" = t."userId") as total
|
||||
FROM "Thread" t
|
||||
LEFT JOIN "User" u ON t."userId" = u.id
|
||||
WHERE t."deletedAt" IS NULL
|
||||
"""
|
||||
params: Dict[str, Any] = {}
|
||||
param_count = 1
|
||||
|
||||
if filters.search:
|
||||
query += f" AND t.name ILIKE ${param_count}"
|
||||
params["name"] = f"%{filters.search}%"
|
||||
param_count += 1
|
||||
|
||||
if filters.userId:
|
||||
query += f' AND t."userId" = ${param_count}'
|
||||
params["user_id"] = filters.userId
|
||||
param_count += 1
|
||||
|
||||
if pagination.cursor:
|
||||
query += f' AND t."updatedAt" < (SELECT "updatedAt" FROM "Thread" WHERE id = ${param_count})'
|
||||
params["cursor"] = pagination.cursor
|
||||
param_count += 1
|
||||
|
||||
query += f' ORDER BY t."updatedAt" DESC LIMIT ${param_count}'
|
||||
params["limit"] = pagination.first + 1
|
||||
|
||||
results = await self.execute_query(query, params)
|
||||
threads = results
|
||||
|
||||
has_next_page = len(threads) > pagination.first
|
||||
if has_next_page:
|
||||
threads = threads[:-1]
|
||||
|
||||
thread_dicts = []
|
||||
for thread in threads:
|
||||
thread_dict = ThreadDict(
|
||||
id=str(thread["id"]),
|
||||
createdAt=thread["updatedAt"].isoformat(),
|
||||
name=thread["name"],
|
||||
userId=str(thread["userId"]) if thread["userId"] else None,
|
||||
userIdentifier=thread["user_identifier"],
|
||||
metadata=json.loads(thread["metadata"]),
|
||||
steps=[],
|
||||
elements=[],
|
||||
tags=[],
|
||||
)
|
||||
thread_dicts.append(thread_dict)
|
||||
|
||||
return PaginatedResponse(
|
||||
pageInfo=PageInfo(
|
||||
hasNextPage=has_next_page,
|
||||
startCursor=thread_dicts[0]["id"] if thread_dicts else None,
|
||||
endCursor=thread_dicts[-1]["id"] if thread_dicts else None,
|
||||
),
|
||||
data=thread_dicts,
|
||||
)
|
||||
|
||||
async def get_thread(self, thread_id: str) -> Optional[ThreadDict]:
|
||||
query = """
|
||||
SELECT t.*, u.identifier as user_identifier
|
||||
FROM "Thread" t
|
||||
LEFT JOIN "User" u ON t."userId" = u.id
|
||||
WHERE t.id = $1 AND t."deletedAt" IS NULL
|
||||
"""
|
||||
results = await self.execute_query(query, {"thread_id": thread_id})
|
||||
|
||||
if not results:
|
||||
return None
|
||||
|
||||
thread = results[0]
|
||||
|
||||
# Get steps and related feedback
|
||||
steps_query = """
|
||||
SELECT s.*,
|
||||
f.id feedback_id,
|
||||
f.value feedback_value,
|
||||
f."comment" feedback_comment
|
||||
FROM "Step" s left join "Feedback" f on s.id = f."stepId"
|
||||
WHERE s."threadId" = $1
|
||||
ORDER BY "startTime"
|
||||
"""
|
||||
steps_results = await self.execute_query(steps_query, {"thread_id": thread_id})
|
||||
|
||||
# Get elements
|
||||
elements_query = """
|
||||
SELECT * FROM "Element"
|
||||
WHERE "threadId" = $1
|
||||
"""
|
||||
elements_results = await self.execute_query(
|
||||
elements_query, {"thread_id": thread_id}
|
||||
)
|
||||
|
||||
if self.storage_client is not None:
|
||||
for elem in elements_results:
|
||||
if not elem["url"] and elem["objectKey"]:
|
||||
elem["url"] = await self.storage_client.get_read_url(
|
||||
object_key=elem["objectKey"],
|
||||
)
|
||||
|
||||
return ThreadDict(
|
||||
id=str(thread["id"]),
|
||||
createdAt=thread["createdAt"].isoformat(),
|
||||
name=thread["name"],
|
||||
userId=str(thread["userId"]) if thread["userId"] else None,
|
||||
userIdentifier=thread["user_identifier"],
|
||||
metadata=json.loads(thread["metadata"]),
|
||||
steps=[self._convert_step_row_to_dict(step) for step in steps_results],
|
||||
elements=[
|
||||
self._convert_element_row_to_dict(elem) for elem in elements_results
|
||||
],
|
||||
tags=[],
|
||||
)
|
||||
|
||||
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,
|
||||
):
|
||||
if self.show_logger:
|
||||
logger.info(f"asyncpg: update_thread, thread_id={thread_id}")
|
||||
|
||||
has_updates = (
|
||||
metadata is not None
|
||||
or name is not None
|
||||
or user_id is not None
|
||||
or tags is not None
|
||||
)
|
||||
|
||||
if metadata is None:
|
||||
metadata = {}
|
||||
|
||||
thread_name = truncate(
|
||||
name
|
||||
if name is not None
|
||||
else (metadata.get("name") if metadata and "name" in metadata else None)
|
||||
)
|
||||
|
||||
existing = await self.execute_query(
|
||||
'SELECT "metadata" FROM "Thread" WHERE id = $1',
|
||||
{"thread_id": thread_id},
|
||||
)
|
||||
|
||||
thread_exists = isinstance(existing, list) and existing
|
||||
if thread_exists and not has_updates:
|
||||
return
|
||||
|
||||
base = {}
|
||||
if thread_exists:
|
||||
raw = existing[0].get("metadata") or {}
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
base = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
base = {}
|
||||
elif isinstance(raw, dict):
|
||||
base = raw
|
||||
to_delete = {k for k, v in metadata.items() if v is None}
|
||||
incoming = {k: v for k, v in metadata.items() if v is not None}
|
||||
base = {k: v for k, v in base.items() if k not in to_delete}
|
||||
metadata = {**base, **incoming}
|
||||
|
||||
data = {
|
||||
"id": thread_id,
|
||||
"name": thread_name,
|
||||
"userId": user_id,
|
||||
"tags": tags,
|
||||
"metadata": json.dumps(metadata),
|
||||
"updatedAt": datetime.now(),
|
||||
}
|
||||
|
||||
# Remove None values
|
||||
data = {k: v for k, v in data.items() if v is not None}
|
||||
|
||||
# Build the query dynamically based on available fields
|
||||
columns = [f'"{k}"' for k in data.keys()]
|
||||
placeholders = [f"${i + 1}" for i in range(len(data))]
|
||||
values = list(data.values())
|
||||
|
||||
update_sets = [f'"{k}" = EXCLUDED."{k}"' for k in data.keys() if k != "id"]
|
||||
|
||||
if update_sets:
|
||||
query = f"""
|
||||
INSERT INTO "Thread" ({", ".join(columns)})
|
||||
VALUES ({", ".join(placeholders)})
|
||||
ON CONFLICT (id) DO UPDATE
|
||||
SET {", ".join(update_sets)};
|
||||
"""
|
||||
else:
|
||||
query = f"""
|
||||
INSERT INTO "Thread" ({", ".join(columns)})
|
||||
VALUES ({", ".join(placeholders)})
|
||||
ON CONFLICT (id) DO NOTHING
|
||||
"""
|
||||
|
||||
await self.execute_query(query, {str(i + 1): v for i, v in enumerate(values)})
|
||||
|
||||
async def get_favorite_steps(self, user_id: str) -> List[StepDict]:
|
||||
query = """
|
||||
SELECT s.*
|
||||
FROM "Step" s
|
||||
JOIN "Thread" t ON s."threadId" = t.id
|
||||
WHERE t."userId" = $1
|
||||
AND s.metadata::jsonb->>'favorite' = 'true'
|
||||
ORDER BY s."createdAt" DESC \
|
||||
"""
|
||||
results = await self.execute_query(query, {"user_id": user_id})
|
||||
return [self._convert_step_row_to_dict(row) for row in results]
|
||||
|
||||
def _extract_feedback_dict_from_step_row(self, row: Dict) -> Optional[FeedbackDict]:
|
||||
if row.get("feedback_id", None) is not None:
|
||||
return FeedbackDict(
|
||||
forId=str(row["id"]),
|
||||
id=str(row["feedback_id"]),
|
||||
value=row["feedback_value"],
|
||||
comment=row["feedback_comment"],
|
||||
)
|
||||
return None
|
||||
|
||||
def _convert_step_row_to_dict(self, row: Dict) -> StepDict:
|
||||
return StepDict(
|
||||
id=str(row["id"]),
|
||||
threadId=str(row["threadId"]) if row.get("threadId") else "",
|
||||
parentId=str(row["parentId"]) if row.get("parentId") else None,
|
||||
name=str(row.get("name")),
|
||||
type=row["type"],
|
||||
input=row.get("input", {}),
|
||||
output=row.get("output", {}),
|
||||
metadata=json.loads(row.get("metadata", "{}")),
|
||||
createdAt=row["createdAt"].isoformat() if row.get("createdAt") else None,
|
||||
start=row["startTime"].isoformat() if row.get("startTime") else None,
|
||||
showInput=row.get("showInput"),
|
||||
isError=row.get("isError"),
|
||||
end=row["endTime"].isoformat() if row.get("endTime") else None,
|
||||
feedback=self._extract_feedback_dict_from_step_row(row),
|
||||
)
|
||||
|
||||
def _convert_element_row_to_dict(self, row: Dict) -> ElementDict:
|
||||
metadata = json.loads(row.get("metadata", "{}"))
|
||||
return ElementDict(
|
||||
id=str(row["id"]),
|
||||
threadId=str(row["threadId"]) if row.get("threadId") else None,
|
||||
type=metadata.get("type", "file"),
|
||||
url=row["url"],
|
||||
name=row["name"],
|
||||
mime=row["mime"],
|
||||
objectKey=row["objectKey"],
|
||||
forId=str(row["stepId"]),
|
||||
chainlitKey=row.get("chainlitKey"),
|
||||
display=row["display"],
|
||||
size=row["size"],
|
||||
language=row["language"],
|
||||
page=row["page"],
|
||||
autoPlay=row.get("autoPlay"),
|
||||
playerConfig=row.get("playerConfig"),
|
||||
props=json.loads(row.get("props") or "{}"),
|
||||
)
|
||||
|
||||
async def build_debug_url(self) -> str:
|
||||
return ""
|
||||
|
||||
async def cleanup(self):
|
||||
"""Cleanup database connections"""
|
||||
if self.pool:
|
||||
logger.debug("Cleaning up connection pool")
|
||||
await self.pool.close()
|
||||
self.pool = None
|
||||
|
||||
async def close(self) -> None:
|
||||
if self.storage_client:
|
||||
await self.storage_client.close()
|
||||
await self.cleanup()
|
||||
|
||||
|
||||
def truncate(text: Optional[str], max_length: int = 255) -> Optional[str]:
|
||||
return None if text is None else text[:max_length]
|
||||
Reference in New Issue
Block a user