mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-12 16:00:14 +08:00
Compare commits
15 Commits
copilot/cr
...
perf/webui
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0293434580 | ||
|
|
39131d2e12 | ||
|
|
735bd43648 | ||
|
|
6a42ad7934 | ||
|
|
b07dbb3d26 | ||
|
|
0b69034491 | ||
|
|
e286da75c4 | ||
|
|
7008a46158 | ||
|
|
5a90b56e45 | ||
|
|
2cb6c84eeb | ||
|
|
1a1d83d3be | ||
|
|
a748264fa4 | ||
|
|
2cda708eba | ||
|
|
7f897887fd | ||
|
|
40076b6aff |
@@ -1,7 +1,7 @@
|
||||
import json
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from ..message import Message, TextPart
|
||||
from ..message import AudioURLPart, ImageURLPart, Message, TextPart, ThinkPart
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
@@ -28,9 +28,19 @@ class TokenCounter(Protocol):
|
||||
...
|
||||
|
||||
|
||||
# 图片/音频 token 开销估算值,参考 OpenAI vision pricing:
|
||||
# low-res ~85 tokens, high-res ~170 per 512px tile, 通常几百到上千。
|
||||
# 这里取一个保守中位数,宁可偏高触发压缩也不要偏低导致 API 报错。
|
||||
IMAGE_TOKEN_ESTIMATE = 765
|
||||
AUDIO_TOKEN_ESTIMATE = 500
|
||||
|
||||
|
||||
class EstimateTokenCounter:
|
||||
"""Estimate token counter implementation.
|
||||
Provides a simple estimation of token count based on character types.
|
||||
|
||||
Supports multimodal content: images, audio, and thinking parts
|
||||
are all counted so that the context compressor can trigger in time.
|
||||
"""
|
||||
|
||||
def count_tokens(
|
||||
@@ -45,12 +55,16 @@ class EstimateTokenCounter:
|
||||
if isinstance(content, str):
|
||||
total += self._estimate_tokens(content)
|
||||
elif isinstance(content, list):
|
||||
# 处理多模态内容
|
||||
for part in content:
|
||||
if isinstance(part, TextPart):
|
||||
total += self._estimate_tokens(part.text)
|
||||
elif isinstance(part, ThinkPart):
|
||||
total += self._estimate_tokens(part.think)
|
||||
elif isinstance(part, ImageURLPart):
|
||||
total += IMAGE_TOKEN_ESTIMATE
|
||||
elif isinstance(part, AudioURLPart):
|
||||
total += AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
# 处理 Tool Calls
|
||||
if msg.tool_calls:
|
||||
for tc in msg.tool_calls:
|
||||
tc_str = json.dumps(tc if isinstance(tc, dict) else tc.model_dump())
|
||||
|
||||
@@ -12,14 +12,50 @@ class ContextTruncator:
|
||||
and len(message.tool_calls) > 0
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _split_system_rest(
|
||||
messages: list[Message],
|
||||
) -> tuple[list[Message], list[Message]]:
|
||||
"""Split messages into system messages and the rest.
|
||||
|
||||
Returns:
|
||||
tuple: (system_messages, non_system_messages)
|
||||
"""
|
||||
first_non_system = 0
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role != "system":
|
||||
first_non_system = i
|
||||
break
|
||||
return messages[:first_non_system], messages[first_non_system:]
|
||||
|
||||
@staticmethod
|
||||
def _ensure_user_message(
|
||||
system_messages: list[Message],
|
||||
truncated: list[Message],
|
||||
original_messages: list[Message],
|
||||
) -> list[Message]:
|
||||
"""Ensure the result always contains the first user message right after
|
||||
system messages. This is required by many LLM APIs (e.g. Zhipu) that
|
||||
mandate a ``user`` message immediately following the ``system`` message.
|
||||
"""
|
||||
if truncated and truncated[0].role == "user":
|
||||
return system_messages + truncated
|
||||
|
||||
# Locate the first user message from the *original* list.
|
||||
first_user = next((m for m in original_messages if m.role == "user"), None)
|
||||
if first_user is None:
|
||||
return system_messages + truncated
|
||||
|
||||
return system_messages + [first_user] + truncated
|
||||
|
||||
def fix_messages(self, messages: list[Message]) -> list[Message]:
|
||||
"""修复消息列表,确保 tool call 和 tool response 的配对关系有效。
|
||||
"""Fix the message list to ensure the validity of tool call and tool response pairing.
|
||||
|
||||
此方法确保:
|
||||
1. 每个 `tool` 消息前面都有一个包含 tool_calls 的 `assistant` 消息
|
||||
2. 每个包含 tool_calls 的 `assistant` 消息后面都有对应的 `tool` 响应
|
||||
This method ensures that:
|
||||
1. Each `tool` message is preceded by an `assistant` message containing `tool_calls`.
|
||||
2. Each `assistant` message containing `tool_calls` is followed by corresponding `
|
||||
|
||||
这是 OpenAI Chat Completions API 规范的要求(Gemini 对此执行严格检查)。
|
||||
This is a requirement of the OpenAI Chat Completions API specification (Gemini enforces this strictly).
|
||||
"""
|
||||
if not messages:
|
||||
return messages
|
||||
@@ -38,24 +74,25 @@ class ContextTruncator:
|
||||
|
||||
for msg in messages:
|
||||
if msg.role == "tool":
|
||||
# 只有在有挂起的 assistant(tool_calls) 时才记录 tool 响应
|
||||
# Only record tool responses when there is a pending assistant(tool_calls)
|
||||
if pending_assistant is not None:
|
||||
pending_tools.append(msg)
|
||||
# else: 孤立的 tool 消息,直接忽略
|
||||
# Isolated tool messages without a preceding assistant(tool_calls) are ignored
|
||||
continue
|
||||
|
||||
if self._has_tool_calls(msg):
|
||||
# 遇到新的 assistant(tool_calls) 前,先处理旧的 pending 链
|
||||
# When encountering a new assistant(tool_calls), first process the old pending chain
|
||||
flush_pending_if_valid()
|
||||
pending_assistant = msg
|
||||
continue
|
||||
|
||||
# 非 tool,且不含 tool_calls 的消息
|
||||
# 先结束任何 pending 链,再正常追加
|
||||
# Non-tool messages that do not contain tool_calls will break the pending chain.
|
||||
# Flush any pending chain first, then append the current message normally.
|
||||
flush_pending_if_valid()
|
||||
fixed_messages.append(msg)
|
||||
|
||||
# 结束时处理最后一个 pending 链
|
||||
# Flush the last pending chain at the end,
|
||||
# ensuring that any remaining valid assistant(tool_calls) and its tools are included in the final list.
|
||||
flush_pending_if_valid()
|
||||
|
||||
return fixed_messages
|
||||
@@ -66,29 +103,23 @@ class ContextTruncator:
|
||||
keep_most_recent_turns: int,
|
||||
drop_turns: int = 1,
|
||||
) -> list[Message]:
|
||||
"""截断上下文列表,确保不超过最大长度。
|
||||
一个 turn 包含一个 user 消息和一个 assistant 消息。
|
||||
这个方法会保证截断后的上下文列表符合 OpenAI 的上下文格式。
|
||||
"""
|
||||
Turn-based truncation strategy, which drops the oldest turns while keeping the most recent N turns.
|
||||
A turn consists of a user message and an assistant message.
|
||||
This method ensures that the truncated context list conforms to OpenAI's context format.
|
||||
|
||||
Args:
|
||||
messages: 上下文列表
|
||||
keep_most_recent_turns: 保留最近的对话轮数
|
||||
drop_turns: 一次性丢弃的对话轮数
|
||||
messages: The original list of messages in the context.
|
||||
keep_most_recent_turns: The number of most recent turns to keep. If set to -1, it means keeping all turns (no truncation).
|
||||
drop_turns: The number of turns to drop from the beginning.
|
||||
|
||||
Returns:
|
||||
截断后的上下文列表
|
||||
The truncated list of messages.
|
||||
"""
|
||||
if keep_most_recent_turns == -1:
|
||||
return messages
|
||||
|
||||
first_non_system = 0
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role != "system":
|
||||
first_non_system = i
|
||||
break
|
||||
|
||||
system_messages = messages[:first_non_system]
|
||||
non_system_messages = messages[first_non_system:]
|
||||
system_messages, non_system_messages = self._split_system_rest(messages)
|
||||
|
||||
if len(non_system_messages) // 2 <= keep_most_recent_turns:
|
||||
return messages
|
||||
@@ -99,7 +130,7 @@ class ContextTruncator:
|
||||
else:
|
||||
truncated_contexts = non_system_messages[-num_to_keep * 2 :]
|
||||
|
||||
# 找到第一个 role 为 user 的索引,确保上下文格式正确
|
||||
# Find the first user message
|
||||
index = next(
|
||||
(i for i, item in enumerate(truncated_contexts) if item.role == "user"),
|
||||
None,
|
||||
@@ -107,8 +138,9 @@ class ContextTruncator:
|
||||
if index is not None and index > 0:
|
||||
truncated_contexts = truncated_contexts[index:]
|
||||
|
||||
result = system_messages + truncated_contexts
|
||||
|
||||
result = self._ensure_user_message(
|
||||
system_messages, truncated_contexts, messages
|
||||
)
|
||||
return self.fix_messages(result)
|
||||
|
||||
def truncate_by_dropping_oldest_turns(
|
||||
@@ -116,53 +148,39 @@ class ContextTruncator:
|
||||
messages: list[Message],
|
||||
drop_turns: int = 1,
|
||||
) -> list[Message]:
|
||||
"""丢弃最旧的 N 个对话轮次。"""
|
||||
"""Drop the oldest N turns, regardless of the number of turns to keep."""
|
||||
if drop_turns <= 0:
|
||||
return messages
|
||||
|
||||
first_non_system = 0
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role != "system":
|
||||
first_non_system = i
|
||||
break
|
||||
|
||||
system_messages = messages[:first_non_system]
|
||||
non_system_messages = messages[first_non_system:]
|
||||
system_messages, non_system_messages = self._split_system_rest(messages)
|
||||
|
||||
if len(non_system_messages) // 2 <= drop_turns:
|
||||
truncated_non_system = []
|
||||
else:
|
||||
truncated_non_system = non_system_messages[drop_turns * 2 :]
|
||||
|
||||
# Find the first user message
|
||||
index = next(
|
||||
(i for i, item in enumerate(truncated_non_system) if item.role == "user"),
|
||||
None,
|
||||
)
|
||||
if index is not None:
|
||||
truncated_non_system = truncated_non_system[index:]
|
||||
elif truncated_non_system:
|
||||
truncated_non_system = []
|
||||
|
||||
result = system_messages + truncated_non_system
|
||||
|
||||
result = self._ensure_user_message(
|
||||
system_messages, truncated_non_system, messages
|
||||
)
|
||||
return self.fix_messages(result)
|
||||
|
||||
def truncate_by_halving(
|
||||
self,
|
||||
messages: list[Message],
|
||||
) -> list[Message]:
|
||||
"""对半砍策略,删除 50% 的消息"""
|
||||
"""Halve the number of messages, keeping the most recent ones."""
|
||||
if len(messages) <= 2:
|
||||
return messages
|
||||
|
||||
first_non_system = 0
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role != "system":
|
||||
first_non_system = i
|
||||
break
|
||||
|
||||
system_messages = messages[:first_non_system]
|
||||
non_system_messages = messages[first_non_system:]
|
||||
system_messages, non_system_messages = self._split_system_rest(messages)
|
||||
|
||||
messages_to_delete = len(non_system_messages) // 2
|
||||
if messages_to_delete == 0:
|
||||
@@ -170,6 +188,7 @@ class ContextTruncator:
|
||||
|
||||
truncated_non_system = non_system_messages[messages_to_delete:]
|
||||
|
||||
# Find the first user message
|
||||
index = next(
|
||||
(i for i, item in enumerate(truncated_non_system) if item.role == "user"),
|
||||
None,
|
||||
@@ -177,6 +196,7 @@ class ContextTruncator:
|
||||
if index is not None:
|
||||
truncated_non_system = truncated_non_system[index:]
|
||||
|
||||
result = system_messages + truncated_non_system
|
||||
|
||||
result = self._ensure_user_message(
|
||||
system_messages, truncated_non_system, messages
|
||||
)
|
||||
return self.fix_messages(result)
|
||||
|
||||
@@ -1112,6 +1112,20 @@ CONFIG_METADATA_2 = {
|
||||
"api_base": "https://api.anthropic.com/v1",
|
||||
"timeout": 120,
|
||||
"proxy": "",
|
||||
"custom_headers": {},
|
||||
"anth_thinking_config": {"type": "", "budget": 0, "effort": ""},
|
||||
},
|
||||
"Kimi Coding Plan": {
|
||||
"id": "kimi-code",
|
||||
"provider": "kimi-code",
|
||||
"type": "kimi_code_chat_completion",
|
||||
"provider_type": "chat_completion",
|
||||
"enable": True,
|
||||
"key": [],
|
||||
"api_base": "https://api.kimi.com/coding/",
|
||||
"timeout": 120,
|
||||
"proxy": "",
|
||||
"custom_headers": {"User-Agent": "claude-code/0.1.0"},
|
||||
"anth_thinking_config": {"type": "", "budget": 0, "effort": ""},
|
||||
},
|
||||
"Moonshot": {
|
||||
@@ -3439,7 +3453,7 @@ CONFIG_METADATA_3 = {
|
||||
"description": "白名单 ID 列表",
|
||||
"type": "list",
|
||||
"items": {"type": "string"},
|
||||
"hint": "使用 /sid 获取 ID。",
|
||||
"hint": "使用 /sid 获取 ID。当白名单列表为空时,代表不启用白名单(即所有 ID 都在白名单内)。",
|
||||
},
|
||||
"platform_settings.id_whitelist_log": {
|
||||
"description": "输出日志",
|
||||
|
||||
@@ -172,6 +172,9 @@ def try_capture_follow_up(event: AstrMessageEvent) -> FollowUpCapture | None:
|
||||
if not active_sender_id or active_sender_id != sender_id:
|
||||
return None
|
||||
|
||||
if runner_event.get_extra("agent_stop_requested"):
|
||||
return None
|
||||
|
||||
ticket = runner.follow_up(message_text=_event_follow_up_text(event))
|
||||
if not ticket:
|
||||
return None
|
||||
|
||||
@@ -86,7 +86,7 @@ class PipelineScheduler:
|
||||
try:
|
||||
await self._process_stages(event)
|
||||
|
||||
# 如果没有发送操作, 则发送一个空消息, 以便于后续的处理
|
||||
# 发送一个空消息, 以便于后续的处理
|
||||
if isinstance(event, WebChatMessageEvent | WecomAIBotMessageEvent):
|
||||
await event.send(None)
|
||||
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Coroutine
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -92,32 +91,37 @@ class WebChatAdapter(Platform):
|
||||
active_request_ids = self._webchat_queue_mgr.list_back_request_ids(
|
||||
conversation_id
|
||||
)
|
||||
subscription_request_ids = [
|
||||
req_id for req_id in active_request_ids if req_id.startswith("ws_sub_")
|
||||
stream_request_ids = [
|
||||
req_id for req_id in active_request_ids if not req_id.startswith("ws_sub_")
|
||||
]
|
||||
target_request_ids = subscription_request_ids or active_request_ids
|
||||
target_request_ids = stream_request_ids or active_request_ids
|
||||
|
||||
if target_request_ids:
|
||||
for request_id in target_request_ids:
|
||||
await WebChatMessageEvent._send(
|
||||
request_id,
|
||||
message_chain,
|
||||
session.session_id,
|
||||
if not target_request_ids:
|
||||
# No active streams to consume this proactive message.
|
||||
# Persist directly and return to avoid creating an unused queue.
|
||||
try:
|
||||
await self._save_proactive_message(conversation_id, message_chain)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WebChatAdapter] Failed to save proactive message: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
else:
|
||||
message_id = f"active_{uuid.uuid4()!s}"
|
||||
await super().send_by_session(session, message_chain)
|
||||
return
|
||||
|
||||
for request_id in target_request_ids:
|
||||
await WebChatMessageEvent._send(
|
||||
message_id,
|
||||
request_id,
|
||||
message_chain,
|
||||
session.session_id,
|
||||
streaming=True,
|
||||
emit_complete=True,
|
||||
)
|
||||
|
||||
should_persist = (
|
||||
bool(subscription_request_ids)
|
||||
or not active_request_ids
|
||||
or all(req_id.startswith("active_") for req_id in active_request_ids)
|
||||
)
|
||||
if should_persist:
|
||||
# If only passive subscription queues exist for this conversation,
|
||||
# keep a proactive save as a fallback since they are not tied to
|
||||
# the normal streaming persistence path.
|
||||
if not stream_request_ids:
|
||||
try:
|
||||
await self._save_proactive_message(conversation_id, message_chain)
|
||||
except Exception as e:
|
||||
|
||||
@@ -34,6 +34,7 @@ class WebChatMessageEvent(AstrMessageEvent):
|
||||
message: MessageChain | None,
|
||||
session_id: str,
|
||||
streaming: bool = False,
|
||||
emit_complete: bool = False,
|
||||
) -> str | None:
|
||||
request_id = str(message_id)
|
||||
conversation_id = _extract_conversation_id(session_id)
|
||||
@@ -127,6 +128,17 @@ class WebChatMessageEvent(AstrMessageEvent):
|
||||
else:
|
||||
logger.debug(f"webchat 忽略: {comp.type}")
|
||||
|
||||
if emit_complete:
|
||||
await web_chat_back_queue.put(
|
||||
{
|
||||
"type": "complete",
|
||||
"data": data,
|
||||
"streaming": streaming,
|
||||
"chain_type": message.type,
|
||||
"message_id": message_id,
|
||||
},
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
async def send(self, message: MessageChain | None) -> None:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""企业微信智能机器人事件处理模块,处理消息事件的发送和接收"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from astrbot.api import logger
|
||||
@@ -14,6 +15,8 @@ from .wecomai_webhook import WecomAIBotWebhookClient
|
||||
class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
"""企业微信智能机器人消息事件"""
|
||||
|
||||
STREAM_FLUSH_INTERVAL = 0.5
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message_str: str,
|
||||
@@ -135,6 +138,8 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
|
||||
async def send(self, message: MessageChain | None) -> None:
|
||||
"""发送消息"""
|
||||
if message is None:
|
||||
return
|
||||
raw = self.message_obj.raw_message
|
||||
assert isinstance(raw, dict), (
|
||||
"wecom_ai_bot platform event raw_message should be a dict"
|
||||
@@ -240,6 +245,7 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
return
|
||||
|
||||
increment_plain = ""
|
||||
last_stream_update_time = 0.0
|
||||
async for chain in generator:
|
||||
if self.webhook_client:
|
||||
await self.webhook_client.send_message_chain(
|
||||
@@ -251,17 +257,20 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
chunk_text = self._extract_plain_text_from_chain(chain)
|
||||
if chunk_text:
|
||||
increment_plain += chunk_text
|
||||
await self.long_connection_sender(
|
||||
req_id,
|
||||
{
|
||||
"msgtype": "stream",
|
||||
"stream": {
|
||||
"id": stream_id,
|
||||
"finish": False,
|
||||
"content": increment_plain,
|
||||
now = asyncio.get_running_loop().time()
|
||||
if now - last_stream_update_time >= self.STREAM_FLUSH_INTERVAL:
|
||||
await self.long_connection_sender(
|
||||
req_id,
|
||||
{
|
||||
"msgtype": "stream",
|
||||
"stream": {
|
||||
"id": stream_id,
|
||||
"finish": False,
|
||||
"content": increment_plain,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
last_stream_update_time = now
|
||||
|
||||
await self.long_connection_sender(
|
||||
req_id,
|
||||
@@ -287,22 +296,31 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
await super().send_streaming(generator, use_fallback)
|
||||
return
|
||||
|
||||
# 企业微信智能机器人不支持增量发送,因此我们需要在这里将增量内容累积起来,积累发送
|
||||
# 企业微信智能机器人不支持增量发送,因此我们需要在这里将增量内容累积起来,按间隔推送
|
||||
increment_plain = ""
|
||||
last_stream_update_time = 0.0
|
||||
|
||||
async def enqueue_stream_plain(text: str) -> None:
|
||||
if not text:
|
||||
return
|
||||
await back_queue.put(
|
||||
{
|
||||
"type": "plain",
|
||||
"data": text,
|
||||
"streaming": True,
|
||||
"session_id": stream_id,
|
||||
},
|
||||
)
|
||||
|
||||
async for chain in generator:
|
||||
if self.webhook_client:
|
||||
await self.webhook_client.send_message_chain(
|
||||
chain, unsupported_only=True
|
||||
)
|
||||
# 累积增量内容,并改写 Plain 段
|
||||
chain.squash_plain()
|
||||
for comp in chain.chain:
|
||||
if isinstance(comp, Plain):
|
||||
comp.text = increment_plain + comp.text
|
||||
increment_plain = comp.text
|
||||
break
|
||||
|
||||
if chain.type == "break" and final_data:
|
||||
if increment_plain:
|
||||
await enqueue_stream_plain(increment_plain)
|
||||
# 分割符
|
||||
await back_queue.put(
|
||||
{
|
||||
@@ -313,15 +331,30 @@ class WecomAIBotMessageEvent(AstrMessageEvent):
|
||||
},
|
||||
)
|
||||
final_data = ""
|
||||
increment_plain = ""
|
||||
continue
|
||||
|
||||
final_data += await WecomAIBotMessageEvent._send(
|
||||
chain,
|
||||
stream_id=stream_id,
|
||||
queue_mgr=self.queue_mgr,
|
||||
streaming=True,
|
||||
suppress_unsupported_log=self.webhook_client is not None,
|
||||
)
|
||||
chunk_text = self._extract_plain_text_from_chain(chain)
|
||||
if chunk_text:
|
||||
increment_plain += chunk_text
|
||||
final_data += chunk_text
|
||||
now = asyncio.get_running_loop().time()
|
||||
if now - last_stream_update_time >= self.STREAM_FLUSH_INTERVAL:
|
||||
await enqueue_stream_plain(increment_plain)
|
||||
last_stream_update_time = now
|
||||
|
||||
for comp in chain.chain:
|
||||
if isinstance(comp, (At, Plain)):
|
||||
continue
|
||||
await WecomAIBotMessageEvent._send(
|
||||
MessageChain([comp]),
|
||||
stream_id=stream_id,
|
||||
queue_mgr=self.queue_mgr,
|
||||
streaming=True,
|
||||
suppress_unsupported_log=self.webhook_client is not None,
|
||||
)
|
||||
|
||||
await enqueue_stream_plain(increment_plain)
|
||||
|
||||
await back_queue.put(
|
||||
{
|
||||
|
||||
@@ -375,6 +375,10 @@ class ProviderManager:
|
||||
from .sources.anthropic_source import (
|
||||
ProviderAnthropic as ProviderAnthropic,
|
||||
)
|
||||
case "kimi_code_chat_completion":
|
||||
from .sources.kimi_code_source import (
|
||||
ProviderKimiCode as ProviderKimiCode,
|
||||
)
|
||||
case "googlegenai_chat_completion":
|
||||
from .sources.gemini_source import (
|
||||
ProviderGoogleGenAI as ProviderGoogleGenAI,
|
||||
|
||||
@@ -16,7 +16,6 @@ from astrbot.core.provider.entities import LLMResponse, TokenUsage
|
||||
from astrbot.core.provider.func_tool_manager import ToolSet
|
||||
from astrbot.core.utils.io import download_image_by_url
|
||||
from astrbot.core.utils.network_utils import (
|
||||
create_proxy_client,
|
||||
is_connection_error,
|
||||
log_connection_failure,
|
||||
)
|
||||
@@ -29,6 +28,30 @@ from ..register import register_provider_adapter
|
||||
"Anthropic Claude API 提供商适配器",
|
||||
)
|
||||
class ProviderAnthropic(Provider):
|
||||
@staticmethod
|
||||
def _normalize_custom_headers(provider_config: dict) -> dict[str, str] | None:
|
||||
custom_headers = provider_config.get("custom_headers", {})
|
||||
if not isinstance(custom_headers, dict) or not custom_headers:
|
||||
return None
|
||||
normalized_headers: dict[str, str] = {}
|
||||
for key, value in custom_headers.items():
|
||||
normalized_headers[str(key)] = str(value)
|
||||
return normalized_headers or None
|
||||
|
||||
@classmethod
|
||||
def _resolve_custom_headers(
|
||||
cls,
|
||||
provider_config: dict,
|
||||
*,
|
||||
required_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, str] | None:
|
||||
merged_headers = cls._normalize_custom_headers(provider_config) or {}
|
||||
if required_headers:
|
||||
for header_name, header_value in required_headers.items():
|
||||
if not merged_headers.get(header_name, "").strip():
|
||||
merged_headers[header_name] = header_value
|
||||
return merged_headers or None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_config,
|
||||
@@ -46,6 +69,7 @@ class ProviderAnthropic(Provider):
|
||||
if isinstance(self.timeout, str):
|
||||
self.timeout = int(self.timeout)
|
||||
self.thinking_config = provider_config.get("anth_thinking_config", {})
|
||||
self.custom_headers = self._resolve_custom_headers(provider_config)
|
||||
|
||||
if use_api_key:
|
||||
self._init_api_key(provider_config)
|
||||
@@ -66,7 +90,12 @@ class ProviderAnthropic(Provider):
|
||||
def _create_http_client(self, provider_config: dict) -> httpx.AsyncClient | None:
|
||||
"""创建带代理的 HTTP 客户端"""
|
||||
proxy = provider_config.get("proxy", "")
|
||||
return create_proxy_client("Anthropic", proxy)
|
||||
if proxy:
|
||||
logger.info(f"[Anthropic] 使用代理: {proxy}")
|
||||
return httpx.AsyncClient(proxy=proxy, headers=self.custom_headers)
|
||||
if self.custom_headers:
|
||||
return httpx.AsyncClient(headers=self.custom_headers)
|
||||
return None
|
||||
|
||||
def _apply_thinking_config(self, payloads: dict) -> None:
|
||||
thinking_type = self.thinking_config.get("type", "")
|
||||
|
||||
27
astrbot/core/provider/sources/kimi_code_source.py
Normal file
27
astrbot/core/provider/sources/kimi_code_source.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from ..register import register_provider_adapter
|
||||
from .anthropic_source import ProviderAnthropic
|
||||
|
||||
KIMI_CODE_API_BASE = "https://api.kimi.com/coding"
|
||||
KIMI_CODE_DEFAULT_MODEL = "kimi-for-coding"
|
||||
KIMI_CODE_USER_AGENT = "claude-code/0.1.0"
|
||||
|
||||
|
||||
@register_provider_adapter(
|
||||
"kimi_code_chat_completion",
|
||||
"Kimi Code Provider Adapter",
|
||||
)
|
||||
class ProviderKimiCode(ProviderAnthropic):
|
||||
def __init__(
|
||||
self,
|
||||
provider_config: dict,
|
||||
provider_settings: dict,
|
||||
) -> None:
|
||||
merged_provider_config = dict(provider_config)
|
||||
merged_provider_config.setdefault("api_base", KIMI_CODE_API_BASE)
|
||||
merged_provider_config.setdefault("model", KIMI_CODE_DEFAULT_MODEL)
|
||||
merged_provider_config["custom_headers"] = self._resolve_custom_headers(
|
||||
merged_provider_config,
|
||||
required_headers={"User-Agent": KIMI_CODE_USER_AGENT},
|
||||
)
|
||||
|
||||
super().__init__(merged_provider_config, provider_settings)
|
||||
@@ -19,17 +19,15 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
|
||||
self.provider_config = provider_config
|
||||
self.provider_settings = provider_settings
|
||||
proxy = provider_config.get("proxy", "")
|
||||
provider_id = provider_config.get("id", "unknown_id")
|
||||
http_client = None
|
||||
if proxy:
|
||||
logger.info(f"[OpenAI Embedding] 使用代理: {proxy}")
|
||||
logger.info(f"[OpenAI Embedding] {provider_id} Using proxy: {proxy}")
|
||||
http_client = httpx.AsyncClient(proxy=proxy)
|
||||
api_base = provider_config.get("embedding_api_base", "").strip()
|
||||
if not api_base:
|
||||
api_base = "https://api.openai.com/v1"
|
||||
else:
|
||||
api_base = api_base.removesuffix("/")
|
||||
if not api_base.endswith("/v1"):
|
||||
api_base = f"{api_base}/v1"
|
||||
api_base = provider_config.get(
|
||||
"embedding_api_base", "https://api.openai.com/v1"
|
||||
).strip()
|
||||
logger.info(f"[OpenAI Embedding] {provider_id} Using API Base: {api_base}")
|
||||
self.client = AsyncOpenAI(
|
||||
api_key=provider_config.get("embedding_api_key"),
|
||||
base_url=api_base,
|
||||
|
||||
@@ -313,7 +313,8 @@ class ProviderOpenAIOfficial(Provider):
|
||||
logger.warning("Saving chunk state error: " + str(e))
|
||||
if not chunk.choices:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
# logger.debug(f"chunk delta: {delta}")
|
||||
# handle the content delta
|
||||
reasoning = self._extract_reasoning_content(chunk)
|
||||
@@ -331,6 +332,11 @@ class ProviderOpenAIOfficial(Provider):
|
||||
_y = True
|
||||
if chunk.usage:
|
||||
llm_response.usage = self._extract_usage(chunk.usage)
|
||||
elif choice_usage := getattr(choice, "usage", None):
|
||||
# Workaround for some providers that only return usage in choices[].usage, e.g. MoonshotAI
|
||||
# See https://github.com/AstrBotDevs/AstrBot/issues/6614
|
||||
llm_response.usage = self._extract_usage(choice_usage)
|
||||
state.current_completion_snapshot.usage = choice_usage
|
||||
if _y:
|
||||
yield llm_response
|
||||
|
||||
@@ -359,13 +365,11 @@ class ProviderOpenAIOfficial(Provider):
|
||||
reasoning_text = str(reasoning_attr)
|
||||
return reasoning_text
|
||||
|
||||
def _extract_usage(self, usage: CompletionUsage) -> TokenUsage:
|
||||
ptd = usage.prompt_tokens_details
|
||||
cached = ptd.cached_tokens if ptd and ptd.cached_tokens else 0
|
||||
prompt_tokens = 0 if usage.prompt_tokens is None else usage.prompt_tokens
|
||||
completion_tokens = (
|
||||
0 if usage.completion_tokens is None else usage.completion_tokens
|
||||
)
|
||||
def _extract_usage(self, usage: CompletionUsage | dict) -> TokenUsage:
|
||||
ptd = getattr(usage, "prompt_tokens_details", None)
|
||||
cached = getattr(ptd, "cached_tokens", 0) if ptd else 0
|
||||
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
||||
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
|
||||
return TokenUsage(
|
||||
input_other=prompt_tokens - cached,
|
||||
input_cached=cached,
|
||||
|
||||
@@ -5,12 +5,14 @@ import importlib.metadata as importlib_metadata
|
||||
import importlib.util
|
||||
import io
|
||||
import logging
|
||||
import ntpath
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import sys
|
||||
import threading
|
||||
from collections import deque
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -30,6 +32,9 @@ logger = logging.getLogger("astrbot")
|
||||
|
||||
_DISTLIB_FINDER_PATCH_ATTEMPTED = False
|
||||
_SITE_PACKAGES_IMPORT_LOCK = threading.RLock()
|
||||
_PIP_IN_PROCESS_ENV_LOCK = threading.RLock()
|
||||
_WINDOWS_UNC_PATH_PREFIXES = ("\\\\?\\UNC\\", "\\??\\UNC\\")
|
||||
_WINDOWS_EXTENDED_PATH_PREFIXES = ("\\\\?\\", "\\??\\")
|
||||
_PIP_FAILURE_PATTERNS = {
|
||||
"error_prefix": re.compile(r"^\s*error:", re.IGNORECASE),
|
||||
"user_requested": re.compile(r"\bthe user requested\b", re.IGNORECASE),
|
||||
@@ -235,6 +240,120 @@ def _run_pip_main_streaming(pip_main, args: list[str]) -> tuple[int, list[str]]:
|
||||
return result_code, stream.lines
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _temporary_environ(updates: Mapping[str, str]):
|
||||
if not updates:
|
||||
yield
|
||||
return
|
||||
|
||||
missing = object()
|
||||
previous_values = {key: os.environ.get(key, missing) for key in updates}
|
||||
|
||||
try:
|
||||
os.environ.update(updates)
|
||||
yield
|
||||
finally:
|
||||
for key, previous_value in previous_values.items():
|
||||
if previous_value is missing:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
assert isinstance(previous_value, str)
|
||||
os.environ[key] = previous_value
|
||||
|
||||
|
||||
def _run_pip_main_with_temporary_environ(
|
||||
pip_main,
|
||||
args: list[str],
|
||||
) -> tuple[int, list[str]]:
|
||||
# os.environ is process-wide; serialize reading current INCLUDE/LIB values
|
||||
# together with the temporary mutation window around the in-process pip
|
||||
# invocation.
|
||||
with _PIP_IN_PROCESS_ENV_LOCK:
|
||||
env_updates = _build_packaged_windows_runtime_build_env(base_env=os.environ)
|
||||
if not env_updates:
|
||||
return _run_pip_main_streaming(pip_main, args)
|
||||
|
||||
with _temporary_environ(env_updates):
|
||||
return _run_pip_main_streaming(pip_main, args)
|
||||
|
||||
|
||||
def _normalize_windows_native_build_path(path: str) -> str:
|
||||
"""Normalize a Windows path returned by native APIs or sys.executable.
|
||||
|
||||
Extended UNC prefixes are converted back to the standard ``\\server`` form,
|
||||
other extended prefixes are stripped, and the remaining path is normalized.
|
||||
"""
|
||||
normalized = path.replace("/", "\\")
|
||||
|
||||
# Extended UNC: \\?\UNC\server\share\... -> \\server\share\...
|
||||
for prefix in _WINDOWS_UNC_PATH_PREFIXES:
|
||||
if normalized.startswith(prefix):
|
||||
return ntpath.normpath(f"\\\\{normalized[len(prefix) :]}")
|
||||
|
||||
# Other extended prefixes are stripped before normalizing the path.
|
||||
for prefix in _WINDOWS_EXTENDED_PATH_PREFIXES:
|
||||
if normalized.startswith(prefix):
|
||||
normalized = normalized[len(prefix) :]
|
||||
break
|
||||
|
||||
return ntpath.normpath(normalized)
|
||||
|
||||
|
||||
def _get_case_insensitive_env_value(
|
||||
env: Mapping[str, str],
|
||||
upper_to_key: Mapping[str, str],
|
||||
name: str,
|
||||
) -> str | None:
|
||||
direct = env.get(name)
|
||||
if direct is not None:
|
||||
return direct
|
||||
|
||||
existing_key = upper_to_key.get(name.upper())
|
||||
if existing_key is not None:
|
||||
return env.get(existing_key)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _build_packaged_windows_runtime_build_env(
|
||||
*,
|
||||
base_env: Mapping[str, str] | None = None,
|
||||
) -> dict[str, str]:
|
||||
if sys.platform != "win32" or not is_packaged_desktop_runtime():
|
||||
return {}
|
||||
|
||||
base_env = os.environ if base_env is None else base_env
|
||||
|
||||
runtime_executable = _normalize_windows_native_build_path(sys.executable)
|
||||
runtime_dir = ntpath.dirname(runtime_executable)
|
||||
if not runtime_dir:
|
||||
return {}
|
||||
|
||||
include_dir = _normalize_windows_native_build_path(
|
||||
ntpath.join(runtime_dir, "include")
|
||||
)
|
||||
libs_dir = _normalize_windows_native_build_path(ntpath.join(runtime_dir, "libs"))
|
||||
include_exists = os.path.isdir(include_dir)
|
||||
libs_exists = os.path.isdir(libs_dir)
|
||||
|
||||
if not (include_exists or libs_exists):
|
||||
return {}
|
||||
|
||||
upper_to_key = {key.upper(): key for key in base_env}
|
||||
env_updates: dict[str, str] = {}
|
||||
|
||||
if include_exists:
|
||||
existing = _get_case_insensitive_env_value(base_env, upper_to_key, "INCLUDE")
|
||||
env_updates["INCLUDE"] = (
|
||||
f"{include_dir};{existing}" if existing else include_dir
|
||||
)
|
||||
if libs_exists:
|
||||
existing = _get_case_insensitive_env_value(base_env, upper_to_key, "LIB")
|
||||
env_updates["LIB"] = f"{libs_dir};{existing}" if existing else libs_dir
|
||||
|
||||
return env_updates
|
||||
|
||||
|
||||
def _matches_pip_failure_pattern(line: str, *pattern_names: str) -> bool:
|
||||
names = pattern_names or tuple(_PIP_FAILURE_PATTERNS)
|
||||
return any(_PIP_FAILURE_PATTERNS[name].search(line) for name in names)
|
||||
@@ -931,7 +1050,9 @@ class PipInstaller:
|
||||
original_handlers = list(logging.getLogger().handlers)
|
||||
try:
|
||||
result_code, output_lines = await asyncio.to_thread(
|
||||
_run_pip_main_streaming, pip_main, args
|
||||
_run_pip_main_with_temporary_environ,
|
||||
pip_main,
|
||||
args,
|
||||
)
|
||||
finally:
|
||||
_cleanup_added_root_handlers(original_handlers)
|
||||
|
||||
@@ -455,7 +455,10 @@ class OpenApiRoute(Route):
|
||||
if msg_type == "end":
|
||||
break
|
||||
if (streaming and msg_type == "complete") or not streaming:
|
||||
if chain_type in ("tool_call", "tool_call_result"):
|
||||
if chain_type in (
|
||||
"tool_call",
|
||||
"tool_call_result",
|
||||
):
|
||||
continue
|
||||
try:
|
||||
refs = self.chat_route._extract_web_search_refs(
|
||||
|
||||
@@ -4,7 +4,10 @@ version: '3.8'
|
||||
|
||||
services:
|
||||
astrbot:
|
||||
image: soulter/astrbot:latest
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
image: astrbot:kimi-code
|
||||
container_name: astrbot
|
||||
restart: always
|
||||
ports: # mappings description: https://github.com/AstrBotDevs/AstrBot/issues/497
|
||||
|
||||
@@ -14,19 +14,23 @@
|
||||
</div>
|
||||
|
||||
<!-- 选择对话框 -->
|
||||
<v-dialog v-model="dialog" max-width="1000px" min-width="800px">
|
||||
<v-dialog
|
||||
v-model="dialog"
|
||||
:max-width="isCompactLayout ? '96vw' : '1000px'"
|
||||
:min-width="isCompactLayout ? undefined : '800px'"
|
||||
>
|
||||
<v-card class="selector-dialog-card">
|
||||
<v-card-title class="dialog-title d-flex align-center py-4 px-5">
|
||||
<v-card-title class="dialog-title d-flex align-center" :class="isCompactLayout ? 'py-3 px-4' : 'py-4 px-5'">
|
||||
<v-icon class="mr-3" color="primary">mdi-account-circle</v-icon>
|
||||
<span>{{ labels.dialogTitle || '选择项目' }}</span>
|
||||
</v-card-title>
|
||||
|
||||
<v-divider />
|
||||
|
||||
<v-card-text class="pa-0" style="height: 600px; max-height: 80vh; overflow: hidden;">
|
||||
<v-card-text class="pa-0 selector-content">
|
||||
<div class="selector-layout">
|
||||
<!-- 左侧文件夹树 -->
|
||||
<div class="folder-sidebar">
|
||||
<div v-if="!isCompactLayout" class="folder-sidebar">
|
||||
<div class="sidebar-header pa-3 pb-2">
|
||||
<span class="text-caption text-medium-emphasis font-weight-medium">
|
||||
<v-icon size="small" class="mr-1">mdi-folder-multiple</v-icon>
|
||||
@@ -58,6 +62,20 @@
|
||||
|
||||
<!-- 右侧项目列表 -->
|
||||
<div class="items-panel">
|
||||
<div v-if="isCompactLayout" class="mobile-folder-bar px-4 py-2">
|
||||
<v-btn icon="mdi-arrow-left" size="small" variant="text"
|
||||
:disabled="currentFolderId === null" @click="navigateToParentFolder" />
|
||||
<v-btn size="small" variant="tonal" color="primary" prepend-icon="mdi-home"
|
||||
@click="navigateToFolder(null)">
|
||||
{{ labels.rootFolder || '根目录' }}
|
||||
</v-btn>
|
||||
<span class="text-caption text-medium-emphasis text-truncate mobile-folder-label">
|
||||
{{ currentFolderLabel }}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<v-divider v-if="isCompactLayout" />
|
||||
|
||||
<!-- 面包屑导航 -->
|
||||
<div class="breadcrumb-bar px-4 py-3">
|
||||
<v-breadcrumbs :items="breadcrumbItems" density="compact" class="pa-0">
|
||||
@@ -245,6 +263,18 @@ export default defineComponent({
|
||||
};
|
||||
},
|
||||
computed: {
|
||||
isCompactLayout(): boolean {
|
||||
return this.$vuetify.display.smAndDown;
|
||||
},
|
||||
|
||||
currentFolderLabel(): string {
|
||||
if (this.currentFolderId === null) {
|
||||
return this.labels.rootFolder || '根目录';
|
||||
}
|
||||
const currentFolder = this.breadcrumbPath[this.breadcrumbPath.length - 1];
|
||||
return currentFolder?.name || this.labels.rootFolder || '根目录';
|
||||
},
|
||||
|
||||
displayValue(): string {
|
||||
if (this.displayValueFormatter) {
|
||||
return this.displayValueFormatter(this.modelValue);
|
||||
@@ -332,6 +362,20 @@ export default defineComponent({
|
||||
this.$emit('navigate', folderId);
|
||||
},
|
||||
|
||||
navigateToParentFolder() {
|
||||
if (this.currentFolderId === null) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.breadcrumbPath.length <= 1) {
|
||||
this.navigateToFolder(null);
|
||||
return;
|
||||
}
|
||||
|
||||
const parent = this.breadcrumbPath[this.breadcrumbPath.length - 2];
|
||||
this.navigateToFolder(parent?.folder_id ?? null);
|
||||
},
|
||||
|
||||
findFolderInTree(folderId: string): FolderTreeNode | null {
|
||||
const findNode = (nodes: FolderTreeNode[]): FolderTreeNode | null => {
|
||||
for (const node of nodes) {
|
||||
@@ -414,6 +458,13 @@ export default defineComponent({
|
||||
.selector-layout {
|
||||
display: flex;
|
||||
height: 100%;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.selector-content {
|
||||
height: 600px;
|
||||
max-height: 80vh;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.folder-sidebar {
|
||||
@@ -450,6 +501,18 @@ export default defineComponent({
|
||||
|
||||
.items-content {
|
||||
background-color: transparent;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.mobile-folder-bar {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
.mobile-folder-label {
|
||||
min-width: 0;
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
.tree-list {
|
||||
@@ -517,22 +580,28 @@ export default defineComponent({
|
||||
background-color: rgba(var(--v-theme-primary), 0.08);
|
||||
}
|
||||
|
||||
@media (max-width: 600px) {
|
||||
@media (max-width: 960px) {
|
||||
.selector-layout {
|
||||
flex-direction: column;
|
||||
height: auto;
|
||||
max-height: 500px;
|
||||
max-height: none;
|
||||
}
|
||||
|
||||
.folder-sidebar {
|
||||
width: 100%;
|
||||
border-right: none;
|
||||
border-bottom: 1px solid rgba(var(--v-border-color), var(--v-border-opacity));
|
||||
max-height: 150px;
|
||||
.selector-content {
|
||||
max-height: 76vh;
|
||||
}
|
||||
|
||||
.items-list {
|
||||
max-height: 300px;
|
||||
min-height: 0;
|
||||
}
|
||||
|
||||
.breadcrumb-bar {
|
||||
overflow-x: auto;
|
||||
}
|
||||
|
||||
.breadcrumb-bar :deep(.v-breadcrumbs) {
|
||||
flex-wrap: nowrap;
|
||||
min-width: max-content;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
@@ -91,6 +91,15 @@ type WsStreamContext = {
|
||||
|
||||
const STREAMING_STORAGE_KEY = 'enableStreaming';
|
||||
const TRANSPORT_MODE_STORAGE_KEY = 'chatTransportMode';
|
||||
const HIDDEN_TOOL_CALL_NAMES = new Set(['send_message_to_user']);
|
||||
|
||||
function isHiddenToolCall(toolCall: ToolCall | { name?: unknown } | null | undefined): boolean {
|
||||
if (!toolCall || typeof toolCall !== 'object') {
|
||||
return false;
|
||||
}
|
||||
const name = toolCall.name;
|
||||
return typeof name === 'string' && HIDDEN_TOOL_CALL_NAMES.has(name);
|
||||
}
|
||||
|
||||
export function useMessages(
|
||||
currSessionId: Ref<string>,
|
||||
@@ -489,6 +498,9 @@ export function useMessages(
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
if (isHiddenToolCall(toolCallData)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const toolCall: ToolCall = {
|
||||
id: toolCallData.id,
|
||||
@@ -528,6 +540,9 @@ export function useMessages(
|
||||
} catch {
|
||||
return;
|
||||
}
|
||||
if (isHiddenToolCall(resultData)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageObj) {
|
||||
for (const part of messageObj.message) {
|
||||
@@ -658,7 +673,18 @@ export function useMessages(
|
||||
|
||||
// 如果 message 是数组 (新格式),遍历并填充 embedded 字段
|
||||
if (Array.isArray(message)) {
|
||||
const filteredMessage: MessagePart[] = [];
|
||||
for (const part of message as MessagePart[]) {
|
||||
if (part.type === 'tool_call' && Array.isArray(part.tool_calls)) {
|
||||
const visibleToolCalls = part.tool_calls.filter(
|
||||
(toolCall) => !isHiddenToolCall(toolCall),
|
||||
);
|
||||
if (!visibleToolCalls.length) {
|
||||
continue;
|
||||
}
|
||||
part.tool_calls = visibleToolCalls;
|
||||
}
|
||||
|
||||
if (part.type === 'image' && part.attachment_id) {
|
||||
part.embedded_url = await getAttachment(part.attachment_id);
|
||||
} else if (part.type === 'record' && part.attachment_id) {
|
||||
@@ -671,7 +697,9 @@ export function useMessages(
|
||||
};
|
||||
}
|
||||
// plain, reply, tool_call, video 保持原样
|
||||
filteredMessage.push(part);
|
||||
}
|
||||
content.message = filteredMessage;
|
||||
}
|
||||
|
||||
// 处理 agent_stats (snake_case -> camelCase)
|
||||
|
||||
@@ -697,11 +697,11 @@
|
||||
"platform_settings": {
|
||||
"enable_id_white_list": {
|
||||
"description": "Enable Whitelist",
|
||||
"hint": "When enabled, only sessions in the whitelist will be responded to."
|
||||
"hint": "When enabled, only sessions in the whitelist will be responded to. If the whitelist is empty, the whitelist is disabled and all IDs are allowed."
|
||||
},
|
||||
"id_whitelist": {
|
||||
"description": "Whitelist ID List",
|
||||
"hint": "Use /sid to get IDs."
|
||||
"hint": "Use /sid to get IDs. If the list is empty, it means whitelist is disabled (all IDs are in the whitelist)."
|
||||
},
|
||||
"id_whitelist_log": {
|
||||
"description": "Output Logs",
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
},
|
||||
"description": {
|
||||
"openai": "Supports all OpenAI API compatible providers.",
|
||||
"kimi_code": "Dedicated Kimi CodingPlan / Kimi Code integration using a custom Anthropic API compatibility layer.",
|
||||
"default": ""
|
||||
}
|
||||
},
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
},
|
||||
"description": {
|
||||
"openai": "Поддерживаются все провайдеры, совместимые с OpenAI API.",
|
||||
"kimi_code": "Специальная интеграция Kimi CodingPlan / Kimi Code с отдельной совместимой адаптацией Anthropic API.",
|
||||
"vllm_rerank": "Также поддерживает Jina AI, Cohere, PPIO и другие.",
|
||||
"default": "Преобразование речи в текст"
|
||||
}
|
||||
@@ -148,4 +149,4 @@
|
||||
"modelId": "ID модели"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -699,11 +699,11 @@
|
||||
"platform_settings": {
|
||||
"enable_id_white_list": {
|
||||
"description": "启用白名单",
|
||||
"hint": "启用后,只有在白名单内的会话会被响应。"
|
||||
"hint": "启用后,只有在白名单内的会话会被响应,白名单列表为空时代表不启用白名单(所有 ID 都会被放行)。"
|
||||
},
|
||||
"id_whitelist": {
|
||||
"description": "白名单 ID 列表",
|
||||
"hint": "使用 /sid 获取 ID。"
|
||||
"hint": "使用 /sid 获取 ID。列表为空时表示该白名单不启用(即所有 ID 都在白名单内)。"
|
||||
},
|
||||
"id_whitelist_log": {
|
||||
"description": "输出日志",
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
},
|
||||
"description": {
|
||||
"openai": "也支持所有兼容 OpenAI API 的模型提供商。",
|
||||
"kimi_code": "Kimi 的 CodingPlan / Kimi Code 专用接入,使用特殊的 Anthropic API 兼容适配。",
|
||||
"vllm_rerank": "也支持 Jina AI, Cohere, PPIO 等提供商。",
|
||||
"default": ""
|
||||
}
|
||||
|
||||
@@ -21,6 +21,8 @@ export function getProviderIcon(type) {
|
||||
'nvidia': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/nvidia-color.svg',
|
||||
'siliconflow': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/siliconcloud.svg',
|
||||
'moonshot': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/kimi.svg',
|
||||
'kimi': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/kimi.svg',
|
||||
'kimi-code': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/kimi.svg',
|
||||
'ppio': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/ppio.svg',
|
||||
'dify': 'https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@latest/icons/dify-color.svg',
|
||||
"coze": "https://cdn.jsdelivr.net/npm/@lobehub/icons-static-svg@1.66.0/icons/coze.svg",
|
||||
@@ -50,9 +52,11 @@ export function getProviderIcon(type) {
|
||||
* @returns {string} 提供商描述
|
||||
*/
|
||||
export function getProviderDescription(template, name, tm) {
|
||||
if (name == 'OpenAI') {
|
||||
if (name === 'OpenAI') {
|
||||
return tm('providers.description.openai', { type: template.type });
|
||||
} else if (name == 'vLLM Rerank') {
|
||||
} else if (template.provider === 'kimi-code') {
|
||||
return tm('providers.description.kimi_code');
|
||||
} else if (name === 'vLLM Rerank') {
|
||||
return tm('providers.description.vllm_rerank', { type: template.type });
|
||||
}
|
||||
return tm('providers.description.default', { type: template.type });
|
||||
|
||||
@@ -9,68 +9,381 @@ Starting from version `v4.12.0`, AstrBot introduced the Agent sandbox environmen
|
||||
|
||||
## Enabling the Sandbox Environment
|
||||
|
||||
Currently, the sandbox environment only supports running through Docker. We are currently using the [Shipyard](https://github.com/AstrBotDevs/shipyard) project as AstrBot's sandbox environment driver. In the future, we will support more types of sandbox environment drivers, such as e2b.
|
||||
AstrBot currently supports the following sandbox drivers:
|
||||
|
||||
- `Shipyard Neo` (recommended)
|
||||
- `Shipyard` (legacy option, still supported)
|
||||
|
||||
In the current AstrBot console, go to **AI Settings** -> **Agent Computer Use** and select:
|
||||
|
||||
- `Computer Use Runtime` = `sandbox`
|
||||
- `Sandbox Driver` = `Shipyard Neo` or `Shipyard`
|
||||
|
||||
`Shipyard Neo` is now the default driver. It consists of Bay, Ship, and Gull:
|
||||
|
||||
- **Bay**: the control-plane API responsible for creating and managing sandboxes
|
||||
- **Ship**: provides Python / Shell / filesystem capabilities
|
||||
- **Gull**: provides browser automation capabilities
|
||||
|
||||
For `Shipyard Neo`, the workspace root is fixed at `/workspace`. When using filesystem tools in AstrBot, you should pass **paths relative to the workspace root**, for example `reports/result.txt`, not `/workspace/reports/result.txt`.
|
||||
|
||||
> [!TIP]
|
||||
> Browser capability is not available in every `Shipyard Neo` profile. AstrBot only mounts browser-related tools when the selected profile supports the `browser` capability. A typical example is `browser-python`.
|
||||
|
||||
## Performance Requirements
|
||||
|
||||
AstrBot limits each sandbox instance to at most 1 CPU and 512 MB of memory.
|
||||
|
||||
We recommend that your host machine have at least 2 CPUs, 4 GB of memory, and swap enabled, so multiple sandbox instances can run more reliably.
|
||||
|
||||
## Recommended: Use Shipyard Neo
|
||||
|
||||
### Deploy Shipyard Neo Separately (Recommended)
|
||||
|
||||
If you plan to use `Shipyard Neo` for the long term, it is generally better to **deploy it separately on a machine with more resources**, such as your homelab, a LAN server, or a dedicated cloud host, and then let AstrBot connect to Bay remotely.
|
||||
|
||||
The reason is that `Shipyard Neo` can become fairly resource-heavy when browser capability is enabled, because it needs to run a full browser runtime. On resource-constrained cloud servers, deploying AstrBot and `Shipyard Neo` on the same machine usually puts significant pressure on CPU and memory, which can negatively affect both stability and overall experience.
|
||||
|
||||
A basic deployment flow looks like this:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/AstrBotDevs/shipyard-neo
|
||||
cd shipyard-neo/deploy/docker
|
||||
# Modify the key settings in config.yaml, such as security.api_key
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
After deployment:
|
||||
|
||||
- Bay listens on `http://<your-host>:8114` by default
|
||||
- In the AstrBot console, choose the `Shipyard Neo` driver
|
||||
- Set `Shipyard Neo API Endpoint` to the corresponding address, for example `http://<your-host>:8114`
|
||||
- Set `Shipyard Neo Access Token` to the Bay API key; if AstrBot can access Bay's `credentials.json`, you may also leave it empty and let AstrBot auto-discover it
|
||||
|
||||
### Reference: Full `config.yaml` Example (with Notes)
|
||||
|
||||
If you want to customize the deployment parameters of `Shipyard Neo`, you can refer to the complete example below, adapted from [`deploy/docker/config.yaml`](https://github.com/AstrBotDevs/shipyard-neo/blob/main/deploy/docker/config.yaml). It keeps the default structure and adds explanatory notes to make each option easier to understand.
|
||||
|
||||
> [!TIP]
|
||||
> The minimum required change is `security.api_key`. If you are not sure what the other options do, it is usually best to keep the defaults first and only adjust profiles, resource limits, and warm pool settings as needed.
|
||||
|
||||
```yaml
|
||||
# Bay Production Config - Docker Compose (container_network mode)
|
||||
#
|
||||
# Bay runs inside Docker and communicates with Ship/Gull containers
|
||||
# through a shared Docker network.
|
||||
# In this mode, sandbox containers do not need to expose ports to the host.
|
||||
#
|
||||
# At minimum, update:
|
||||
# 1. security.api_key — set a strong random secret
|
||||
|
||||
server:
|
||||
# Bay API listen address
|
||||
host: "0.0.0.0"
|
||||
# Bay API listen port
|
||||
port: 8114
|
||||
|
||||
database:
|
||||
# SQLite is the default for single-node deployment.
|
||||
# For multi-instance / HA deployments, you can switch to PostgreSQL, for example:
|
||||
# url: "postgresql+asyncpg://user:pass@db-host:5432/bay"
|
||||
url: "sqlite+aiosqlite:///./data/bay.db"
|
||||
echo: false
|
||||
|
||||
driver:
|
||||
# Docker is the default driver
|
||||
type: docker
|
||||
|
||||
# Whether to pull images when creating new sandboxes.
|
||||
# In production, always is usually recommended so you get the latest images.
|
||||
image_pull_policy: always
|
||||
|
||||
docker:
|
||||
# Docker Socket endpoint
|
||||
socket: "unix:///var/run/docker.sock"
|
||||
|
||||
# When Bay, Ship, and Gull all run in containers,
|
||||
# container_network is recommended for direct container-network communication.
|
||||
connect_mode: container_network
|
||||
|
||||
# Shared network name; must match the network in docker-compose.yaml
|
||||
network: "bay-network"
|
||||
|
||||
# Whether to expose sandbox container ports to the host.
|
||||
# Disabling this is generally recommended in production.
|
||||
publish_ports: false
|
||||
host_port: null
|
||||
|
||||
cargo:
|
||||
# Cargo storage root path on the Bay side
|
||||
root_path: "/var/lib/bay/cargos"
|
||||
# Default workspace size limit (MB)
|
||||
default_size_limit_mb: 1024
|
||||
# Path mounted inside the sandbox. This is AstrBot/Neo's workspace root.
|
||||
mount_path: "/workspace"
|
||||
|
||||
security:
|
||||
# Required: set a strong random secret, for example openssl rand -hex 32
|
||||
api_key: "CHANGE-ME"
|
||||
# Whether anonymous access is allowed. false is recommended for production.
|
||||
allow_anonymous: false
|
||||
|
||||
# Proxy environment variable injection for containers.
|
||||
# When enabled, Bay injects HTTP(S)_PROXY and NO_PROXY into sandbox containers.
|
||||
proxy:
|
||||
enabled: false
|
||||
# http_proxy: "http://proxy.example.com:7890"
|
||||
# https_proxy: "http://proxy.example.com:7890"
|
||||
# no_proxy: "my-internal.service"
|
||||
|
||||
# Warm Pool: keep standby sandboxes pre-warmed to reduce cold-start latency.
|
||||
# When a user creates a sandbox, Bay will first try to claim a pre-warmed instance.
|
||||
warm_pool:
|
||||
enabled: true
|
||||
# Number of warmup queue workers
|
||||
warmup_queue_workers: 2
|
||||
# Maximum warmup queue size
|
||||
warmup_queue_max_size: 256
|
||||
# Policy when the queue is full
|
||||
warmup_queue_drop_policy: "drop_newest"
|
||||
# Useful threshold for operational alerts
|
||||
warmup_queue_drop_alert_threshold: 50
|
||||
# Warm pool maintenance interval (seconds)
|
||||
interval_seconds: 30
|
||||
# Whether to start warm-pool maintenance when Bay starts
|
||||
run_on_startup: true
|
||||
|
||||
profiles:
|
||||
# ── Standard Python sandbox ────────────────────────
|
||||
- id: python-default
|
||||
description: "Standard Python sandbox with filesystem and shell access"
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "1g"
|
||||
capabilities:
|
||||
- filesystem # includes upload/download
|
||||
- shell
|
||||
- python
|
||||
# Idle timeout (seconds)
|
||||
idle_timeout: 1800
|
||||
# Keep 1 warm instance ready
|
||||
warm_pool_size: 1
|
||||
env: {}
|
||||
# Optional profile-level proxy override
|
||||
# proxy:
|
||||
# enabled: false
|
||||
|
||||
# ── Data-science sandbox (more resources) ──────────
|
||||
- id: python-data
|
||||
description: "Data science sandbox with extra CPU and memory"
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 2.0
|
||||
memory: "4g"
|
||||
capabilities:
|
||||
- filesystem # includes upload/download
|
||||
- shell
|
||||
- python
|
||||
idle_timeout: 1800
|
||||
warm_pool_size: 1
|
||||
env: {}
|
||||
|
||||
# ── Browser + Python multi-container sandbox ───────
|
||||
- id: browser-python
|
||||
description: "Browser automation with Python backend"
|
||||
containers:
|
||||
- name: ship
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "1g"
|
||||
capabilities:
|
||||
- python
|
||||
- shell
|
||||
- filesystem # includes upload/download
|
||||
# These capabilities are primarily handled by the ship container
|
||||
primary_for:
|
||||
- filesystem
|
||||
- python
|
||||
- shell
|
||||
env: {}
|
||||
- name: browser
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-gull:latest"
|
||||
runtime_type: gull
|
||||
runtime_port: 8115
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "2g"
|
||||
capabilities:
|
||||
- browser
|
||||
env: {}
|
||||
idle_timeout: 1800
|
||||
warm_pool_size: 1
|
||||
|
||||
gc:
|
||||
# Automatic GC is recommended in production
|
||||
enabled: true
|
||||
run_on_startup: true
|
||||
# GC interval (seconds)
|
||||
interval_seconds: 300
|
||||
|
||||
# Must be unique in multi-instance deployments
|
||||
instance_id: "bay-prod"
|
||||
|
||||
idle_session:
|
||||
enabled: true
|
||||
expired_sandbox:
|
||||
enabled: true
|
||||
orphan_cargo:
|
||||
enabled: true
|
||||
orphan_container:
|
||||
# Recommended in production to clean up leaked containers
|
||||
enabled: true
|
||||
```
|
||||
|
||||
A practical way to think about this file:
|
||||
|
||||
- **Minimum required change**: `security.api_key`
|
||||
- **Most commonly adjusted options**: resource limits, `warm_pool_size`, and `idle_timeout` under `profiles`
|
||||
- **If you need browser capability**: use or customize the `browser-python` profile
|
||||
- **If you want to reduce cold-start time**: keep `warm_pool.enabled: true` and increase `warm_pool_size` for frequently used profiles
|
||||
- **If resources are limited**: reduce `warm_pool_size`, or even disable `warm_pool`
|
||||
- **If outbound proxy access is needed**: configure the top-level `proxy`, or override it per profile
|
||||
|
||||
### About Shipyard Neo Reuse and Persistence
|
||||
|
||||
`Shipyard Neo` has several important concepts:
|
||||
|
||||
- **Sandbox**: the stable, externally visible resource unit
|
||||
- **Session**: the actual running container session, which may be stopped or rebuilt
|
||||
- **Cargo**: the persistent workspace volume mounted at `/workspace`
|
||||
|
||||
From AstrBot's perspective, the current implementation caches the sandbox booter by request `session_id`; in the default main-agent flow, this `session_id` usually equals the message-session identifier `unified_msg_origin`. As a result, follow-up requests from the same message session will usually continue using the same Neo sandbox; if the sandbox becomes unavailable, it will be rebuilt automatically.
|
||||
|
||||
For more detailed explanations of TTL and persistence behavior, see the later sections on “`Shipyard Neo Sandbox TTL`” and “Data Persistence in the Sandbox Environment”.
|
||||
|
||||
## Legacy Option: Shipyard
|
||||
|
||||
The following content describes the older `Shipyard` driver. It is kept for compatibility with existing legacy deployments.
|
||||
|
||||
### Deploying AstrBot and Shipyard with Docker Compose
|
||||
|
||||
If you haven't deployed AstrBot yet, or want to switch to our recommended deployment method with sandbox environment, we recommend using Docker Compose to deploy AstrBot with the following code:
|
||||
If you have not deployed AstrBot yet, or want to use the older recommended deployment method with sandbox support, you can still deploy AstrBot with Docker Compose using the following commands:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/AstrBotDevs/AstrBot
|
||||
cd AstrBot
|
||||
# Modify the environment variable configuration in the compose-with-shipyard.yml file, such as Shipyard's access token, etc.
|
||||
# Modify the environment variables in compose-with-shipyard.yml, such as the Shipyard access token
|
||||
docker compose -f compose-with-shipyard.yml up -d
|
||||
docker pull soulter/shipyard-ship:latest
|
||||
```
|
||||
|
||||
This will start a Docker Compose service that includes AstrBot main program and the sandbox environment.
|
||||
This starts a Docker Compose stack containing the AstrBot main program and the sandbox environment.
|
||||
|
||||
### Deploying Shipyard Separately
|
||||
|
||||
If you have already deployed AstrBot but haven't deployed the sandbox environment, you can deploy Shipyard separately.
|
||||
|
||||
Code as follows:
|
||||
If AstrBot is already deployed but the sandbox environment is not, you can deploy Shipyard separately.
|
||||
|
||||
```bash
|
||||
mkdir astrbot-shipyard
|
||||
cd astrbot-shipyard
|
||||
wget https://raw.githubusercontent.com/AstrBotDevs/shipyard/refs/heads/main/pkgs/bay/docker-compose.yml -O docker-compose.yml
|
||||
# Modify the environment variable configuration in the compose-with-shipyard.yml file, such as Shipyard's access token, etc.
|
||||
# Modify the environment variables in docker-compose.yml, such as the Shipyard access token
|
||||
docker compose -f docker-compose.yml up -d
|
||||
# pull the latest Shipyard ship image
|
||||
docker pull soulter/shipyard-ship:latest
|
||||
```
|
||||
|
||||
After successful deployment, the above command will start a Shipyard service that listens on `http://<your-host>:8156` by default.
|
||||
After successful deployment, Shipyard listens on `http://<your-host>:8156` by default.
|
||||
|
||||
> [!TIP]
|
||||
> If you deploy AstrBot using Docker, you can also modify the Compose file above to place Shipyard's network in the same Docker network as AstrBot, so you don't need to expose Shipyard's port to the host machine.
|
||||
> If you deploy AstrBot with Docker, you can also place Shipyard on the same Docker network as AstrBot so you do not need to expose Shipyard's port to the host.
|
||||
|
||||
## Configuring AstrBot to Use the Sandbox Environment
|
||||
|
||||
In the AstrBot console, go to the "Configuration Files" page and find "Agent Sandbox Environment", then enable the sandbox environment switch.
|
||||
> [!TIP]
|
||||
> Please make sure your AstrBot version is `v4.12.0` or later.
|
||||
|
||||
In the configuration options that appear:
|
||||
In the AstrBot console, go to **AI Settings** -> **Agent Computer Use**.
|
||||
|
||||
For `Shipyard API Endpoint`, if you use the Docker Compose deployment method above, fill in `http://shipyard:8156`. If you deployed Shipyard separately, please fill in the corresponding address, such as `http://<your-host>:8156`.
|
||||
1. Set `Computer Use Runtime` to `sandbox`
|
||||
2. Select `Shipyard Neo` or `Shipyard` as the sandbox driver
|
||||
3. Fill in the corresponding configuration values for the selected driver
|
||||
4. Click **Save**
|
||||
|
||||
For `Shipyard Access Token`, please fill in the access token you configured when deploying Shipyard.
|
||||
### Configuring Shipyard Neo
|
||||
|
||||
For `Shipyard Ship Lifetime (seconds)`, this defines the lifetime of each sandbox environment instance, with a default value of 3600 seconds (1 hour). You can adjust this value as needed.
|
||||
If you choose `Shipyard Neo`, the main configuration items are:
|
||||
|
||||
For `Shipyard Ship Session Reuse Limit`, this defines the maximum number of sessions that each sandbox environment instance can reuse, with a default value of 10. This means that 10 sessions will share the same sandbox environment instance. You can adjust this value as needed.
|
||||
- `Shipyard Neo API Endpoint`
|
||||
- For a separated deployment, use the actual address, such as `http://<your-host>:8114`
|
||||
- `Shipyard Neo Access Token`
|
||||
- Fill in the Bay API key
|
||||
- If AstrBot can access Bay's `credentials.json`, you may leave it empty and let AstrBot auto-discover it
|
||||
- `Shipyard Neo Profile`
|
||||
- For example `python-default` or `browser-python`
|
||||
- If not explicitly specified, AstrBot will try to choose a profile with richer capabilities, preferring one that includes the `browser` capability, and fall back to `python-default` if needed
|
||||
- `Shipyard Neo Sandbox TTL`
|
||||
- The upper lifetime limit of the sandbox, defaulting to 3600 seconds (1 hour)
|
||||
|
||||
After configuring, click the "Save" button at the bottom of the page to save the configuration.
|
||||
### Configuring Shipyard (Legacy)
|
||||
|
||||
If you choose the legacy `Shipyard` driver, the relevant configuration items are:
|
||||
|
||||
- `Shipyard API Endpoint`
|
||||
- If you use the Docker Compose deployment above, set it to `http://shipyard:8156`
|
||||
- If Shipyard is deployed separately, use the corresponding address, such as `http://<your-host>:8156`
|
||||
- `Shipyard Access Token`
|
||||
- Fill in the access token you configured when deploying Shipyard
|
||||
- `Shipyard Ship Lifetime (seconds)`
|
||||
- Defines the lifetime of each sandbox instance, default 3600 seconds (1 hour)
|
||||
- `Shipyard Ship Session Reuse Limit`
|
||||
- Defines the maximum number of sessions that can reuse the same sandbox instance, default 10
|
||||
|
||||
## About `Shipyard Neo Sandbox TTL`
|
||||
|
||||
In `Shipyard Neo`:
|
||||
|
||||
- TTL represents the upper lifetime bound of the sandbox
|
||||
- The selected profile also defines a separate idle timeout (`idle_timeout`)
|
||||
- Capability calls from AstrBot usually refresh the idle timeout, rather than directly extending the TTL
|
||||
- `keepalive` only extends the idle timeout; it does not automatically start a new session and does not extend the TTL
|
||||
|
||||
## About `Shipyard Ship Lifetime (seconds)`
|
||||
|
||||
The lifetime of a sandbox environment instance defines the maximum time each instance can exist before being destroyed. This setting needs to be determined based on your usage scenario and resources.
|
||||
The following explanation applies only to the legacy `Shipyard` driver:
|
||||
|
||||
- When a new session joins an existing sandbox environment instance, the instance will automatically extend its lifetime to the TTL requested by this session.
|
||||
- When an operation is performed on a sandbox environment instance, the instance will automatically extend its lifetime to the current time plus the TTL.
|
||||
The lifetime of a sandbox instance defines the maximum amount of time that instance can exist before being destroyed. This value should be chosen according to your use case and available resources.
|
||||
|
||||
- When a new session joins an existing sandbox instance, the instance automatically extends its lifetime to the TTL requested by that session
|
||||
- When an operation is performed on a sandbox instance, the instance automatically extends its lifetime to the current time plus TTL
|
||||
|
||||
## About Data Persistence in the Sandbox Environment
|
||||
|
||||
Shipyard allocates a working directory for each session under the `/home/<unique session ID>` directory.
|
||||
### Shipyard Neo
|
||||
|
||||
Shipyard automatically mounts the /home directory in the sandbox environment to the `${PWD}/data/shipyard/ship_mnt_data` directory on the host machine. When a sandbox environment instance is destroyed, if a session continues to request the sandbox, Shipyard will recreate a new sandbox environment instance and remount the previously persisted data to ensure data continuity.
|
||||
The workspace root of `Shipyard Neo` is fixed at `/workspace`.
|
||||
|
||||
Persistence is provided by Cargo:
|
||||
|
||||
- Filesystem data is stored in Cargo and mounted at `/workspace`
|
||||
- Even if the underlying Session is stopped or rebuilt, the data in Cargo is usually retained
|
||||
- For profiles with browser capability, browser state may also be persisted together, for example under `/workspace/.browser/profile/`
|
||||
|
||||
### Shipyard (Legacy)
|
||||
|
||||
Shipyard allocates a working directory for each session under `/home/<unique session ID>`.
|
||||
|
||||
Shipyard automatically mounts the `/home` directory from the sandbox environment to `${PWD}/data/shipyard/ship_mnt_data` on the host. When a sandbox instance is destroyed and a session later requests the sandbox again, Shipyard recreates a new instance and remounts the previously persisted data to preserve continuity.
|
||||
|
||||
## Other Community Plugins
|
||||
|
||||
### luosheng520qaq/astrobot_plugin_code_executor
|
||||
|
||||
If your resources are limited and you do not want to use the sandbox environment for code execution, you can try the [astrobot_plugin_code_executor](https://github.com/luosheng520qaq/astrobot_plugin_code_executor) plugin developed by luosheng520qaq. This plugin executes code directly on the host machine. It tries to improve safety as much as possible, but you should still pay close attention to code-execution security.
|
||||
|
||||
@@ -9,7 +9,26 @@
|
||||
|
||||
## 启用沙盒环境
|
||||
|
||||
目前,沙盒环境仅支持通过 Docker 来运行。我们目前使用了 [Shipyard](https://github.com/AstrBotDevs/shipyard) 项目作为 AstrBot 的沙盒环境驱动器。未来,我们会支持更多类型的沙盒环境驱动器,如 e2b。
|
||||
目前,AstrBot 的沙盒环境驱动器支持:
|
||||
|
||||
- `Shipyard Neo`(当前推荐)
|
||||
- `Shipyard`(旧方案,仍可继续使用)
|
||||
|
||||
在当前版本的 AstrBot 控制台中,可在“AI 配置” -> “Agent Computer Use”中选择:
|
||||
|
||||
- `Computer Use Runtime` = `sandbox`
|
||||
- `沙箱环境驱动器` = `Shipyard Neo` 或 `Shipyard`
|
||||
|
||||
其中,`Shipyard Neo` 是当前默认驱动器。它由 Bay、Ship、Gull 三部分组成:
|
||||
|
||||
- **Bay**:控制面 API,负责创建和管理 sandbox
|
||||
- **Ship**:负责 Python / Shell / 文件系统能力
|
||||
- **Gull**:负责浏览器自动化能力
|
||||
|
||||
对于 `Shipyard Neo`,工作区根目录固定为 `/workspace`。在 AstrBot 中调用文件系统工具时,应当传入**相对于工作区根目录**的路径,例如 `reports/result.txt`,而不是 `/workspace/reports/result.txt`。
|
||||
|
||||
> [!TIP]
|
||||
> `Shipyard Neo` 下浏览器能力并不是所有 profile 都有。只有 profile 支持 `browser` capability 时,AstrBot 才会挂载浏览器相关工具。典型 profile 如 `browser-python`。
|
||||
|
||||
## 性能要求
|
||||
|
||||
@@ -17,6 +36,242 @@ AstrBot 给每个沙盒环境限制最高 1 CPU 和 512 MB 内存。
|
||||
|
||||
我们建议您的宿主机至少有 2 个 CPU 和 4 GB 内存,并开启 Swap,以保证多个沙盒环境实例可以稳定运行。
|
||||
|
||||
## 推荐:使用 Shipyard Neo
|
||||
|
||||
### 单独部署 Shipyard Neo(推荐)
|
||||
|
||||
如果您准备长期使用 `Shipyard Neo`,更推荐将它**单独部署在一台资源更充足的机器上**,例如您的 homelab、局域网服务器,或独立云主机,然后再让 AstrBot 远程接入 Bay。
|
||||
|
||||
原因是:`Shipyard Neo` 在启用浏览器能力时需要运行较重的浏览器运行时。对于资源紧张的云服务器,把 AstrBot 和 `Shipyard Neo` 部署在同一台机器上,通常会让 CPU 和内存压力都比较大,稳定性和体验都不理想。
|
||||
|
||||
大致步骤如下:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/AstrBotDevs/shipyard-neo
|
||||
cd shipyard-neo/deploy/docker
|
||||
# 修改 config.yaml 中的关键配置,例如 security.api_key
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
部署完成后:
|
||||
|
||||
- Bay 默认监听在 `http://<your-host>:8114`
|
||||
- 在 AstrBot 控制台中选择 `Shipyard Neo` 驱动器
|
||||
- `Shipyard Neo API Endpoint` 填写对应地址,例如 `http://<your-host>:8114`
|
||||
- `Shipyard Neo Access Token` 填写 Bay API Key;如果 AstrBot 能访问 Bay 的 `credentials.json`,也可以留空让 AstrBot 自动发现
|
||||
|
||||
### 参考:`config.yaml` 完整示例(附说明)
|
||||
|
||||
如果您准备自行调整 `Shipyard Neo` 的部署参数,可以直接参考下面这份基于 [`deploy/docker/config.yaml`](https://github.com/AstrBotDevs/shipyard-neo/blob/main/deploy/docker/config.yaml) 整理的完整示例。它保留了默认结构,并额外加上了中文注释,便于理解每个配置项的用途。
|
||||
|
||||
> [!TIP]
|
||||
> 其中最少需要修改的是 `security.api_key`。如果不清楚其他参数的作用,建议先保持默认值,仅按需调整 profile、资源限制和 warm pool 配置。
|
||||
|
||||
```yaml
|
||||
# Bay Production Config - Docker Compose (container_network mode)
|
||||
#
|
||||
# Bay 运行在 Docker 容器中,并通过共享 Docker 网络与 Ship/Gull 容器通信。
|
||||
# 这种模式下,sandbox 容器不需要向宿主机暴露端口。
|
||||
#
|
||||
# 部署前至少需要修改:
|
||||
# 1. security.api_key —— 设置强随机密钥
|
||||
|
||||
server:
|
||||
# Bay API 监听地址
|
||||
host: "0.0.0.0"
|
||||
# Bay API 监听端口
|
||||
port: 8114
|
||||
|
||||
database:
|
||||
# 单机部署默认使用 SQLite。
|
||||
# 如果要做多实例 / 高可用,可改用 PostgreSQL,例如:
|
||||
# url: "postgresql+asyncpg://user:pass@db-host:5432/bay"
|
||||
url: "sqlite+aiosqlite:///./data/bay.db"
|
||||
echo: false
|
||||
|
||||
driver:
|
||||
# 当前默认使用 Docker 驱动
|
||||
type: docker
|
||||
|
||||
# 创建新 sandbox 时是否拉取镜像。
|
||||
# 生产环境通常建议 always,以便拿到最新镜像。
|
||||
image_pull_policy: always
|
||||
|
||||
docker:
|
||||
# Docker Socket 地址
|
||||
socket: "unix:///var/run/docker.sock"
|
||||
|
||||
# Bay 在容器内运行,Ship/Gull 也在容器内运行时,
|
||||
# 推荐使用 container_network 通过容器网络直接通信。
|
||||
connect_mode: container_network
|
||||
|
||||
# 共享网络名,必须与 docker-compose.yaml 中的网络一致
|
||||
network: "bay-network"
|
||||
|
||||
# 是否将 sandbox 容器端口暴露到宿主机。
|
||||
# 生产环境建议关闭,以减少攻击面。
|
||||
publish_ports: false
|
||||
host_port: null
|
||||
|
||||
cargo:
|
||||
# Cargo 在 Bay 侧的存储根路径
|
||||
root_path: "/var/lib/bay/cargos"
|
||||
# 默认工作区大小限制(MB)
|
||||
default_size_limit_mb: 1024
|
||||
# Cargo 挂载到 sandbox 内的路径。AstrBot/Neo 的工作区根目录就是这里。
|
||||
mount_path: "/workspace"
|
||||
|
||||
security:
|
||||
# 必改项:设置一个强随机密钥,例如 openssl rand -hex 32
|
||||
api_key: "CHANGE-ME"
|
||||
# 是否允许匿名访问。生产环境建议 false。
|
||||
allow_anonymous: false
|
||||
|
||||
# 容器代理环境变量注入。
|
||||
# 启用后,Bay 会把 HTTP(S)_PROXY 和 NO_PROXY 注入到 sandbox 容器。
|
||||
proxy:
|
||||
enabled: false
|
||||
# http_proxy: "http://proxy.example.com:7890"
|
||||
# https_proxy: "http://proxy.example.com:7890"
|
||||
# no_proxy: "my-internal.service"
|
||||
|
||||
# Warm Pool:预热一批待命 sandbox,减少冷启动延迟。
|
||||
# 当用户创建 sandbox 时,Bay 会优先尝试领取一个已预热实例。
|
||||
warm_pool:
|
||||
enabled: true
|
||||
# 预热队列 worker 数量
|
||||
warmup_queue_workers: 2
|
||||
# 预热队列最大长度
|
||||
warmup_queue_max_size: 256
|
||||
# 队列满时的丢弃策略
|
||||
warmup_queue_drop_policy: "drop_newest"
|
||||
# 超过这个阈值时便于运维告警
|
||||
warmup_queue_drop_alert_threshold: 50
|
||||
# 预热池维护扫描周期(秒)
|
||||
interval_seconds: 30
|
||||
# Bay 启动时是否立即运行预热逻辑
|
||||
run_on_startup: true
|
||||
|
||||
profiles:
|
||||
# ── 标准 Python 沙箱 ────────────────────────
|
||||
- id: python-default
|
||||
description: "Standard Python sandbox with filesystem and shell access"
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "1g"
|
||||
capabilities:
|
||||
- filesystem # 包含 upload/download
|
||||
- shell
|
||||
- python
|
||||
# 空闲超时(秒)
|
||||
idle_timeout: 1800
|
||||
# 保持 1 个预热实例
|
||||
warm_pool_size: 1
|
||||
env: {}
|
||||
# 可选:profile 级代理覆盖
|
||||
# proxy:
|
||||
# enabled: false
|
||||
|
||||
# ── 数据科学沙箱(更多资源) ──────────
|
||||
- id: python-data
|
||||
description: "Data science sandbox with extra CPU and memory"
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 2.0
|
||||
memory: "4g"
|
||||
capabilities:
|
||||
- filesystem # 包含 upload/download
|
||||
- shell
|
||||
- python
|
||||
idle_timeout: 1800
|
||||
warm_pool_size: 1
|
||||
env: {}
|
||||
|
||||
# ── 浏览器 + Python 多容器沙箱 ───────
|
||||
- id: browser-python
|
||||
description: "Browser automation with Python backend"
|
||||
containers:
|
||||
- name: ship
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-ship:latest"
|
||||
runtime_type: ship
|
||||
runtime_port: 8123
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "1g"
|
||||
capabilities:
|
||||
- python
|
||||
- shell
|
||||
- filesystem # 包含 upload/download
|
||||
# 这些能力优先由 ship 容器提供
|
||||
primary_for:
|
||||
- filesystem
|
||||
- python
|
||||
- shell
|
||||
env: {}
|
||||
- name: browser
|
||||
image: "ghcr.io/astrbotdevs/shipyard-neo-gull:latest"
|
||||
runtime_type: gull
|
||||
runtime_port: 8115
|
||||
resources:
|
||||
cpus: 1.0
|
||||
memory: "2g"
|
||||
capabilities:
|
||||
- browser
|
||||
env: {}
|
||||
idle_timeout: 1800
|
||||
warm_pool_size: 1
|
||||
|
||||
gc:
|
||||
# 生产环境建议启用自动 GC
|
||||
enabled: true
|
||||
run_on_startup: true
|
||||
# GC 扫描周期(秒)
|
||||
interval_seconds: 300
|
||||
|
||||
# 多实例部署时必须保证唯一
|
||||
instance_id: "bay-prod"
|
||||
|
||||
idle_session:
|
||||
enabled: true
|
||||
expired_sandbox:
|
||||
enabled: true
|
||||
orphan_cargo:
|
||||
enabled: true
|
||||
orphan_container:
|
||||
# 建议在生产环境开启,用于清理遗留容器
|
||||
enabled: true
|
||||
```
|
||||
|
||||
通常可以按下面的思路理解和修改:
|
||||
|
||||
- **最小必改项**:`security.api_key`
|
||||
- **最常改项**:`profiles` 里的资源限制、`warm_pool_size`、`idle_timeout`
|
||||
- **需要浏览器能力时**:使用或调整 `browser-python` profile
|
||||
- **希望减少冷启动时间时**:保留 `warm_pool.enabled: true`,并适当提高常用 profile 的 `warm_pool_size`
|
||||
- **资源较紧张时**:可先把 `warm_pool_size` 改小,甚至关闭 `warm_pool`
|
||||
- **如果需要代理访问外网**:配置顶层 `proxy`,或按 profile 单独覆盖
|
||||
|
||||
### 关于 Shipyard Neo 的复用与持久化
|
||||
|
||||
`Shipyard Neo` 中有几个重要概念:
|
||||
|
||||
- **Sandbox**:对外稳定可见的资源单元
|
||||
- **Session**:实际运行中的容器会话,可被停止或重建
|
||||
- **Cargo**:持久化工作区卷,挂载到 `/workspace`
|
||||
|
||||
对 AstrBot 而言,当前会按请求的 `session_id` 维度缓存沙箱 booter;在主 Agent 默认流程下,这个 `session_id` 通常等于消息会话标识 `unified_msg_origin`。因此,同一消息会话的后续请求通常会继续复用同一个 Neo sandbox;如果沙箱失效,则会自动重建。
|
||||
|
||||
关于 TTL 与数据持久化的更详细说明,请参考下文的“关于 `Shipyard Neo Sandbox TTL`”与“关于沙盒环境的数据持久化”小节。
|
||||
|
||||
## 旧方案:Shipyard
|
||||
|
||||
以下内容为旧版 `Shipyard` 驱动器的部署与配置说明,仍然保留,供兼容旧部署方案时参考。
|
||||
|
||||
### 使用 Docker Compose 部署 AstrBot 和 Shipyard
|
||||
|
||||
如果您还没有部署 AstrBot,或者想更换为我们推荐的带沙盒环境的部署方式,推荐使用 Docker Compose 来部署 AstrBot,代码如下:
|
||||
@@ -56,22 +311,56 @@ docker pull soulter/shipyard-ship:latest
|
||||
> [!TIP]
|
||||
> 请确保您的 AstrBot 版本在 `v4.12.0` 及之后。
|
||||
|
||||
在 AstrBot 控制台,进入 “配置文件” 页面,找到 “Agent 沙箱环境”,启用沙箱环境开关。
|
||||
在 AstrBot 控制台,进入 “AI 配置” -> “Agent Computer Use”。
|
||||
|
||||
在出现的配置项中,
|
||||
1. 将 `Computer Use Runtime` 设为 `sandbox`
|
||||
2. 在 `沙箱环境驱动器` 中选择 `Shipyard Neo` 或 `Shipyard`
|
||||
3. 根据驱动器填写对应配置项
|
||||
4. 点击右下角“保存”
|
||||
|
||||
对于 `Shipyard API Endpoint`,如果您使用上述的 Docker Compose 部署方式,填写 `http://shipyard:8156` 即可。如果您是单独部署的 Shipyard,请填写对应的地址,例如 `http://<your-host>:8156`。
|
||||
### 配置 Shipyard Neo
|
||||
|
||||
对于 `Shipyard Access Token`,请填写您在部署 Shipyard 时配置的访问令牌。
|
||||
如果您选择的是 `Shipyard Neo`,主要配置项如下:
|
||||
|
||||
对于 `Shipyard Ship 存活时间(秒)`,这个定义了每个沙箱环境实例的存活时间,默认值为 3600 秒(1 小时)。您可以根据需要调整这个值。
|
||||
- `Shipyard Neo API Endpoint`
|
||||
- 联合部署时可填写 `http://bay:8114`
|
||||
- 单独部署时填写实际地址,例如 `http://<your-host>:8114`
|
||||
- `Shipyard Neo Access Token`
|
||||
- 填写 Bay API Key
|
||||
- 如果是官方联合部署,且 AstrBot 能访问 Bay 的 `credentials.json`,可以留空自动发现
|
||||
- `Shipyard Neo Profile`
|
||||
- 例如 `python-default`、`browser-python`
|
||||
- 如果未手动指定,AstrBot 会优先尝试选择能力更完整、且优先带有 `browser` capability 的 profile,失败时再回退到 `python-default`
|
||||
- `Shipyard Neo Sandbox TTL`
|
||||
- sandbox 生命周期上限,默认值为 3600 秒(1 小时)
|
||||
|
||||
对于 `Shipyard Ship 会话复用上限`,这个定义了每个沙箱环境实例可以复用的最大会话数,默认值为 10。也就是 10 个会话会共享同一个沙箱环境实例。您可以根据需要调整这个值。
|
||||
### 配置 Shipyard(旧方案)
|
||||
|
||||
填写好之后,点击右下角 “保存” 即可。
|
||||
如果您选择的是旧版 `Shipyard`,配置项如下:
|
||||
|
||||
- `Shipyard API Endpoint`
|
||||
- 如果您使用上述 Docker Compose 部署方式,填写 `http://shipyard:8156` 即可
|
||||
- 如果您是单独部署的 Shipyard,请填写对应地址,例如 `http://<your-host>:8156`
|
||||
- `Shipyard Access Token`
|
||||
- 请填写部署 Shipyard 时配置的访问令牌
|
||||
- `Shipyard Ship 存活时间(秒)`
|
||||
- 定义每个沙箱环境实例的存活时间,默认值为 3600 秒(1 小时)
|
||||
- `Shipyard Ship 会话复用上限`
|
||||
- 定义每个沙箱环境实例可以复用的最大会话数,默认值为 10
|
||||
|
||||
## 关于 `Shipyard Neo Sandbox TTL`
|
||||
|
||||
在 `Shipyard Neo` 中:
|
||||
|
||||
- TTL 表示 sandbox 生命周期上限
|
||||
- profile 还会定义一个独立的空闲超时(`idle_timeout`)
|
||||
- AstrBot 发起能力调用时,通常会刷新空闲超时,而不是直接延长 TTL
|
||||
- `keepalive` 只会延长空闲超时,不会自动启动新的 session,也不会延长 TTL
|
||||
|
||||
## 关于 `Shipyard Ship 存活时间(秒)`
|
||||
|
||||
以下说明仅适用于旧版 `Shipyard`:
|
||||
|
||||
沙箱环境实例的存活时间定义了每个实例在被销毁之前可以存在的最长时间,这个时间的设置需要根据您的使用场景以及资源来决定。
|
||||
|
||||
- 新的会话加入已有的沙箱环境实例时,该实例会自动延长存活时间到这个会话请求的 TTL。
|
||||
@@ -79,6 +368,18 @@ docker pull soulter/shipyard-ship:latest
|
||||
|
||||
## 关于沙盒环境的数据持久化
|
||||
|
||||
### Shipyard Neo
|
||||
|
||||
`Shipyard Neo` 的工作区根目录固定为 `/workspace`。
|
||||
|
||||
其持久化由 Cargo 提供:
|
||||
|
||||
- 文件系统数据保存在 Cargo 中,并挂载到 `/workspace`
|
||||
- 即使底层 Session 被停止或重建,Cargo 中的数据通常仍可保留
|
||||
- 对于带浏览器能力的 profile,浏览器状态也可能会一起持久化,例如 `/workspace/.browser/profile/`
|
||||
|
||||
### Shipyard(旧方案)
|
||||
|
||||
Shipyard 会给每个会话分配一个工作目录,在 `/home/<会话唯一 ID>` 目录下。
|
||||
|
||||
Shipyard 会自动将沙盒环境中的 /home 目录挂载到宿主机的 `${PWD}/data/shipyard/ship_mnt_data` 目录下,当沙盒环境实例被销毁后,如果某个会话继续请求调用沙箱,Shipyard 会重新创建一个新的沙盒环境实例,并将之前持久化的数据重新挂载进去,保证数据的连续性。
|
||||
|
||||
103
tests/agent/test_token_counter.py
Normal file
103
tests/agent/test_token_counter.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""Tests for EstimateTokenCounter multimodal support."""
|
||||
|
||||
from astrbot.core.agent.context.token_counter import (
|
||||
AUDIO_TOKEN_ESTIMATE,
|
||||
IMAGE_TOKEN_ESTIMATE,
|
||||
EstimateTokenCounter,
|
||||
)
|
||||
from astrbot.core.agent.message import (
|
||||
AudioURLPart,
|
||||
ImageURLPart,
|
||||
Message,
|
||||
TextPart,
|
||||
ThinkPart,
|
||||
)
|
||||
|
||||
|
||||
counter = EstimateTokenCounter()
|
||||
|
||||
|
||||
def _msg(role: str, content) -> Message:
|
||||
return Message(role=role, content=content)
|
||||
|
||||
|
||||
class TestTextCounting:
|
||||
def test_plain_string(self):
|
||||
tokens = counter.count_tokens([_msg("user", "hello world")])
|
||||
assert tokens > 0
|
||||
|
||||
def test_chinese(self):
|
||||
# 中文字符权重更高
|
||||
en = counter.count_tokens([_msg("user", "abc")])
|
||||
zh = counter.count_tokens([_msg("user", "你好啊")])
|
||||
assert zh > en
|
||||
|
||||
def test_text_part(self):
|
||||
msg = _msg("user", [TextPart(text="hello")])
|
||||
assert counter.count_tokens([msg]) > 0
|
||||
|
||||
|
||||
class TestMultimodalCounting:
|
||||
def test_image_counted(self):
|
||||
msg = _msg("user", [
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,abc")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == IMAGE_TOKEN_ESTIMATE
|
||||
|
||||
def test_audio_counted(self):
|
||||
msg = _msg("user", [
|
||||
AudioURLPart(audio_url=AudioURLPart.AudioURL(url="https://x.com/a.mp3")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == AUDIO_TOKEN_ESTIMATE
|
||||
|
||||
def test_think_counted(self):
|
||||
msg = _msg("assistant", [ThinkPart(think="let me think about this")])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens > 0
|
||||
|
||||
def test_mixed_content(self):
|
||||
"""文本 + 图片的多模态消息,token 数 = 文本 token + 图片估算。"""
|
||||
text_only = _msg("user", [TextPart(text="describe this image")])
|
||||
mixed = _msg("user", [
|
||||
TextPart(text="describe this image"),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,x")),
|
||||
])
|
||||
text_tokens = counter.count_tokens([text_only])
|
||||
mixed_tokens = counter.count_tokens([mixed])
|
||||
assert mixed_tokens == text_tokens + IMAGE_TOKEN_ESTIMATE
|
||||
|
||||
def test_multiple_images(self):
|
||||
"""多张图片应该各自计算。"""
|
||||
msg = _msg("user", [
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,a")),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,b")),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,c")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg])
|
||||
assert tokens == IMAGE_TOKEN_ESTIMATE * 3
|
||||
|
||||
|
||||
class TestTrustedUsage:
|
||||
def test_trusted_overrides(self):
|
||||
"""如果 API 返回了 token 数,直接用它不做估算。"""
|
||||
msg = _msg("user", [
|
||||
TextPart(text="hello"),
|
||||
ImageURLPart(image_url=ImageURLPart.ImageURL(url="data:image/png;base64,x")),
|
||||
])
|
||||
tokens = counter.count_tokens([msg], trusted_token_usage=42)
|
||||
assert tokens == 42
|
||||
|
||||
|
||||
class TestToolCalls:
|
||||
def test_tool_calls_counted(self):
|
||||
msg = Message(
|
||||
role="assistant",
|
||||
content="calling tool",
|
||||
tool_calls=[{"type": "function", "id": "1", "function": {"name": "get_weather", "arguments": '{"city": "Beijing"}'}}],
|
||||
)
|
||||
tokens = counter.count_tokens([msg])
|
||||
# 文本 + tool call JSON 都应被计算
|
||||
text_only = counter.count_tokens([_msg("assistant", "calling tool")])
|
||||
assert tokens > text_only
|
||||
@@ -104,8 +104,9 @@ class TestContextTruncator:
|
||||
messages, keep_most_recent_turns=0, drop_turns=1
|
||||
)
|
||||
|
||||
# Should result in empty or minimal list
|
||||
assert len(result) == 0
|
||||
# 截断后至少保留一条 user 消息 (#6196)
|
||||
assert len(result) >= 1
|
||||
assert result[0].role == "user"
|
||||
|
||||
def test_truncate_by_turns_below_threshold(self):
|
||||
"""Test truncate_by_turns when messages are below threshold."""
|
||||
@@ -201,8 +202,9 @@ class TestContextTruncator:
|
||||
messages = self.create_messages(4)
|
||||
result = truncator.truncate_by_dropping_oldest_turns(messages, drop_turns=2)
|
||||
|
||||
# Should drop all turns
|
||||
assert len(result) == 0
|
||||
# 即使 drop 掉所有 turn,也会把 user 消息补回来 (#6196)
|
||||
assert len(result) >= 1
|
||||
assert result[0].role == "user"
|
||||
|
||||
def test_truncate_by_dropping_oldest_turns_drop_more_than_available(self):
|
||||
"""Test truncate_by_dropping_oldest_turns with drop_turns > available turns."""
|
||||
@@ -211,8 +213,9 @@ class TestContextTruncator:
|
||||
messages = self.create_messages(4)
|
||||
result = truncator.truncate_by_dropping_oldest_turns(messages, drop_turns=5)
|
||||
|
||||
# Should result in empty list
|
||||
assert len(result) == 0
|
||||
# 同理,user 消息会被保留 (#6196)
|
||||
assert len(result) >= 1
|
||||
assert result[0].role == "user"
|
||||
|
||||
def test_truncate_by_dropping_oldest_turns_ensures_user_first(self):
|
||||
"""Test that result starts with user message after dropping."""
|
||||
@@ -372,3 +375,60 @@ class TestContextTruncator:
|
||||
assert len(result) >= 0 # May keep system messages or clear all
|
||||
if len(result) > 0:
|
||||
assert all(msg.role == "system" for msg in result)
|
||||
|
||||
# ==================== #6196: 长 tool chain 只有一条 user 消息 ====================
|
||||
|
||||
def _build_tool_chain(self, tool_rounds: int = 20) -> list[Message]:
|
||||
"""构造 system -> user -> (assistant -> tool) * N 的长链,只有一条 user。"""
|
||||
msgs = [
|
||||
self.create_message("system", "You are a helpful assistant."),
|
||||
self.create_message("user", "帮我查一下天气"),
|
||||
]
|
||||
for i in range(tool_rounds):
|
||||
msgs.append(self.create_message("assistant", f"调用工具 {i}"))
|
||||
msgs.append(self.create_message("tool", f"工具结果 {i}"))
|
||||
return msgs
|
||||
|
||||
def test_drop_oldest_preserves_sole_user(self):
|
||||
"""#6196: drop 1 turn 不应丢掉唯一的 user 消息。"""
|
||||
truncator = ContextTruncator()
|
||||
msgs = self._build_tool_chain(20) # 1 system + 1 user + 40 asst/tool = 42
|
||||
result = truncator.truncate_by_dropping_oldest_turns(msgs, drop_turns=1)
|
||||
roles = [m.role for m in result]
|
||||
assert "user" in roles, "唯一的 user 消息被丢掉了"
|
||||
assert roles[0] == "system"
|
||||
|
||||
def test_halving_preserves_sole_user(self):
|
||||
"""#6196: 对半砍不应丢掉唯一的 user 消息。"""
|
||||
truncator = ContextTruncator()
|
||||
msgs = self._build_tool_chain(20)
|
||||
result = truncator.truncate_by_halving(msgs)
|
||||
roles = [m.role for m in result]
|
||||
assert "user" in roles, "唯一的 user 消息被丢掉了"
|
||||
|
||||
def test_truncate_by_turns_preserves_sole_user(self):
|
||||
"""#6196: keep_most_recent_turns 也不应丢掉唯一的 user 消息。"""
|
||||
truncator = ContextTruncator()
|
||||
msgs = self._build_tool_chain(20)
|
||||
result = truncator.truncate_by_turns(
|
||||
msgs, keep_most_recent_turns=3, drop_turns=1
|
||||
)
|
||||
roles = [m.role for m in result]
|
||||
assert "user" in roles, "唯一的 user 消息被丢掉了"
|
||||
|
||||
def test_drop_oldest_heavy_drops_still_has_user(self):
|
||||
"""#6196: 大量 drop 也不会丢 user。"""
|
||||
truncator = ContextTruncator()
|
||||
msgs = self._build_tool_chain(30)
|
||||
result = truncator.truncate_by_dropping_oldest_turns(msgs, drop_turns=10)
|
||||
roles = [m.role for m in result]
|
||||
assert "user" in roles
|
||||
|
||||
def test_normal_multi_user_not_affected(self):
|
||||
"""正常多 user 对话不受影响。"""
|
||||
truncator = ContextTruncator()
|
||||
msgs = self.create_messages(20, include_system=True)
|
||||
result_before = truncator.truncate_by_dropping_oldest_turns(msgs, drop_turns=2)
|
||||
# 多 user 场景下截断后仍有 user
|
||||
roles = [m.role for m in result_before]
|
||||
assert "user" in roles
|
||||
|
||||
81
tests/test_anthropic_kimi_code_provider.py
Normal file
81
tests/test_anthropic_kimi_code_provider.py
Normal file
@@ -0,0 +1,81 @@
|
||||
import httpx
|
||||
|
||||
import astrbot.core.provider.sources.anthropic_source as anthropic_source
|
||||
import astrbot.core.provider.sources.kimi_code_source as kimi_code_source
|
||||
|
||||
|
||||
class _FakeAsyncAnthropic:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
def test_anthropic_provider_injects_custom_headers_into_http_client(monkeypatch):
|
||||
monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)
|
||||
|
||||
provider = anthropic_source.ProviderAnthropic(
|
||||
provider_config={
|
||||
"id": "anthropic-test",
|
||||
"type": "anthropic_chat_completion",
|
||||
"model": "claude-test",
|
||||
"key": ["test-key"],
|
||||
"custom_headers": {
|
||||
"User-Agent": "custom-agent/1.0",
|
||||
"X-Test-Header": 123,
|
||||
},
|
||||
},
|
||||
provider_settings={},
|
||||
)
|
||||
|
||||
assert provider.custom_headers == {
|
||||
"User-Agent": "custom-agent/1.0",
|
||||
"X-Test-Header": "123",
|
||||
}
|
||||
assert isinstance(provider.client.kwargs["http_client"], httpx.AsyncClient)
|
||||
assert provider.client.kwargs["http_client"].headers["User-Agent"] == "custom-agent/1.0"
|
||||
assert provider.client.kwargs["http_client"].headers["X-Test-Header"] == "123"
|
||||
|
||||
|
||||
def test_kimi_code_provider_sets_defaults_and_preserves_custom_headers(monkeypatch):
|
||||
monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)
|
||||
|
||||
provider = kimi_code_source.ProviderKimiCode(
|
||||
provider_config={
|
||||
"id": "kimi-code",
|
||||
"type": "kimi_code_chat_completion",
|
||||
"key": ["test-key"],
|
||||
"custom_headers": {"X-Trace-Id": "trace-1"},
|
||||
},
|
||||
provider_settings={},
|
||||
)
|
||||
|
||||
assert provider.base_url == kimi_code_source.KIMI_CODE_API_BASE
|
||||
assert provider.get_model() == kimi_code_source.KIMI_CODE_DEFAULT_MODEL
|
||||
assert provider.custom_headers == {
|
||||
"User-Agent": kimi_code_source.KIMI_CODE_USER_AGENT,
|
||||
"X-Trace-Id": "trace-1",
|
||||
}
|
||||
assert provider.client.kwargs["http_client"].headers["User-Agent"] == (
|
||||
kimi_code_source.KIMI_CODE_USER_AGENT
|
||||
)
|
||||
assert provider.client.kwargs["http_client"].headers["X-Trace-Id"] == "trace-1"
|
||||
|
||||
|
||||
def test_kimi_code_provider_restores_required_user_agent_when_blank(monkeypatch):
|
||||
monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)
|
||||
|
||||
provider = kimi_code_source.ProviderKimiCode(
|
||||
provider_config={
|
||||
"id": "kimi-code",
|
||||
"type": "kimi_code_chat_completion",
|
||||
"key": ["test-key"],
|
||||
"custom_headers": {"User-Agent": " "},
|
||||
},
|
||||
provider_settings={},
|
||||
)
|
||||
|
||||
assert provider.custom_headers == {
|
||||
"User-Agent": kimi_code_source.KIMI_CODE_USER_AGENT,
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import ntpath
|
||||
import threading
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
@@ -9,6 +10,15 @@ from astrbot.core.utils import pip_installer as pip_installer_module
|
||||
from astrbot.core.utils import requirements_utils
|
||||
from astrbot.core.utils.pip_installer import PipInstaller
|
||||
|
||||
WINDOWS_RUNTIME_ROOT = ntpath.join(r"C:\astrbot-test", "backend", "python")
|
||||
WINDOWS_RUNTIME_EXECUTABLE = ntpath.join(WINDOWS_RUNTIME_ROOT, "python.exe")
|
||||
WINDOWS_PACKAGED_RUNTIME_EXECUTABLE = f"\\\\?\\{WINDOWS_RUNTIME_EXECUTABLE}"
|
||||
WINDOWS_RUNTIME_INCLUDE_DIR = ntpath.join(WINDOWS_RUNTIME_ROOT, "include")
|
||||
WINDOWS_RUNTIME_LIBS_DIR = ntpath.join(WINDOWS_RUNTIME_ROOT, "libs")
|
||||
EXISTING_WINDOWS_INCLUDE_DIR = ntpath.join(r"C:\toolchain", "include")
|
||||
EXISTING_WINDOWS_LIB_DIR = ntpath.join(r"C:\toolchain", "lib")
|
||||
_ENV_MISSING = object()
|
||||
|
||||
|
||||
def _make_run_pip_mock(
|
||||
code: int = 0,
|
||||
@@ -24,6 +34,64 @@ def _make_run_pip_mock(
|
||||
return AsyncMock(side_effect=run_pip)
|
||||
|
||||
|
||||
def _configure_run_pip_in_process_capture(
|
||||
monkeypatch,
|
||||
*,
|
||||
platform: str,
|
||||
packaged_runtime: bool,
|
||||
runtime_executable: str = WINDOWS_PACKAGED_RUNTIME_EXECUTABLE,
|
||||
include_value: str | object = _ENV_MISSING,
|
||||
lib_value: str | object = _ENV_MISSING,
|
||||
existing_runtime_dirs: set[str] | None = None,
|
||||
) -> dict[str, str | None]:
|
||||
observed_env: dict[str, str | None] = {}
|
||||
|
||||
def fake_pip_main(args):
|
||||
del args
|
||||
observed_env["INCLUDE"] = pip_installer_module.os.environ.get("INCLUDE")
|
||||
observed_env["LIB"] = pip_installer_module.os.environ.get("LIB")
|
||||
return 0
|
||||
|
||||
if packaged_runtime:
|
||||
monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1")
|
||||
else:
|
||||
monkeypatch.delenv("ASTRBOT_DESKTOP_CLIENT", raising=False)
|
||||
|
||||
if include_value is _ENV_MISSING:
|
||||
monkeypatch.delenv("INCLUDE", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("INCLUDE", include_value)
|
||||
|
||||
if lib_value is _ENV_MISSING:
|
||||
monkeypatch.delenv("LIB", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("LIB", lib_value)
|
||||
|
||||
monkeypatch.setattr(pip_installer_module.sys, "platform", platform)
|
||||
monkeypatch.setattr(pip_installer_module.sys, "executable", runtime_executable)
|
||||
|
||||
if existing_runtime_dirs is not None:
|
||||
monkeypatch.setattr(
|
||||
pip_installer_module.os.path,
|
||||
"isdir",
|
||||
lambda path: path in existing_runtime_dirs,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"astrbot.core.utils.pip_installer._get_pip_main",
|
||||
lambda: fake_pip_main,
|
||||
)
|
||||
return observed_env
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def configure_run_pip_in_process_capture(monkeypatch):
|
||||
def _configure(**kwargs):
|
||||
return _configure_run_pip_in_process_capture(monkeypatch, **kwargs)
|
||||
|
||||
return _configure
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_install_targets_site_packages_for_desktop_client(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1")
|
||||
@@ -226,6 +294,316 @@ async def test_run_pip_in_process_normalizes_crlf_without_extra_blank_lines(
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "expected"),
|
||||
[
|
||||
(
|
||||
WINDOWS_RUNTIME_EXECUTABLE,
|
||||
WINDOWS_RUNTIME_EXECUTABLE,
|
||||
),
|
||||
(
|
||||
WINDOWS_PACKAGED_RUNTIME_EXECUTABLE,
|
||||
WINDOWS_RUNTIME_EXECUTABLE,
|
||||
),
|
||||
(
|
||||
f"\\??\\{WINDOWS_RUNTIME_EXECUTABLE}",
|
||||
WINDOWS_RUNTIME_EXECUTABLE,
|
||||
),
|
||||
(
|
||||
r"\\?\UNC\server\share\include",
|
||||
r"\\server\share\include",
|
||||
),
|
||||
(
|
||||
r"\??\UNC\server\share\libs",
|
||||
r"\\server\share\libs",
|
||||
),
|
||||
(
|
||||
r"\\server\share\include",
|
||||
r"\\server\share\include",
|
||||
),
|
||||
(
|
||||
"C:/astrbot-test/backend/python/libs",
|
||||
WINDOWS_RUNTIME_LIBS_DIR,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_normalize_windows_native_build_path_variants(path, expected):
|
||||
assert pip_installer_module._normalize_windows_native_build_path(path) == expected
|
||||
|
||||
|
||||
def test_temporary_environ_restores_previous_values(monkeypatch):
|
||||
monkeypatch.setenv("INCLUDE", EXISTING_WINDOWS_INCLUDE_DIR)
|
||||
monkeypatch.delenv("LIB", raising=False)
|
||||
|
||||
with pip_installer_module._temporary_environ(
|
||||
{
|
||||
"INCLUDE": WINDOWS_RUNTIME_INCLUDE_DIR,
|
||||
"LIB": WINDOWS_RUNTIME_LIBS_DIR,
|
||||
}
|
||||
):
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == WINDOWS_RUNTIME_INCLUDE_DIR
|
||||
assert pip_installer_module.os.environ["LIB"] == WINDOWS_RUNTIME_LIBS_DIR
|
||||
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR
|
||||
assert "LIB" not in pip_installer_module.os.environ
|
||||
|
||||
|
||||
def test_build_packaged_windows_runtime_build_env_uses_base_env_snapshot(
|
||||
monkeypatch,
|
||||
):
|
||||
snapshot_include = ntpath.join(r"C:\snapshot-toolchain", "include")
|
||||
snapshot_lib = ntpath.join(r"C:\snapshot-toolchain", "lib")
|
||||
process_include = ntpath.join(r"C:\process-toolchain", "include")
|
||||
process_lib = ntpath.join(r"C:\process-toolchain", "lib")
|
||||
|
||||
monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1")
|
||||
monkeypatch.setenv("INCLUDE", process_include)
|
||||
monkeypatch.setenv("LIB", process_lib)
|
||||
monkeypatch.setattr(pip_installer_module.sys, "platform", "win32")
|
||||
monkeypatch.setattr(
|
||||
pip_installer_module.sys,
|
||||
"executable",
|
||||
WINDOWS_PACKAGED_RUNTIME_EXECUTABLE,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
pip_installer_module.os.path,
|
||||
"isdir",
|
||||
lambda path: path in {WINDOWS_RUNTIME_INCLUDE_DIR, WINDOWS_RUNTIME_LIBS_DIR},
|
||||
)
|
||||
|
||||
env_updates = pip_installer_module._build_packaged_windows_runtime_build_env(
|
||||
base_env={
|
||||
"INCLUDE": snapshot_include,
|
||||
"LIB": snapshot_lib,
|
||||
}
|
||||
)
|
||||
|
||||
assert env_updates == {
|
||||
"INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{snapshot_include}",
|
||||
"LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{snapshot_lib}",
|
||||
}
|
||||
|
||||
|
||||
def test_build_packaged_windows_runtime_build_env_matches_snapshot_keys_case_insensitively(
|
||||
monkeypatch,
|
||||
):
|
||||
snapshot_include = ntpath.join(r"C:\snapshot-toolchain", "include")
|
||||
snapshot_lib = ntpath.join(r"C:\snapshot-toolchain", "lib")
|
||||
|
||||
monkeypatch.setenv("ASTRBOT_DESKTOP_CLIENT", "1")
|
||||
monkeypatch.setattr(pip_installer_module.sys, "platform", "win32")
|
||||
monkeypatch.setattr(
|
||||
pip_installer_module.sys,
|
||||
"executable",
|
||||
WINDOWS_PACKAGED_RUNTIME_EXECUTABLE,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
pip_installer_module.os.path,
|
||||
"isdir",
|
||||
lambda path: path in {WINDOWS_RUNTIME_INCLUDE_DIR, WINDOWS_RUNTIME_LIBS_DIR},
|
||||
)
|
||||
|
||||
env_updates = pip_installer_module._build_packaged_windows_runtime_build_env(
|
||||
base_env={
|
||||
"include": snapshot_include,
|
||||
"lib": snapshot_lib,
|
||||
}
|
||||
)
|
||||
|
||||
assert env_updates == {
|
||||
"INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{snapshot_include}",
|
||||
"LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{snapshot_lib}",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("include_exists", "libs_exists"),
|
||||
[
|
||||
(True, True),
|
||||
(True, False),
|
||||
(False, True),
|
||||
],
|
||||
)
|
||||
async def test_run_pip_in_process_injects_windows_runtime_build_env(
|
||||
configure_run_pip_in_process_capture, include_exists, libs_exists
|
||||
):
|
||||
existing_runtime_dirs = set()
|
||||
if include_exists:
|
||||
existing_runtime_dirs.add(WINDOWS_RUNTIME_INCLUDE_DIR)
|
||||
if libs_exists:
|
||||
existing_runtime_dirs.add(WINDOWS_RUNTIME_LIBS_DIR)
|
||||
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="win32",
|
||||
packaged_runtime=True,
|
||||
include_value=EXISTING_WINDOWS_INCLUDE_DIR,
|
||||
lib_value=EXISTING_WINDOWS_LIB_DIR,
|
||||
existing_runtime_dirs=existing_runtime_dirs,
|
||||
)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
expected_include = EXISTING_WINDOWS_INCLUDE_DIR
|
||||
expected_lib = EXISTING_WINDOWS_LIB_DIR
|
||||
if include_exists:
|
||||
expected_include = (
|
||||
f"{WINDOWS_RUNTIME_INCLUDE_DIR};{EXISTING_WINDOWS_INCLUDE_DIR}"
|
||||
)
|
||||
if libs_exists:
|
||||
expected_lib = f"{WINDOWS_RUNTIME_LIBS_DIR};{EXISTING_WINDOWS_LIB_DIR}"
|
||||
assert observed_env == {"INCLUDE": expected_include, "LIB": expected_lib}
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR
|
||||
assert pip_installer_module.os.environ["LIB"] == EXISTING_WINDOWS_LIB_DIR
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("include_exists", "libs_exists"),
|
||||
[
|
||||
(True, True),
|
||||
(True, False),
|
||||
(False, True),
|
||||
],
|
||||
)
|
||||
async def test_run_pip_in_process_injects_windows_runtime_build_env_without_existing_paths(
|
||||
configure_run_pip_in_process_capture, include_exists, libs_exists
|
||||
):
|
||||
existing_runtime_dirs = set()
|
||||
if include_exists:
|
||||
existing_runtime_dirs.add(WINDOWS_RUNTIME_INCLUDE_DIR)
|
||||
if libs_exists:
|
||||
existing_runtime_dirs.add(WINDOWS_RUNTIME_LIBS_DIR)
|
||||
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="win32",
|
||||
packaged_runtime=True,
|
||||
existing_runtime_dirs=existing_runtime_dirs,
|
||||
)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
assert observed_env == {
|
||||
"INCLUDE": WINDOWS_RUNTIME_INCLUDE_DIR if include_exists else None,
|
||||
"LIB": WINDOWS_RUNTIME_LIBS_DIR if libs_exists else None,
|
||||
}
|
||||
if include_exists:
|
||||
assert ";" not in observed_env["INCLUDE"]
|
||||
if libs_exists:
|
||||
assert ";" not in observed_env["LIB"]
|
||||
assert "INCLUDE" not in pip_installer_module.os.environ
|
||||
assert "LIB" not in pip_installer_module.os.environ
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_pip_in_process_does_not_inject_when_runtime_dirs_missing(
|
||||
configure_run_pip_in_process_capture,
|
||||
):
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="win32",
|
||||
packaged_runtime=True,
|
||||
include_value=EXISTING_WINDOWS_INCLUDE_DIR,
|
||||
lib_value=EXISTING_WINDOWS_LIB_DIR,
|
||||
existing_runtime_dirs=set(),
|
||||
)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
assert observed_env == {
|
||||
"INCLUDE": EXISTING_WINDOWS_INCLUDE_DIR,
|
||||
"LIB": EXISTING_WINDOWS_LIB_DIR,
|
||||
}
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == EXISTING_WINDOWS_INCLUDE_DIR
|
||||
assert pip_installer_module.os.environ["LIB"] == EXISTING_WINDOWS_LIB_DIR
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_pip_in_process_uses_latest_env_when_building_runtime_paths(
|
||||
monkeypatch,
|
||||
configure_run_pip_in_process_capture,
|
||||
):
|
||||
updated_include = ntpath.join(r"C:\new-toolchain", "include")
|
||||
updated_lib = ntpath.join(r"C:\new-toolchain", "lib")
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="win32",
|
||||
packaged_runtime=True,
|
||||
include_value=EXISTING_WINDOWS_INCLUDE_DIR,
|
||||
lib_value=EXISTING_WINDOWS_LIB_DIR,
|
||||
existing_runtime_dirs={
|
||||
WINDOWS_RUNTIME_INCLUDE_DIR,
|
||||
WINDOWS_RUNTIME_LIBS_DIR,
|
||||
},
|
||||
)
|
||||
|
||||
async def fake_to_thread(func, *args):
|
||||
pip_installer_module.os.environ["INCLUDE"] = updated_include
|
||||
pip_installer_module.os.environ["LIB"] = updated_lib
|
||||
return func(*args)
|
||||
|
||||
monkeypatch.setattr(pip_installer_module.asyncio, "to_thread", fake_to_thread)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
assert observed_env == {
|
||||
"INCLUDE": f"{WINDOWS_RUNTIME_INCLUDE_DIR};{updated_include}",
|
||||
"LIB": f"{WINDOWS_RUNTIME_LIBS_DIR};{updated_lib}",
|
||||
}
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == updated_include
|
||||
assert pip_installer_module.os.environ["LIB"] == updated_lib
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_pip_in_process_does_not_modify_env_on_non_windows(
|
||||
configure_run_pip_in_process_capture,
|
||||
):
|
||||
existing_include = "/toolchain/include"
|
||||
existing_lib = "/toolchain/lib"
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="linux",
|
||||
packaged_runtime=True,
|
||||
include_value=existing_include,
|
||||
lib_value=existing_lib,
|
||||
)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
assert observed_env == {"INCLUDE": existing_include, "LIB": existing_lib}
|
||||
assert pip_installer_module.os.environ["INCLUDE"] == existing_include
|
||||
assert pip_installer_module.os.environ["LIB"] == existing_lib
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_pip_in_process_does_not_inject_env_when_not_packaged(
|
||||
configure_run_pip_in_process_capture,
|
||||
):
|
||||
observed_env = configure_run_pip_in_process_capture(
|
||||
platform="win32",
|
||||
packaged_runtime=False,
|
||||
existing_runtime_dirs={
|
||||
WINDOWS_RUNTIME_INCLUDE_DIR,
|
||||
WINDOWS_RUNTIME_LIBS_DIR,
|
||||
},
|
||||
)
|
||||
|
||||
installer = PipInstaller("")
|
||||
result = await installer._run_pip_in_process(["install", "demo-package"])
|
||||
|
||||
assert result == 0
|
||||
assert observed_env == {"INCLUDE": None, "LIB": None}
|
||||
assert "INCLUDE" not in pip_installer_module.os.environ
|
||||
assert "LIB" not in pip_installer_module.os.environ
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_pip_in_process_classifies_nonstandard_conflict_output(monkeypatch):
|
||||
def fake_pip_main(args):
|
||||
|
||||
Reference in New Issue
Block a user