mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
fdc310
583 lines
22 KiB
Python
583 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
import abc
|
|
import typing
|
|
import time
|
|
|
|
from ...core import app
|
|
from ...api.http.context import ExecutionContext
|
|
from ...entity.persistence import model as persistence_model
|
|
from ...workspace.errors import WorkspaceInvariantError
|
|
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
|
from . import token
|
|
from . import reasoning
|
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
|
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
|
|
|
|
|
LLM_USAGE_QUERY_VARIABLE = '_llm_usage'
|
|
STREAM_USAGE_QUERY_VARIABLE = '_stream_usage'
|
|
|
|
|
|
def _ensure_same_execution_scope(
|
|
expected: ExecutionContext,
|
|
actual: ExecutionContext,
|
|
*,
|
|
resource: str,
|
|
) -> None:
|
|
if (
|
|
actual.instance_uuid != expected.instance_uuid
|
|
or actual.workspace_uuid != expected.workspace_uuid
|
|
or actual.placement_generation != expected.placement_generation
|
|
):
|
|
raise WorkspaceInvariantError(f'{resource} belongs to another Workspace execution scope')
|
|
|
|
|
|
def _store_llm_usage(query: pipeline_query.Query | None, usage_info: dict | None) -> None:
|
|
"""Store the latest provider usage on the query for upstream action handlers."""
|
|
if query is None or not usage_info:
|
|
return
|
|
if query.variables is None:
|
|
query.variables = {}
|
|
query.variables[LLM_USAGE_QUERY_VARIABLE] = dict(usage_info)
|
|
|
|
|
|
class RuntimeProvider:
|
|
"""运行时模型提供商"""
|
|
|
|
provider_entity: persistence_model.ModelProvider
|
|
"""提供商数据"""
|
|
|
|
token_mgr: token.TokenManager
|
|
"""api key管理器"""
|
|
|
|
requester: ProviderAPIRequester
|
|
"""请求器实例"""
|
|
|
|
def __init__(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
provider_entity: persistence_model.ModelProvider,
|
|
token_mgr: token.TokenManager,
|
|
requester: ProviderAPIRequester,
|
|
):
|
|
if provider_entity.workspace_uuid != execution_context.workspace_uuid:
|
|
raise WorkspaceInvariantError('Provider belongs to another Workspace')
|
|
self.execution_context = execution_context
|
|
self.provider_entity = provider_entity
|
|
self.token_mgr = token_mgr
|
|
self.requester = requester
|
|
|
|
def _validate_invocation(
|
|
self,
|
|
model: RuntimeLLMModel | RuntimeEmbeddingModel | RuntimeRerankModel,
|
|
execution_context: ExecutionContext,
|
|
) -> None:
|
|
_ensure_same_execution_scope(self.execution_context, execution_context, resource='Provider invocation')
|
|
_ensure_same_execution_scope(self.execution_context, model.execution_context, resource='Runtime model')
|
|
if model.provider is not self:
|
|
raise WorkspaceInvariantError('Runtime model is attached to another provider')
|
|
|
|
def _resolve_llm_execution_context(
|
|
self,
|
|
query: pipeline_query.Query | None,
|
|
execution_context: ExecutionContext | None,
|
|
) -> ExecutionContext:
|
|
if query is not None:
|
|
from ...pipeline.pool import get_query_execution_context
|
|
|
|
query_context = get_query_execution_context(query)
|
|
if execution_context is not None:
|
|
_ensure_same_execution_scope(
|
|
query_context,
|
|
execution_context,
|
|
resource='Explicit LLM invocation context',
|
|
)
|
|
return query_context
|
|
if execution_context is None:
|
|
raise WorkspaceInvariantError('LLM invocation requires an ExecutionContext when query is absent')
|
|
return execution_context
|
|
|
|
async def invoke_llm(
|
|
self,
|
|
query: pipeline_query.Query | None,
|
|
model: RuntimeLLMModel,
|
|
messages: typing.List[provider_message.Message],
|
|
funcs: typing.List[resource_tool.LLMTool] = None,
|
|
extra_args: dict[str, typing.Any] = {},
|
|
remove_think: bool = False,
|
|
execution_context: ExecutionContext | None = None,
|
|
) -> provider_message.Message:
|
|
"""Bridge method for invoking LLM with monitoring"""
|
|
invocation_context = self._resolve_llm_execution_context(query, execution_context)
|
|
self._validate_invocation(model, invocation_context)
|
|
# Start timing for monitoring
|
|
start_time = time.time()
|
|
input_tokens = 0
|
|
output_tokens = 0
|
|
status = 'success'
|
|
error_message = None
|
|
|
|
try:
|
|
# Call the underlying requester
|
|
result = await self.requester.invoke_llm(
|
|
query=query,
|
|
model=model,
|
|
messages=messages,
|
|
funcs=funcs,
|
|
extra_args=extra_args,
|
|
remove_think=remove_think,
|
|
)
|
|
|
|
# Try to extract token usage if the requester returns it
|
|
# For requesters that return tuple (message, usage_info)
|
|
if isinstance(result, tuple):
|
|
msg, usage_info = result
|
|
if usage_info:
|
|
_store_llm_usage(query, usage_info)
|
|
input_tokens = usage_info.get('prompt_tokens', 0)
|
|
output_tokens = usage_info.get('completion_tokens', 0)
|
|
return msg
|
|
else:
|
|
return result
|
|
|
|
except Exception as e:
|
|
status = 'error'
|
|
error_message = str(e)
|
|
raise
|
|
finally:
|
|
# Record LLM call monitoring data (only if query is provided)
|
|
if query is not None:
|
|
duration_ms = int((time.time() - start_time) * 1000)
|
|
|
|
# Import monitoring helper
|
|
try:
|
|
from ...pipeline import monitoring_helper
|
|
|
|
# Get monitoring metadata from query variables
|
|
if query.variables:
|
|
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
|
|
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
|
|
message_id = query.variables.get('_monitoring_message_id')
|
|
else:
|
|
bot_name = 'Unknown'
|
|
pipeline_name = 'Unknown'
|
|
message_id = None
|
|
|
|
await monitoring_helper.MonitoringHelper.record_llm_call(
|
|
ap=self.requester.ap,
|
|
query=query,
|
|
bot_id=query.bot_uuid or 'unknown',
|
|
bot_name=bot_name,
|
|
pipeline_id=query.pipeline_uuid or 'unknown',
|
|
pipeline_name=pipeline_name,
|
|
model_name=model.model_entity.name,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
duration_ms=duration_ms,
|
|
status=status,
|
|
error_message=error_message,
|
|
message_id=message_id,
|
|
)
|
|
except Exception as monitor_err:
|
|
self.requester.ap.logger.error(f'[Monitoring] Failed to record LLM call: {monitor_err}')
|
|
|
|
async def invoke_llm_stream(
|
|
self,
|
|
query: pipeline_query.Query | None,
|
|
model: RuntimeLLMModel,
|
|
messages: typing.List[provider_message.Message],
|
|
funcs: typing.List[resource_tool.LLMTool] = None,
|
|
extra_args: dict[str, typing.Any] = {},
|
|
remove_think: bool = False,
|
|
execution_context: ExecutionContext | None = None,
|
|
) -> provider_message.MessageChunk:
|
|
"""Bridge method for invoking LLM stream with monitoring"""
|
|
invocation_context = self._resolve_llm_execution_context(query, execution_context)
|
|
self._validate_invocation(model, invocation_context)
|
|
# Start timing for monitoring
|
|
start_time = time.time()
|
|
status = 'success'
|
|
error_message = None
|
|
input_tokens = 0
|
|
output_tokens = 0
|
|
|
|
try:
|
|
# Stream the response
|
|
async for chunk in self.requester.invoke_llm_stream(
|
|
query=query,
|
|
model=model,
|
|
messages=messages,
|
|
funcs=funcs,
|
|
extra_args=extra_args,
|
|
remove_think=remove_think,
|
|
):
|
|
yield chunk
|
|
# Extract usage from stream if available (stored by LiteLLM requester)
|
|
if query:
|
|
if query.variables is None:
|
|
query.variables = {}
|
|
if STREAM_USAGE_QUERY_VARIABLE in query.variables:
|
|
usage_info = query.variables[STREAM_USAGE_QUERY_VARIABLE]
|
|
_store_llm_usage(query, usage_info)
|
|
input_tokens = usage_info.get('prompt_tokens', 0)
|
|
output_tokens = usage_info.get('completion_tokens', 0)
|
|
del query.variables[STREAM_USAGE_QUERY_VARIABLE]
|
|
except Exception as e:
|
|
status = 'error'
|
|
error_message = str(e)
|
|
raise
|
|
finally:
|
|
# Record LLM call monitoring data (only if query is provided)
|
|
if query is not None:
|
|
duration_ms = int((time.time() - start_time) * 1000)
|
|
|
|
# Import monitoring helper
|
|
try:
|
|
from ...pipeline import monitoring_helper
|
|
|
|
# Get monitoring metadata from query variables
|
|
if query.variables:
|
|
bot_name = query.variables.get('_monitoring_bot_name', 'Unknown')
|
|
pipeline_name = query.variables.get('_monitoring_pipeline_name', 'Unknown')
|
|
message_id = query.variables.get('_monitoring_message_id')
|
|
else:
|
|
bot_name = 'Unknown'
|
|
pipeline_name = 'Unknown'
|
|
message_id = None
|
|
|
|
await monitoring_helper.MonitoringHelper.record_llm_call(
|
|
ap=self.requester.ap,
|
|
query=query,
|
|
bot_id=query.bot_uuid or 'unknown',
|
|
bot_name=bot_name,
|
|
pipeline_id=query.pipeline_uuid or 'unknown',
|
|
pipeline_name=pipeline_name,
|
|
model_name=model.model_entity.name,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
duration_ms=duration_ms,
|
|
status=status,
|
|
error_message=error_message,
|
|
message_id=message_id,
|
|
)
|
|
except Exception as monitor_err:
|
|
self.requester.ap.logger.error(f'[Monitoring] Failed to record LLM stream call: {monitor_err}')
|
|
|
|
async def invoke_embedding(
|
|
self,
|
|
model: RuntimeEmbeddingModel,
|
|
input_text: typing.List[str],
|
|
extra_args: dict[str, typing.Any] = {},
|
|
*,
|
|
execution_context: ExecutionContext,
|
|
knowledge_base_id: str | None = None,
|
|
query_text: str | None = None,
|
|
session_id: str | None = None,
|
|
message_id: str | None = None,
|
|
call_type: str | None = None,
|
|
) -> typing.List[typing.List[float]]:
|
|
"""Bridge method for invoking embedding with monitoring"""
|
|
self._validate_invocation(model, execution_context)
|
|
# Start timing for monitoring
|
|
start_time = time.time()
|
|
prompt_tokens = 0
|
|
total_tokens = 0
|
|
status = 'success'
|
|
error_message = None
|
|
|
|
try:
|
|
# Call the underlying requester
|
|
result = await self.requester.invoke_embedding(
|
|
model=model,
|
|
input_text=input_text,
|
|
extra_args=extra_args,
|
|
)
|
|
|
|
# Handle both old format (list only) and new format (tuple with usage)
|
|
if isinstance(result, tuple):
|
|
embeddings, usage_info = result
|
|
if usage_info:
|
|
prompt_tokens = usage_info.get('prompt_tokens', 0)
|
|
total_tokens = usage_info.get('total_tokens', 0)
|
|
return embeddings
|
|
else:
|
|
return result
|
|
|
|
except Exception as e:
|
|
status = 'error'
|
|
error_message = str(e)
|
|
raise
|
|
finally:
|
|
# Record embedding call monitoring data
|
|
duration_ms = int((time.time() - start_time) * 1000)
|
|
|
|
try:
|
|
await self.requester.ap.monitoring_service.record_embedding_call(
|
|
execution_context,
|
|
model_name=model.model_entity.name,
|
|
prompt_tokens=prompt_tokens,
|
|
total_tokens=total_tokens,
|
|
duration=duration_ms,
|
|
input_count=len(input_text),
|
|
status=status,
|
|
error_message=error_message,
|
|
knowledge_base_id=knowledge_base_id,
|
|
query_text=query_text,
|
|
session_id=session_id,
|
|
message_id=message_id,
|
|
call_type=call_type,
|
|
)
|
|
except Exception as monitor_err:
|
|
self.requester.ap.logger.error(f'[Monitoring] Failed to record embedding call: {monitor_err}')
|
|
|
|
async def invoke_rerank(
|
|
self,
|
|
model: RuntimeRerankModel,
|
|
query: str,
|
|
documents: typing.List[str],
|
|
extra_args: dict[str, typing.Any] = {},
|
|
*,
|
|
execution_context: ExecutionContext,
|
|
) -> typing.List[dict]:
|
|
"""Bridge method for invoking rerank with monitoring"""
|
|
self._validate_invocation(model, execution_context)
|
|
start_time = time.time()
|
|
status = 'success'
|
|
|
|
try:
|
|
result = await self.requester.invoke_rerank(
|
|
model=model,
|
|
query=query,
|
|
documents=documents,
|
|
extra_args=extra_args,
|
|
)
|
|
return result
|
|
|
|
except Exception:
|
|
status = 'error'
|
|
raise
|
|
finally:
|
|
duration_ms = int((time.time() - start_time) * 1000)
|
|
|
|
try:
|
|
self.requester.ap.logger.debug(
|
|
f'[Rerank] model={model.model_entity.name} docs={len(documents)} '
|
|
f'duration={duration_ms}ms status={status}'
|
|
)
|
|
except Exception as monitor_err:
|
|
self.requester.ap.logger.error(f'[Monitoring] Failed to record rerank call: {monitor_err}')
|
|
|
|
|
|
class RuntimeLLMModel:
|
|
"""运行时模型"""
|
|
|
|
model_entity: persistence_model.LLMModel
|
|
"""模型数据"""
|
|
|
|
provider: RuntimeProvider
|
|
"""提供商实例"""
|
|
|
|
reasoning_config_override: dict[str, str] | None
|
|
"""Request-scoped reasoning policy supplied by the active pipeline."""
|
|
|
|
def __init__(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
model_entity: persistence_model.LLMModel,
|
|
provider: RuntimeProvider,
|
|
reasoning_config_override: dict[str, str] | None = None,
|
|
):
|
|
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='LLM model')
|
|
if model_entity.workspace_uuid != execution_context.workspace_uuid:
|
|
raise WorkspaceInvariantError('LLM model belongs to another Workspace')
|
|
if model_entity.provider_uuid != provider.provider_entity.uuid:
|
|
raise WorkspaceInvariantError('LLM model references another provider')
|
|
self.execution_context = execution_context
|
|
self.model_entity = model_entity
|
|
self.provider = provider
|
|
self.reasoning_config_override = reasoning_config_override
|
|
|
|
|
|
class RuntimeEmbeddingModel:
|
|
"""运行时 Embedding 模型"""
|
|
|
|
model_entity: persistence_model.EmbeddingModel
|
|
"""模型数据"""
|
|
|
|
provider: RuntimeProvider
|
|
"""提供商实例"""
|
|
|
|
def __init__(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
model_entity: persistence_model.EmbeddingModel,
|
|
provider: RuntimeProvider,
|
|
):
|
|
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='Embedding model')
|
|
if model_entity.workspace_uuid != execution_context.workspace_uuid:
|
|
raise WorkspaceInvariantError('Embedding model belongs to another Workspace')
|
|
if model_entity.provider_uuid != provider.provider_entity.uuid:
|
|
raise WorkspaceInvariantError('Embedding model references another provider')
|
|
self.execution_context = execution_context
|
|
self.model_entity = model_entity
|
|
self.provider = provider
|
|
|
|
|
|
class RuntimeRerankModel:
|
|
"""运行时 Rerank 模型"""
|
|
|
|
model_entity: persistence_model.RerankModel
|
|
"""模型数据"""
|
|
|
|
provider: RuntimeProvider
|
|
"""提供商实例"""
|
|
|
|
def __init__(
|
|
self,
|
|
execution_context: ExecutionContext,
|
|
model_entity: persistence_model.RerankModel,
|
|
provider: RuntimeProvider,
|
|
):
|
|
_ensure_same_execution_scope(provider.execution_context, execution_context, resource='Rerank model')
|
|
if model_entity.workspace_uuid != execution_context.workspace_uuid:
|
|
raise WorkspaceInvariantError('Rerank model belongs to another Workspace')
|
|
if model_entity.provider_uuid != provider.provider_entity.uuid:
|
|
raise WorkspaceInvariantError('Rerank model references another provider')
|
|
self.execution_context = execution_context
|
|
self.model_entity = model_entity
|
|
self.provider = provider
|
|
|
|
|
|
class ProviderAPIRequester(metaclass=abc.ABCMeta):
|
|
"""Provider API请求器"""
|
|
|
|
name: str = None
|
|
init_api_key: str = 'langbot-init-placeholder'
|
|
|
|
ap: app.Application
|
|
|
|
default_config: dict[str, typing.Any] = {}
|
|
|
|
requester_cfg: dict[str, typing.Any] = {}
|
|
|
|
def __init__(self, ap: app.Application, config: dict[str, typing.Any]):
|
|
self.ap = ap
|
|
self.requester_cfg = {**self.default_config}
|
|
self.requester_cfg.update(config)
|
|
|
|
async def initialize(self):
|
|
pass
|
|
|
|
async def aclose(self) -> None:
|
|
"""Release requester-owned clients when its runtime provider retires.
|
|
|
|
Most built-in requesters are currently stateless, but provider
|
|
extensions may own connection pools or background resources. Keeping
|
|
the lifecycle hook on the base class lets Workspace generation changes
|
|
and application shutdown retire them deterministically.
|
|
"""
|
|
|
|
return None
|
|
|
|
async def scan_models(self, api_key: str | None = None) -> dict[str, typing.Any] | list[dict[str, typing.Any]]:
|
|
"""Scan models supported by the provider.
|
|
|
|
The default implementation does not support scanning. Requesters that
|
|
can enumerate remote models should override this method.
|
|
"""
|
|
raise NotImplementedError('This provider does not support model scanning')
|
|
|
|
def get_reasoning_capabilities(self, model: RuntimeLLMModel) -> dict[str, typing.Any]:
|
|
"""Return normalized reasoning controls supported by a model."""
|
|
return reasoning.default_reasoning_capabilities(
|
|
supported='reasoning' in (model.model_entity.abilities or []),
|
|
source='manual' if 'reasoning' in (model.model_entity.abilities or []) else 'unknown',
|
|
)
|
|
|
|
@abc.abstractmethod
|
|
async def invoke_llm(
|
|
self,
|
|
query: pipeline_query.Query,
|
|
model: RuntimeLLMModel,
|
|
messages: typing.List[provider_message.Message],
|
|
funcs: typing.List[resource_tool.LLMTool] = None,
|
|
extra_args: dict[str, typing.Any] = {},
|
|
remove_think: bool = False,
|
|
) -> provider_message.Message:
|
|
"""调用API
|
|
|
|
Args:
|
|
model (RuntimeLLMModel): 使用的模型信息
|
|
messages (typing.List[llm_entities.Message]): 消息对象列表
|
|
funcs (typing.List[tools_entities.LLMFunction], optional): 使用的工具函数列表. Defaults to None.
|
|
extra_args (dict[str, typing.Any], optional): 额外的参数. Defaults to {}.
|
|
remove_think (bool, optional): 是否移思考中的消息. Defaults to False.
|
|
|
|
Returns:
|
|
llm_entities.Message: 返回消息对象
|
|
"""
|
|
pass
|
|
|
|
async def invoke_llm_stream(
|
|
self,
|
|
query: pipeline_query.Query,
|
|
model: RuntimeLLMModel,
|
|
messages: typing.List[provider_message.Message],
|
|
funcs: typing.List[resource_tool.LLMTool] = None,
|
|
extra_args: dict[str, typing.Any] = {},
|
|
remove_think: bool = False,
|
|
) -> provider_message.MessageChunk:
|
|
"""调用API
|
|
|
|
Args:
|
|
model (RuntimeLLMModel): 使用的模型信息
|
|
messages (typing.List[provider_message.Message]): 消息对象列表
|
|
funcs (typing.List[resource_tool.LLMTool], optional): 使用的工具函数列表. Defaults to None.
|
|
extra_args (dict[str, typing.Any], optional): 额外的参数. Defaults to {}.
|
|
remove_think (bool, optional): 是否移除思考中的消息. Defaults to False.
|
|
|
|
Returns:
|
|
typing.AsyncGenerator[provider_message.MessageChunk]: 返回消息对象
|
|
"""
|
|
pass
|
|
|
|
async def invoke_embedding(
|
|
self,
|
|
model: RuntimeEmbeddingModel,
|
|
input_text: typing.List[str],
|
|
extra_args: dict[str, typing.Any] = {},
|
|
) -> typing.Union[typing.List[typing.List[float]], tuple[typing.List[typing.List[float]], dict]]:
|
|
"""调用 Embedding API
|
|
|
|
Args:
|
|
model (RuntimeEmbeddingModel): 使用的模型信息
|
|
input_text (typing.List[str]): 输入文本
|
|
extra_args (dict[str, typing.Any], optional): 额外的参数. Defaults to {}.
|
|
|
|
Returns:
|
|
typing.List[typing.List[float]]: 返回的 embedding 向量
|
|
或者 tuple[typing.List[typing.List[float]], dict]: 返回 (embedding 向量, usage_info)
|
|
"""
|
|
pass
|
|
|
|
async def invoke_rerank(
|
|
self,
|
|
model: RuntimeRerankModel,
|
|
query: str,
|
|
documents: typing.List[str],
|
|
extra_args: dict[str, typing.Any] = {},
|
|
) -> typing.List[dict]:
|
|
"""调用 Rerank API
|
|
|
|
Args:
|
|
model (RuntimeRerankModel): 使用的模型信息
|
|
query (str): 查询文本
|
|
documents (typing.List[str]): 待重排序的文档列表
|
|
extra_args (dict[str, typing.Any], optional): 额外的参数. Defaults to {}.
|
|
|
|
Returns:
|
|
typing.List[dict]: [{"index": int, "relevance_score": float}, ...]
|
|
"""
|
|
raise NotImplementedError('This requester does not support rerank')
|