mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-12 09:12:40 +02:00
92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
"""Agent setup helpers for orchestrated chat streaming."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
_PREFLIGHT_TIMEOUT_SEC: float = 2.5
|
|
_PREFLIGHT_MAX_TOKENS: int = 1
|
|
|
|
|
|
async def preflight_llm(
|
|
llm: Any,
|
|
*,
|
|
is_provider_rate_limited: Callable[[BaseException], bool],
|
|
) -> None:
|
|
"""Issue a minimal completion probe to catch immediate provider 429s."""
|
|
from litellm import acompletion
|
|
|
|
model = getattr(llm, "model", None)
|
|
if not model or model == "auto":
|
|
return
|
|
|
|
try:
|
|
await acompletion(
|
|
model=model,
|
|
messages=[{"role": "user", "content": "ping"}],
|
|
api_key=getattr(llm, "api_key", None),
|
|
api_base=getattr(llm, "api_base", None),
|
|
max_tokens=_PREFLIGHT_MAX_TOKENS,
|
|
timeout=_PREFLIGHT_TIMEOUT_SEC,
|
|
stream=False,
|
|
metadata={"tags": ["surfsense:internal", "auto-pin-preflight"]},
|
|
)
|
|
except Exception as exc:
|
|
if is_provider_rate_limited(exc):
|
|
raise
|
|
logging.getLogger(__name__).debug(
|
|
"auto_pin_preflight non_rate_limit_error model=%s err=%s",
|
|
model,
|
|
exc,
|
|
)
|
|
|
|
|
|
async def build_main_agent_for_thread(
|
|
agent_factory: Any,
|
|
*,
|
|
llm: Any,
|
|
search_space_id: int,
|
|
db_session: Any,
|
|
connector_service: Any,
|
|
checkpointer: Any,
|
|
user_id: str | None,
|
|
thread_id: int | None,
|
|
agent_config: Any,
|
|
firecrawl_api_key: str | None,
|
|
thread_visibility: Any,
|
|
filesystem_selection: Any,
|
|
disabled_tools: list[str] | None = None,
|
|
mentioned_document_ids: list[int] | None = None,
|
|
) -> Any:
|
|
"""Run one canonical agent-build call for a single thread."""
|
|
return await agent_factory(
|
|
llm=llm,
|
|
search_space_id=search_space_id,
|
|
db_session=db_session,
|
|
connector_service=connector_service,
|
|
checkpointer=checkpointer,
|
|
user_id=user_id,
|
|
thread_id=thread_id,
|
|
agent_config=agent_config,
|
|
firecrawl_api_key=firecrawl_api_key,
|
|
thread_visibility=thread_visibility,
|
|
filesystem_selection=filesystem_selection,
|
|
disabled_tools=disabled_tools,
|
|
mentioned_document_ids=mentioned_document_ids,
|
|
)
|
|
|
|
|
|
async def settle_speculative_agent_build(task: Any) -> None:
|
|
"""Wait for a discarded speculative build and swallow its outcome."""
|
|
with contextlib.suppress(BaseException):
|
|
await task
|
|
|
|
|
|
__all__ = [
|
|
"build_main_agent_for_thread",
|
|
"preflight_llm",
|
|
"settle_speculative_agent_build",
|
|
]
|