File size: 4,836 Bytes
bcddfe2 59f9af0 bcddfe2 59f9af0 5d5a5b6 bcddfe2 4156529 bcddfe2 9f1c9d5 5d5a5b6 9f1c9d5 bcddfe2 9f1c9d5 bcddfe2 9f1c9d5 4156529 8194821 bcddfe2 d79d2fc bcddfe2 84314bc 45831fa 5d5a5b6 84314bc 45831fa 84314bc 45831fa 5d5a5b6 45831fa bcddfe2 5d5a5b6 d79d2fc bcddfe2 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """LangGraph single-node graph template.
Returns a predefined response. Replace logic and configuration as needed.
"""
from __future__ import annotations
from typing import Any, Dict
from utils.state import State
from utils.context import Context
from langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph
from langgraph.runtime import Runtime
import logging
from web.login import get_or_create_agent
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
async def call_model(
state: State,
runtime: Runtime[Context],
config: RunnableConfig,
) -> Dict[str, Any]:
"""Process input and returns output.
The supervisor agent is a shared singleton built with the server-wide
``HF_TOKEN`` env var (see ``web/login.py``).
"""
# Reuse the shared supervisor instance.
supervisor_agent = await get_or_create_agent("supervisor")
logger.info(f"State:{state}")
messages = state["messages"]
content = ""
# Depth of nested tool execution. While > 0 we are inside a subagent
# (navigate/search) tool call, and any chat-model chunks belong to the
# subagent — NOT to the supervisor's own answer. Tracking this prevents
# subagent output from being interleaved/duplicated into the final text.
tool_depth = 0
# Token accounting: sum usage_metadata from EVERY chat-model chunk —
# including subagent chunks — so the user is billed for all model calls
# their run triggers (supervisor + navigate/search subagents).
tokens_used = 0
async for stream in supervisor_agent.astream_events(
{"messages": messages},
version="v2",
):
logger.debug(f"Stream:{stream}")
event = stream['event']
if event == 'on_tool_start':
tool_depth += 1
continue
if event == 'on_tool_end':
tool_depth = max(0, tool_depth - 1)
continue
if event != 'on_chat_model_stream':
continue
chunk = stream['data']['chunk']
# Accumulate token usage before any filtering so subagent and
# tool-call chunks are billed too.
usage = getattr(chunk, "usage_metadata", None)
if usage:
tokens_used += usage.get("total_tokens") or (
(usage.get("input_tokens") or 0) + (usage.get("output_tokens") or 0)
)
# Ignore chunks produced by subagents running inside a tool call.
if tool_depth > 0:
continue
# Skip chunks that are part of a tool call. Qwen emits raw tool-call
# markup (e.g. `<function=...>` / `</function>`) as content tokens;
# these carry tool-call metadata and must not be treated as
# user-facing text.
if getattr(chunk, "additional_kwargs", {}).get("tool_calls"):
continue
text = chunk.content
if not text:
continue
# Defensive filter: drop any residual tool-markup fragments that
# leaked into plain content (the source of the `se}</function>`
# artifacts).
if any(tok in text for tok in ("<function", "</function", "<tool", "</tool")):
continue
content += text
logger.info(f"Response:{content}")
# Bill the run's tokens to the authenticated user. The auth handler
# (web/auth.py) injects the user into configurable before execution.
# Aegra passes it as a User object (attribute access), while LangGraph
# Platform passes a plain dict — handle both.
if tokens_used > 0:
user = (config.get("configurable") or {}).get("langgraph_auth_user")
if user is not None and not isinstance(user, dict):
user_id = getattr(user, "identity", None)
else:
user_id = (user or {}).get("identity")
# Skip billing for unauthenticated/anonymous runs (no DB user row).
if user_id and user_id != "anonymous":
try:
from web import db # local import to avoid a cycle at module load
total = await db.add_token_usage(user_id, tokens_used)
logger.info(
f"Recorded {tokens_used} tokens for user {user_id} "
f"(monthly total: {total})"
)
except Exception:
# Never fail a completed run because accounting failed.
logger.exception("Failed to record token usage")
# Return only the new assistant message; the add_messages reducer appends
# it to the persisted thread history.
return {"messages": [{"role": "assistant", "content": content}]}
# Define the graph
graph = (
StateGraph(State, context_schema=Context)
.add_node(call_model)
.add_edge("__start__", "call_model")
.compile(name="New Graph")
)
|