Spaces:

Cripple-Lee
/

Personal-Assistant-Server

Paused

Personal-Assistant-Server / src /agent /graph.py

Sync from GitHub via hub-sync

5d5a5b6 verified 3 months ago

4.84 kB

"""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")
)

	"""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")
	)