mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-15 23:10:59 +00:00
Junyan Qin
124 lines
5.1 KiB
Python
124 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import traceback
|
|
|
|
from ..core import app
|
|
from ..core import entities as core_entities
|
|
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError
|
|
|
|
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
|
from .pool import get_query_execution_context
|
|
|
|
|
|
class Controller:
|
|
"""总控制器"""
|
|
|
|
ap: app.Application
|
|
|
|
semaphore: asyncio.Semaphore = None
|
|
"""请求并发控制信号量"""
|
|
|
|
def __init__(self, ap: app.Application):
|
|
self.ap = ap
|
|
self.semaphore = asyncio.Semaphore(self.ap.instance_config.data['concurrency']['pipeline'])
|
|
|
|
async def _assert_query_execution_active(
|
|
self,
|
|
query: pipeline_query.Query,
|
|
):
|
|
"""Revalidate a queued query immediately before runtime work starts."""
|
|
|
|
execution_context = get_query_execution_context(query)
|
|
binding = await self.ap.workspace_service.get_execution_binding(
|
|
execution_context.workspace_uuid,
|
|
expected_generation=execution_context.placement_generation,
|
|
)
|
|
if binding.instance_uuid != execution_context.instance_uuid:
|
|
raise WorkspaceInvariantError('Queued query instance does not match the active Workspace binding')
|
|
return execution_context
|
|
|
|
async def _process_query(self, selected_query: pipeline_query.Query) -> None:
|
|
"""Run one selected query and always release its scheduling slot."""
|
|
|
|
try:
|
|
async with self.semaphore:
|
|
execution_context = await self._assert_query_execution_active(selected_query)
|
|
pipeline_uuid = selected_query.pipeline_uuid
|
|
|
|
if pipeline_uuid:
|
|
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
|
|
execution_context,
|
|
pipeline_uuid,
|
|
)
|
|
if pipeline:
|
|
await pipeline.run(selected_query)
|
|
else:
|
|
self.ap.logger.warning(
|
|
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
|
|
)
|
|
else:
|
|
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
|
|
except WorkspaceError as exc:
|
|
self.ap.logger.info(
|
|
f'Dropped query {selected_query.query_id} because its Workspace execution binding is stale: {exc}'
|
|
)
|
|
finally:
|
|
await self.ap.query_pool.remove_query(selected_query)
|
|
async with self.ap.query_pool:
|
|
(await self.ap.sess_mgr.get_session(selected_query))._semaphore.release()
|
|
self.ap.query_pool.condition.notify_all()
|
|
|
|
async def consumer(self):
|
|
"""事件处理循环"""
|
|
try:
|
|
while True:
|
|
selected_query: pipeline_query.Query = None
|
|
|
|
# 取请求
|
|
async with self.ap.query_pool:
|
|
queries: list[pipeline_query.Query] = self.ap.query_pool.queries
|
|
|
|
for query in queries:
|
|
session = await self.ap.sess_mgr.get_session(query)
|
|
# Debug logging removed from tight loop to prevent excessive log generation
|
|
# that can cause memory overflow in high-traffic scenarios
|
|
|
|
if not session._semaphore.locked():
|
|
selected_query = query
|
|
await session._semaphore.acquire()
|
|
# Only log when actually selecting a query
|
|
self.ap.logger.debug(f'Selected query {query.query_id} for processing')
|
|
|
|
break
|
|
|
|
if selected_query: # 找到了
|
|
queries.remove(selected_query)
|
|
else: # 没找到 说明：没有请求 或者 所有query对应的session都已达到并发上限
|
|
await self.ap.query_pool.condition.wait()
|
|
continue
|
|
|
|
if selected_query:
|
|
execution_context = get_query_execution_context(selected_query)
|
|
self.ap.task_mgr.create_task(
|
|
self._process_query(selected_query),
|
|
kind='query',
|
|
name=f'query-{selected_query.query_id}',
|
|
scopes=[
|
|
core_entities.LifecycleControlScope.APPLICATION,
|
|
core_entities.LifecycleControlScope.PLATFORM,
|
|
],
|
|
instance_uuid=execution_context.instance_uuid,
|
|
workspace_uuid=execution_context.workspace_uuid,
|
|
placement_generation=execution_context.placement_generation,
|
|
)
|
|
|
|
except Exception as e:
|
|
# traceback.print_exc()
|
|
self.ap.logger.error(f'控制器循环出错: {e}')
|
|
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
|
|
|
|
async def run(self):
|
|
"""运行控制器"""
|
|
await self.consumer()
|