File size: 1,499 Bytes
732b14f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
"""Shared LangChain LCEL invoke helpers (async + thread-pool sync)."""

from __future__ import annotations

from typing import Any

from app.async_executor import run_sync_in_executor
from app.config import settings
from app.llm.llm_throttle import throttled_llm_call


async def ainvoke_lcel_chain(
    chain: Any,
    variables: dict[str, Any],
    *,
    phase: str,
    section_id: str | None = None,
) -> str:
    """Run ``chain.ainvoke`` with throttling when the async pipeline is enabled."""

    async def _call() -> str:
        out = await chain.ainvoke(variables)
        return (out or "").strip() if isinstance(out, str) else str(out or "").strip()

    if settings.enable_async_pipeline:
        return await throttled_llm_call(
            phase=phase,
            section_id=section_id,
            cache_hit=None,
            call=_call,
        )
    return await _call()


async def invoke_lcel_chain(
    chain: Any,
    variables: dict[str, Any],
    *,
    phase: str,
    section_id: str | None = None,
) -> str:
    """Async-safe LCEL invoke: ``ainvoke`` on the async pipeline, else ``invoke`` in executor."""
    if settings.enable_async_pipeline:
        return await ainvoke_lcel_chain(
            chain, variables, phase=phase, section_id=section_id
        )

    def _sync_invoke() -> str:
        out = chain.invoke(variables)
        return (out or "").strip() if isinstance(out, str) else str(out or "").strip()

    return await run_sync_in_executor(_sync_invoke)