319 lines
9.0 KiB
Python
319 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
STREAMING_THROTTLE_MS = 160
|
|
SIGNIFICANT_DELTA_CHARS = 18
|
|
PRINT_FREQUENCY_MS = 50
|
|
INITIAL_PLACEHOLDER = "\u23f3 Thinking..."
|
|
|
|
def set_streaming_throttle_ms(value: int) -> None:
|
|
global STREAMING_THROTTLE_MS
|
|
STREAMING_THROTTLE_MS = value
|
|
|
|
def set_significant_delta_chars(value: int) -> None:
|
|
global SIGNIFICANT_DELTA_CHARS
|
|
SIGNIFICANT_DELTA_CHARS = value
|
|
|
|
def set_print_frequency_ms(value: int) -> None:
|
|
global PRINT_FREQUENCY_MS
|
|
PRINT_FREQUENCY_MS = value
|
|
|
|
def set_initial_placeholder(value: str) -> None:
|
|
global INITIAL_PLACEHOLDER
|
|
INITIAL_PLACEHOLDER = value
|
|
|
|
NATURAL_BOUNDARIES = frozenset({"\n", "\u3002", "!", "\uff01", "?", "\uff1f", ";", "\uff1b", ":", "\uff1a"})
|
|
|
|
BACKOFF_DURATION = 60
|
|
|
|
|
|
@dataclass
|
|
class CardKitState:
|
|
card_id: str = ""
|
|
message_id: str = ""
|
|
target_id: str = ""
|
|
reply_to_id: str | None = None
|
|
thread_id: str | None = None
|
|
version: int = 0
|
|
|
|
|
|
_creation_backoff: dict[str, float] = {}
|
|
|
|
|
|
def _get_backoff_key(app_id: str, target_id: str) -> str:
|
|
return f"{app_id}:{target_id}"
|
|
|
|
|
|
def _is_backoff_active(backoff_key: str) -> bool:
|
|
last_fail = _creation_backoff.get(backoff_key, 0)
|
|
return last_fail > 0 and time.monotonic() - last_fail < BACKOFF_DURATION
|
|
|
|
|
|
def _set_backoff(backoff_key: str) -> None:
|
|
_creation_backoff[backoff_key] = time.monotonic()
|
|
|
|
|
|
def _clear_backoff(backoff_key: str) -> None:
|
|
_creation_backoff.pop(backoff_key, None)
|
|
|
|
|
|
async def _api_request(client: Any, method: str, url_path: str, body: dict | None = None) -> dict:
|
|
import httpx
|
|
|
|
base_url = getattr(client, "_open_api_url", "https://open.feishu.cn")
|
|
|
|
async with httpx.AsyncClient(timeout=30.0) as http:
|
|
access_token = await _get_tenant_token(client)
|
|
headers = {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Content-Type": "application/json; charset=utf-8",
|
|
}
|
|
full_url = f"{base_url}/open-apis{url_path}"
|
|
|
|
if method == "POST":
|
|
resp = await http.post(full_url, headers=headers, json=body)
|
|
elif method == "PUT":
|
|
resp = await http.put(full_url, headers=headers, json=body)
|
|
elif method == "PATCH":
|
|
resp = await http.patch(full_url, headers=headers, json=body)
|
|
else:
|
|
resp = await http.get(full_url, headers=headers)
|
|
|
|
if resp.status_code >= 400:
|
|
logger.warning("Card Kit API %s %s failed: %d %s", method, url_path, resp.status_code, resp.text[:200])
|
|
return {"success": False, "error": resp.text, "code": resp.status_code}
|
|
|
|
data = resp.json()
|
|
if data.get("code", -1) != 0:
|
|
logger.warning("Card Kit API error: code=%d msg=%s", data.get("code"), data.get("msg", ""))
|
|
return {"success": False, "error": data.get("msg", ""), "code": data.get("code")}
|
|
|
|
return {"success": True, "data": data.get("data", {})}
|
|
|
|
|
|
async def _get_tenant_token(client: Any) -> str:
|
|
try:
|
|
token = await client.auth.tenant_access_token_internal()
|
|
if hasattr(token, "token"):
|
|
return token.token
|
|
return ""
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
async def create_card_entity(client: Any) -> dict:
|
|
return await _api_request(
|
|
client,
|
|
"POST",
|
|
"/cardkit/v1/cards",
|
|
{"body": {"elements": [{"tag": "markdown", "content": INITIAL_PLACEHOLDER}]}, "config": {"width_mode": "fill"}},
|
|
)
|
|
|
|
|
|
async def update_card_content(client: Any, card_id: str, elements: list[dict], version: int) -> dict:
|
|
return await _api_request(
|
|
client,
|
|
"PUT",
|
|
f"/cardkit/v1/cards/{card_id}/elements/content/content",
|
|
{"elements": elements, "version": version},
|
|
)
|
|
|
|
|
|
async def update_card_settings(client: Any, card_id: str, settings: dict) -> dict:
|
|
return await _api_request(
|
|
client,
|
|
"PATCH",
|
|
f"/cardkit/v1/cards/{card_id}/settings",
|
|
settings,
|
|
)
|
|
|
|
|
|
async def send_card_message(
|
|
client: Any,
|
|
target_id: str,
|
|
card_id: str,
|
|
*,
|
|
reply_to_id: str | None = None,
|
|
thread_id: str | None = None,
|
|
content: str = INITIAL_PLACEHOLDER,
|
|
) -> dict:
|
|
|
|
card_content = json.dumps(
|
|
{
|
|
"schema": "2.0",
|
|
"config": {"width_mode": "fill", "streaming_mode": "on"},
|
|
"card_id": card_id,
|
|
"body": {"elements": [{"tag": "markdown", "content": content}]},
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
body: dict[str, Any] = {
|
|
"receive_id": target_id,
|
|
"msg_type": "interactive",
|
|
"content": card_content,
|
|
}
|
|
if reply_to_id:
|
|
body["root_id"] = reply_to_id
|
|
|
|
import lark_oapi
|
|
|
|
req = lark_oapi.im.v1.CreateMessageRequest(
|
|
receive_id_type="chat_id",
|
|
request_body=lark_oapi.im.v1.CreateMessageRequestBody(**body),
|
|
)
|
|
resp = await client.im.v1.message.create_async(req)
|
|
|
|
if resp.success():
|
|
return {"success": True, "msg_id": getattr(resp.data, "message_id", "")}
|
|
else:
|
|
return {"success": False, "error": getattr(resp, "msg", ""), "code": getattr(resp, "code", 0)}
|
|
|
|
|
|
async def open_card_stream(
|
|
client: Any,
|
|
target_id: str,
|
|
*,
|
|
reply_to_id: str | None = None,
|
|
thread_id: str | None = None,
|
|
app_id: str = "",
|
|
) -> CardKitState | None:
|
|
|
|
backoff_key = _get_backoff_key(app_id, target_id)
|
|
if _is_backoff_active(backoff_key):
|
|
logger.debug("Card Kit backoff active for %s", backoff_key)
|
|
return None
|
|
|
|
try:
|
|
create_result = await create_card_entity(client)
|
|
if not create_result.get("success"):
|
|
_set_backoff(backoff_key)
|
|
return None
|
|
|
|
card_id = create_result["data"].get("card_id", "")
|
|
if not card_id:
|
|
_set_backoff(backoff_key)
|
|
return None
|
|
|
|
send_result = await send_card_message(
|
|
client, target_id, card_id, reply_to_id=reply_to_id, thread_id=thread_id
|
|
)
|
|
if not send_result.get("success"):
|
|
_set_backoff(backoff_key)
|
|
return None
|
|
|
|
_clear_backoff(backoff_key)
|
|
|
|
return CardKitState(
|
|
card_id=card_id,
|
|
message_id=send_result.get("msg_id", ""),
|
|
target_id=target_id,
|
|
reply_to_id=reply_to_id,
|
|
thread_id=thread_id,
|
|
version=0,
|
|
)
|
|
except Exception:
|
|
logger.exception("Card Kit open_card_stream failed")
|
|
_set_backoff(backoff_key)
|
|
return None
|
|
|
|
|
|
async def push_card_update(
|
|
client: Any,
|
|
state: CardKitState,
|
|
content: str,
|
|
*,
|
|
throttle_ms: int = STREAMING_THROTTLE_MS,
|
|
) -> bool:
|
|
|
|
version = state.version + 1
|
|
elements = [{"tag": "markdown", "content": content}]
|
|
|
|
await asyncio.sleep(throttle_ms / 1000.0)
|
|
|
|
result = await update_card_content(client, state.card_id, elements, version)
|
|
if result.get("success"):
|
|
state.version = version
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
async def close_card_stream(client: Any, state: CardKitState) -> bool:
|
|
|
|
result = await update_card_settings(client, state.card_id, {"streaming_mode": "off"})
|
|
return result.get("success", False)
|
|
|
|
|
|
async def send_card_fallback(
|
|
client: Any,
|
|
state: CardKitState,
|
|
content: str,
|
|
) -> dict:
|
|
|
|
card_json = {
|
|
"schema": "2.0",
|
|
"config": {"width_mode": "fill"},
|
|
"body": {"elements": [{"tag": "markdown", "content": content}]},
|
|
}
|
|
|
|
import lark_oapi
|
|
|
|
body: dict[str, Any] = {
|
|
"receive_id": state.target_id,
|
|
"msg_type": "interactive",
|
|
"content": json.dumps(card_json, ensure_ascii=False),
|
|
}
|
|
if state.reply_to_id:
|
|
body["root_id"] = state.reply_to_id
|
|
|
|
req = lark_oapi.im.v1.CreateMessageRequest(
|
|
receive_id_type="chat_id",
|
|
request_body=lark_oapi.im.v1.CreateMessageRequestBody(**body),
|
|
)
|
|
resp = await client.im.v1.message.create_async(req)
|
|
|
|
if resp.success():
|
|
return {"success": True, "msg_id": getattr(resp.data, "message_id", "")}
|
|
return {"success": False, "error": getattr(resp, "msg", ""), "code": getattr(resp, "code", 0)}
|
|
|
|
|
|
def should_send_update(new_content: str, previous_content: str) -> bool:
|
|
if not new_content:
|
|
return False
|
|
delta = len(new_content) - len(previous_content)
|
|
if delta >= SIGNIFICANT_DELTA_CHARS:
|
|
return True
|
|
if delta > 0 and new_content[-1] in NATURAL_BOUNDARIES:
|
|
return True
|
|
return False
|
|
|
|
|
|
def merge_content(previous: str, next_text: str) -> str:
|
|
if not previous:
|
|
return next_text
|
|
if not next_text:
|
|
return previous
|
|
if next_text.startswith(previous):
|
|
return next_text
|
|
if previous in next_text:
|
|
return next_text
|
|
|
|
overlap = 0
|
|
max_overlap = min(len(previous), len(next_text))
|
|
for i in range(min(max_overlap, 20), 0, -1):
|
|
if previous.endswith(next_text[:i]):
|
|
overlap = i
|
|
break
|
|
|
|
if overlap > 0:
|
|
return previous[: len(previous) - overlap] + next_text
|
|
return next_text |