zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-16 04:37:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(agent): report accurate aggregated token usage and propagate session/user + input/output to Langfuse for agent runs (#16420)

Browse Source 
### What problem does this PR solve?

_Briefly describe what this PR aims to solve. Include background context
that will help reviewers understand the purpose of the PR._

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
- [x] Other (please describe):

## Summary

Agent (Canvas) runs previously did not surface token usage in the SSE
stream, and RAGFlow's own Langfuse generations for agent runs were
missing the prompt/completion split and the session/user correlation.
This made it impossible for an external caller (or Langfuse) to
reconcile an agent turn's cost with the upstream provider (e.g.
OpenRouter), because a single turn can issue several distinct LLM calls
(query rewriting / cross-language translation, multi-round tool
reasoning, nested sub-agents, and the final answer).

This PR introduces a per-run token usage sink so that **every** LLM call
in a run is aggregated and reported once, and enriches Langfuse
generations with the prompt/completion split plus session/user
attributes.

## What changes

### 1. Per-run token usage sink (`common/token_utils.py`)

- Adds two `contextvars`: `token_usage_sink` (a mutable per-run
accumulator) and `langfuse_run_attrs` (session_id/user_id for the run).
- Adds `record_run_token_usage(...)` (thread-safe via a lock, because
`thread_pool_exec` copies the context into worker threads that share the
sink dict) and `usage_from_response(...)` which extracts a
`{prompt_tokens, completion_tokens, total_tokens}` split from
OpenAI/OpenRouter-style responses.

### 2. Provider layer captures the prompt/completion split
(`rag/llm/chat_model.py`)

- `LiteLLMBase` and `Base` now store `self.last_usage`
(prompt/completion/total) for the most recent chat call, in both the
plain and tool-calling paths.
- Streaming requests set `stream_options.include_usage = True` (LiteLLM
path) so the authoritative usage arrives on the final chunk; this is
read even on the usage-only chunk that carries no `choices`.
- Fixes a multi-round accounting bug in `*_with_tools`: token totals
were **overwritten** by each round (`total_tokens = tol`) instead of
accumulated, undercounting multi-round tool conversations. Each round is
now committed to a running aggregate.

### 3. LLMBundle reports usage once, per call
(`api/db/services/llm_service.py`)

- New `_report_usage(total_tokens)` records the call's usage into the
active run sink and returns the prompt/completion/total split for
Langfuse. The split is only used when it is consistent with the
authoritative total; otherwise only the total is reported.
- All three chat entry points (`async_chat`, `async_chat_streamly`,
`async_chat_streamly_delta`) now emit `usage_details` with
`input`/`output`/`total` instead of total-only.
- `_start_langfuse_observation` now applies `session_id`/`user_id` from
the per-run context (`langfuse_run_attrs`) so agent-run generations are
correctly grouped, even though agent LLMBundles are constructed without
those attributes.

### 4. Canvas installs the sink and emits the aggregate
(`agent/canvas.py`)

- `Canvas.run()` installs a fresh `token_usage_sink` and
`langfuse_run_attrs` (from `user_id`/`session_id`) at the start of every
turn.
- `message_end` now includes an aggregated `usage` object:
`{prompt_tokens, completion_tokens, total_tokens, calls}` covering all
LLM calls in the run.

### 5. Pass session id into the run
(`api/db/services/canvas_service.py`)

- `completion()` forwards `session_id` to `Canvas.run()` for Langfuse
session correlation.

## Why a context variable

LLM calls in an agent run originate from many places that each build
their own `LLMBundle` (e.g. `cross_languages`/`keyword_extraction`
helpers, the Agent component, and nested sub-agents invoked as tools). A
run-scoped context variable is the only non-invasive chokepoint that
captures all of them exactly once, including nested agents (which run in
the same async context) and thread-pool tools (the executor copies the
context).

## Behavior / compatibility

- No public API or wire-format removal: `message_end` gains an
additional optional `usage` field; existing consumers are unaffected.
- When a provider does not return authoritative usage, behavior falls
back to the previous token estimate (total only, no split).
- Non-agent flows (Dataflow `Pipeline`, sync `Graph.run`) are untouched.

## Testing
- [x] Simple agent answer: `message_end.usage.total_tokens` matches
provider usage.
- [x] Agent with cross-language retrieval: aggregate equals the sum of
both provider calls.
- [x] Tool-calling agent (multi-round): total accumulates across rounds.
- [x] Nested agent (agent-as-tool): sub-agent tokens included in the
parent run total.
- [x] Langfuse: agent generations show input/output split and are
grouped by session/user.

---------

Co-authored-by: yzc <yuzhichang@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Öndery
2026-07-02 04:35:28 +03:00
committed by GitHub
parent 42a0faad18
commit 742188c3bb
6 changed files with 655 additions and 422 deletions
Download Patch File Download Diff File Expand all files Collapse all files

244
api/apps/restful_apis/agent_api.py
View File

@@ -88,29 +88,32 @@ def _canvas_json_default(obj):
def _require_canvas_access_sync(func):
@wraps(func)
def wrapper(*args, **kwargs):
if not UserCanvasService.accessible(kwargs.get('agent_id'), kwargs.get('tenant_id')):
if not UserCanvasService.accessible(kwargs.get("agent_id"), kwargs.get("tenant_id")):
return get_json_result(data=False, message="Make sure you have permission to access the agent.", code=RetCode.OPERATING_ERROR)
return func(*args, **kwargs)
return wrapper
def _require_canvas_access_async(func):
@wraps(func)
async def wrapper(*args, **kwargs):
agent_id = kwargs.get('agent_id')
tenant_id = kwargs.get('tenant_id')
agent_id = kwargs.get("agent_id")
tenant_id = kwargs.get("tenant_id")
if not await thread_pool_exec(UserCanvasService.accessible, agent_id, tenant_id):
return get_json_result(data=False, message="Make sure you have permission to access the agent.", code=RetCode.OPERATING_ERROR)
return await func(*args, **kwargs)
return wrapper
def _require_canvas_owner_sync(func):
@wraps(func)
def wrapper(*args, **kwargs):
if not UserCanvasService.query(user_id=kwargs.get('tenant_id'), id=kwargs.get('agent_id')):
if not UserCanvasService.query(user_id=kwargs.get("tenant_id"), id=kwargs.get("agent_id")):
return get_json_result(data=False, message="Only the owner of the agent is authorized for this operation.", code=RetCode.OPERATING_ERROR)
return func(*args, **kwargs)
return wrapper
@@ -261,9 +264,7 @@ async def _run_workflow_session(
if "chunks" in workflow_conv["reference"]:
workflow_conv["reference"] = [workflow_conv["reference"]]
else:
workflow_conv["reference"] = [
value for _, value in sorted(workflow_conv["reference"].items(), key=lambda item: int(item[0]))
]
workflow_conv["reference"] = [value for _, value in sorted(workflow_conv["reference"].items(), key=lambda item: int(item[0]))]
elif not isinstance(workflow_conv.get("reference"), list):
workflow_conv["reference"] = []
workflow_conv["reference"] = [_normalize_agent_reference_entry(reference) for reference in workflow_conv["reference"]]
@@ -344,22 +345,16 @@ async def _run_workflow_session(
# bare [DONE] (fixes #15169).
logging.info(
"empty agent output - returning session_id (agent_id=%s session_id=%s stream=%s)",
agent_id, session_id, True,
)
yield (
"data:"
+ json.dumps({"session_id": session_id, "data": {}}, ensure_ascii=False)
+ "\n\n"
agent_id,
session_id,
True,
)
yield ("data:" + json.dumps({"session_id": session_id, "data": {}}, ensure_ascii=False) + "\n\n")
await persist_workflow_session()
except Exception as exc:
logging.exception(exc)
canvas.cancel_task()
yield (
"data:"
+ json.dumps({"code": 500, "message": str(exc), "data": False}, ensure_ascii=False)
+ "\n\n"
)
yield ("data:" + json.dumps({"code": 500, "message": str(exc), "data": False}, ensure_ascii=False) + "\n\n")
finally:
if not done_sent:
done_sent = True
@@ -400,7 +395,9 @@ async def _run_workflow_session(
# (fixes #15169).
logging.info(
"empty agent output - returning session_id (agent_id=%s session_id=%s stream=%s)",
agent_id, session_id, False,
agent_id,
session_id,
False,
)
await commit_runtime_replica()
return get_result(data={"session_id": session_id})
@@ -559,16 +556,13 @@ async def delete_agent_session(tenant_id, agent_id):
if errors:
if success_count > 0:
return get_result(data={"success_count": success_count, "errors": errors},
message=f"Partially deleted {success_count} sessions with {len(errors)} errors")
return get_result(data={"success_count": success_count, "errors": errors}, message=f"Partially deleted {success_count} sessions with {len(errors)} errors")
else:
return get_error_data_result(message="; ".join(errors))
if duplicate_messages:
if success_count > 0:
return get_result(
message=f"Partially deleted {success_count} sessions with {len(duplicate_messages)} errors",
data={"success_count": success_count, "errors": duplicate_messages})
return get_result(message=f"Partially deleted {success_count} sessions with {len(duplicate_messages)} errors", data={"success_count": success_count, "errors": duplicate_messages})
else:
return get_error_data_result(message=";".join(duplicate_messages))
@@ -611,8 +605,8 @@ async def _iter_session_completion_events(tenant_id, agent_id, req, return_trace
yield ans
continue
if event in ["message", "message_end", "user_inputs", "workflow_finished"]:
if event in ["user_inputs", "workflow_finished"]:
if event in ["message", "message_end", "user_inputs"]:
if event == "user_inputs":
logging.debug(
"Forwarding session completion event: tenant_id=%s agent_id=%s event=%s",
tenant_id,
@@ -620,6 +614,22 @@ async def _iter_session_completion_events(tenant_id, agent_id, req, return_trace
event,
)
yield ans
continue
if event == "workflow_finished":
# Forward only the run-level aggregated token usage, not the whole terminal
# payload (inputs/outputs), so the session completion stream surface stays
# limited to what the usage contract needs.
logging.debug(
"Forwarding session completion event: tenant_id=%s agent_id=%s event=%s",
tenant_id,
agent_id,
event,
)
usage = ans.get("data", {}).get("usage")
if usage is not None:
yield {**ans, "data": {"usage": usage}}
continue
@manager.route("/agents/templates", methods=["GET"]) # noqa: F821
@@ -760,7 +770,7 @@ async def update_agent_tags(tenant_id, canvas_id):
@add_tenant_id_to_kwargs
async def create_agent(tenant_id):
req = {k: v for k, v in (await get_request_json()).items() if v is not None}
req["canvas_type"] = req.get("canvas_type","")
req["canvas_type"] = req.get("canvas_type", "")
req["user_id"] = tenant_id
req["canvas_category"] = req.get("canvas_category") or CanvasCategory.Agent
req["release"] = bool(req.get("release", ""))
@@ -837,13 +847,9 @@ async def upload_agent_file(agent_id, tenant_id):
)
try:
if len(file_objs) == 1:
uploaded = await thread_pool_exec(
FileService.upload_info, tenant_id, file_objs[0], request.args.get("url")
)
uploaded = await thread_pool_exec(FileService.upload_info, tenant_id, file_objs[0], request.args.get("url"))
return get_json_result(data=uploaded)
results = await asyncio.gather(
*(thread_pool_exec(FileService.upload_info, tenant_id, file_obj) for file_obj in file_objs)
)
results = await asyncio.gather(*(thread_pool_exec(FileService.upload_info, tenant_id, file_obj) for file_obj in file_objs))
return get_json_result(data=results)
except Exception as exc:
logging.exception(
@@ -1015,7 +1021,7 @@ def delete_agent(agent_id, tenant_id):
@_require_canvas_access_async
async def update_agent(agent_id, tenant_id):
req = {k: v for k, v in (await get_request_json()).items() if v is not None}
req["canvas_type"] = req.get("canvas_type","")
req["canvas_type"] = req.get("canvas_type", "")
req["release"] = bool(req.get("release", ""))
if req.get("dsl") is not None:
@@ -1038,10 +1044,7 @@ async def update_agent(agent_id, tenant_id):
return get_data_error_result(message=f"{req['title']} already exists.")
agent_title_for_version = req.get("title") or (current_agent.title if current_agent else "")
canvas_category = (
req.get("canvas_category")
or (current_agent.canvas_category if current_agent else CanvasCategory.Agent)
)
canvas_category = req.get("canvas_category") or (current_agent.canvas_category if current_agent else CanvasCategory.Agent)
owner_nickname = _get_user_nickname(tenant_id)
UserCanvasService.update_by_id(agent_id, req)
@@ -1153,13 +1156,19 @@ async def test_db_connection():
except ValueError as exc:
logging.warning(
"Rejected test_db_connection: unsafe host %r (db_type=%s, user=%s): %s",
req.get("host"), req.get("db_type"), current_user.id, exc,
req.get("host"),
req.get("db_type"),
current_user.id,
exc,
)
return get_data_error_result(message=str(exc))
except OSError as exc:
logging.warning(
"Rejected test_db_connection: cannot resolve host %r (db_type=%s, user=%s): %s",
req.get("host"), req.get("db_type"), current_user.id, exc,
req.get("host"),
req.get("db_type"),
current_user.id,
exc,
)
logging.debug("Full resolver exception for host %r", req.get("host"), exc_info=True)
return get_data_error_result(message=f"Could not resolve host {req.get('host')!r}.")
@@ -1198,13 +1207,7 @@ async def test_db_connection():
elif req["db_type"] == "mssql":
import pyodbc
connection_string = (
f"DRIVER={{ODBC Driver 17 for SQL Server}};"
f"SERVER={safe_host},{req['port']};"
f"DATABASE={req['database']};"
f"UID={req['username']};"
f"PWD={req['password']};"
)
connection_string = f"DRIVER={{ODBC Driver 17 for SQL Server}};SERVER={safe_host},{req['port']};DATABASE={req['database']};UID={req['username']};PWD={req['password']};"
db = pyodbc.connect(connection_string)
try:
cursor = db.cursor()
@@ -1217,14 +1220,7 @@ async def test_db_connection():
elif req["db_type"] == "IBM DB2":
import ibm_db
conn_str = (
f"DATABASE={req['database']};"
f"HOSTNAME={safe_host};"
f"PORT={req['port']};"
f"PROTOCOL=TCPIP;"
f"UID={req['username']};"
f"PWD={req['password']};"
)
conn_str = f"DATABASE={req['database']};HOSTNAME={safe_host};PORT={req['port']};PROTOCOL=TCPIP;UID={req['username']};PWD={req['password']};"
logging.info(
"DATABASE=%s;HOSTNAME=%s;PORT=%s;PROTOCOL=TCPIP;UID=%s;PWD=****;",
req["database"],
@@ -1387,9 +1383,7 @@ async def agent_chat_completion(tenant_id, agent_id=None):
if "chunks" in workflow_conv["reference"]:
workflow_conv["reference"] = [workflow_conv["reference"]]
else:
workflow_conv["reference"] = [
value for _, value in sorted(workflow_conv["reference"].items(), key=lambda item: int(item[0]))
]
workflow_conv["reference"] = [value for _, value in sorted(workflow_conv["reference"].items(), key=lambda item: int(item[0]))]
elif not isinstance(workflow_conv.get("reference"), list):
workflow_conv["reference"] = []
workflow_conv["reference"] = [_normalize_agent_reference_entry(reference) for reference in workflow_conv["reference"]]
@@ -1598,13 +1592,11 @@ async def agent_chat_completion(tenant_id, agent_id=None):
# seeing only a bare [DONE] (fixes #15169).
logging.info(
"empty agent output - returning session_id (agent_id=%s session_id=%s stream=%s)",
agent_id, session_id, True,
)
yield (
"data:"
+ json.dumps({"session_id": session_id, "data": {}}, ensure_ascii=False)
+ "\n\n"
agent_id,
session_id,
True,
)
yield ("data:" + json.dumps({"session_id": session_id, "data": {}}, ensure_ascii=False) + "\n\n")
yield "data:[DONE]\n\n"
return _build_sse_response(generate())
@@ -1614,6 +1606,7 @@ async def agent_chat_completion(tenant_id, agent_id=None):
final_ans = {}
trace_items = []
structured_output = {}
run_usage = None
async for ans in _iter_session_completion_events(tenant_id, agent_id, req, return_trace):
try:
if ans["event"] == "message":
@@ -1633,6 +1626,11 @@ async def agent_chat_completion(tenant_id, agent_id=None):
"trace": [copy.deepcopy(data)],
}
)
if ans.get("event") == "workflow_finished":
# Capture the run-level usage but keep message_end/user_inputs as
# final_ans so the non-stream response shape stays unchanged.
run_usage = ans.get("data", {}).get("usage")
continue
if ans.get("event") == "message_end":
final_ans = ans
elif ans.get("event") == "user_inputs" and not final_ans:
@@ -1647,7 +1645,9 @@ async def agent_chat_completion(tenant_id, agent_id=None):
# (fixes #15169).
logging.info(
"empty agent output - returning session_id (agent_id=%s session_id=%s stream=%s)",
agent_id, session_id, False,
agent_id,
session_id,
False,
)
return get_result(data={"session_id": session_id})
@@ -1655,6 +1655,8 @@ async def agent_chat_completion(tenant_id, agent_id=None):
final_ans["data"] = {}
final_ans["data"]["content"] = full_content
final_ans["data"]["reference"] = reference
if run_usage:
final_ans["data"]["usage"] = run_usage
if structured_output:
final_ans["data"]["structured"] = structured_output
if return_trace and final_ans:
@@ -1688,16 +1690,16 @@ async def _webhook_impl(agent_id: str, is_test: bool):
# 1. Fetch canvas by agent_id
exists, cvs = UserCanvasService.get_by_id(agent_id)
if not exists:
return get_data_error_result(code=RetCode.BAD_REQUEST,message="Canvas not found."),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message="Canvas not found."), RetCode.BAD_REQUEST
# 2. Check canvas category
if cvs.canvas_category == CanvasCategory.DataFlow:
return get_data_error_result(code=RetCode.BAD_REQUEST,message="Dataflow can not be triggered by webhook."),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message="Dataflow can not be triggered by webhook."), RetCode.BAD_REQUEST
# 3. Load DSL from canvas
dsl = getattr(cvs, "dsl", None)
if not isinstance(dsl, dict):
return get_data_error_result(code=RetCode.BAD_REQUEST,message="Invalid DSL format."),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message="Invalid DSL format."), RetCode.BAD_REQUEST
# 4. Check webhook configuration in DSL
webhook_cfg = {}
@@ -1708,15 +1710,13 @@ async def _webhook_impl(agent_id: str, is_test: bool):
webhook_cfg = cpn_obj["params"]
if not webhook_cfg:
return get_data_error_result(code=RetCode.BAD_REQUEST,message="Webhook not configured for this agent."),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message="Webhook not configured for this agent."), RetCode.BAD_REQUEST
# 5. Validate request method against webhook_cfg.methods
allowed_methods = webhook_cfg.get("methods", [])
request_method = request.method.upper()
if allowed_methods and request_method not in allowed_methods:
return get_data_error_result(
code=RetCode.BAD_REQUEST,message=f"HTTP method '{request_method}' not allowed for this webhook."
),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message=f"HTTP method '{request_method}' not allowed for this webhook."), RetCode.BAD_REQUEST
async def validate_webhook_security(security_cfg: dict):
"""Validate webhook security rules based on security configuration."""
@@ -1795,7 +1795,6 @@ async def _webhook_impl(agent_id: str, is_test: bool):
client_ip = request.remote_addr
for rule in whitelist:
if "/" in rule:
# CIDR notation
@@ -1854,7 +1853,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
def _validate_token_auth(security_cfg):
"""Validate header-based token authentication."""
token_cfg = security_cfg.get("token",{})
token_cfg = security_cfg.get("token", {})
header = token_cfg.get("token_header")
token_value = token_cfg.get("token_value")
@@ -1883,7 +1882,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
if not auth_header.startswith("Bearer "):
raise Exception("Missing Bearer token")
token = auth_header[len("Bearer "):].strip()
token = auth_header[len("Bearer ") :].strip()
if not token:
raise Exception("Empty Bearer token")
@@ -1922,10 +1921,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
else:
required_claims = []
required_claims = [
c for c in required_claims
if isinstance(c, str) and c.strip()
]
required_claims = [c for c in required_claims if isinstance(c, str) and c.strip()]
RESERVED_CLAIMS = {"exp", "sub", "aud", "iss", "nbf", "iat"}
for claim in required_claims:
@@ -1939,10 +1935,10 @@ async def _webhook_impl(agent_id: str, is_test: bool):
return decoded
try:
security_config=webhook_cfg.get("security", {})
security_config = webhook_cfg.get("security", {})
await validate_webhook_security(security_config)
except Exception as e:
return get_data_error_result(code=RetCode.BAD_REQUEST,message=str(e)),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message=str(e)), RetCode.BAD_REQUEST
if not isinstance(cvs.dsl, str):
dsl = json.dumps(cvs.dsl, ensure_ascii=False)
try:
@@ -1950,7 +1946,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
canvas = Canvas(dsl, cvs.user_id, agent_id, canvas_id=agent_id)
except Exception as e:
resp=get_data_error_result(code=RetCode.BAD_REQUEST,message=str(e))
resp = get_data_error_result(code=RetCode.BAD_REQUEST, message=str(e))
resp.status_code = RetCode.BAD_REQUEST
return resp
@@ -1967,9 +1963,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
# 3. Body
ctype = request.headers.get("Content-Type", "").split(";")[0].strip()
if ctype and ctype != content_type:
raise ValueError(
f"Invalid Content-Type: expect '{content_type}', got '{ctype}'"
)
raise ValueError(f"Invalid Content-Type: expect '{content_type}', got '{ctype}'")
body_data: dict = {}
@@ -1991,11 +1985,11 @@ async def _webhook_impl(agent_id: str, is_test: bool):
raise Exception("Too many uploaded files")
for key, file in files.items():
desc = FileService.upload_info(
cvs.user_id, # user
file, # FileStorage
None # url (None for webhook)
cvs.user_id, # user
file, # FileStorage
None, # url (None for webhook)
)
file_parsed= await canvas.get_files_async([desc])
file_parsed = await canvas.get_files_async([desc])
body_data[key] = file_parsed
elif ctype == "application/x-www-form-urlencoded":
@@ -2057,15 +2051,12 @@ async def _webhook_impl(agent_id: str, is_test: bool):
# 4. Type validation
if not validate_type(value, field_type):
raise Exception(
f"{name}.{field} type mismatch: expected {field_type}, got {type(value).__name__}"
)
raise Exception(f"{name}.{field} type mismatch: expected {field_type}, got {type(value).__name__}")
extracted[field] = value
return extracted
def default_for_type(t):
"""Return default value for the given schema type."""
if t == "file":
@@ -2145,7 +2136,6 @@ async def _webhook_impl(agent_id: str, is_test: bool):
# Default: do nothing
return value
def validate_type(value, t):
"""Validate value type against schema type t."""
if t == "file":
@@ -2179,28 +2169,24 @@ async def _webhook_impl(agent_id: str, is_test: bool):
return True
return True
parsed = await parse_webhook_request(webhook_cfg.get("content_types"))
SCHEMA = webhook_cfg.get("schema", {"query": {}, "headers": {}, "body": {}})
# Extract strictly by schema
try:
query_clean = extract_by_schema(parsed["query"], SCHEMA.get("query", {}), name="query")
query_clean = extract_by_schema(parsed["query"], SCHEMA.get("query", {}), name="query")
header_clean = extract_by_schema(parsed["headers"], SCHEMA.get("headers", {}), name="headers")
body_clean = extract_by_schema(parsed["body"], SCHEMA.get("body", {}), name="body")
body_clean = extract_by_schema(parsed["body"], SCHEMA.get("body", {}), name="body")
except Exception as e:
return get_data_error_result(code=RetCode.BAD_REQUEST,message=str(e)),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message=str(e)), RetCode.BAD_REQUEST
clean_request = {
"query": query_clean,
"headers": header_clean,
"body": body_clean,
"input": parsed
}
clean_request = {"query": query_clean, "headers": header_clean, "body": body_clean, "input": parsed}
execution_mode = webhook_cfg.get("execution_mode", "Immediately")
response_cfg = webhook_cfg.get("response", {})
def append_webhook_trace(agent_id: str, start_ts: float,event: dict, ttl=600):
def append_webhook_trace(agent_id: str, start_ts: float, event: dict, ttl=600):
from rag.utils.redis_conn import REDIS_CONN
key = f"webhook-trace-{agent_id}-logs"
@@ -2208,15 +2194,9 @@ async def _webhook_impl(agent_id: str, is_test: bool):
raw = REDIS_CONN.get(key)
obj = json.loads(raw) if raw else {"webhooks": {}}
ws = obj["webhooks"].setdefault(
str(start_ts),
{"start_ts": start_ts, "events": []}
)
ws = obj["webhooks"].setdefault(str(start_ts), {"start_ts": start_ts, "events": []})
ws["events"].append({
"ts": time.time(),
**event
})
ws["events"].append({"ts": time.time(), **event})
REDIS_CONN.set_obj(key, obj, ttl)
@@ -2225,10 +2205,10 @@ async def _webhook_impl(agent_id: str, is_test: bool):
try:
status = int(status)
except (TypeError, ValueError):
return get_data_error_result(code=RetCode.BAD_REQUEST,message=str(f"Invalid response status code: {status}")),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message=str(f"Invalid response status code: {status}")), RetCode.BAD_REQUEST
if not (200 <= status <= 399):
return get_data_error_result(code=RetCode.BAD_REQUEST,message=str(f"Invalid response status code: {status}, must be between 200 and 399")),RetCode.BAD_REQUEST
return get_data_error_result(code=RetCode.BAD_REQUEST, message=str(f"Invalid response status code: {status}, must be between 200 and 399")), RetCode.BAD_REQUEST
body_tpl = response_cfg.get("body_template", "")
@@ -2242,7 +2222,6 @@ async def _webhook_impl(agent_id: str, is_test: bool):
except (json.JSONDecodeError, TypeError):
return body, "text/plain"
body, content_type = parse_body(body_tpl)
resp = Response(
json.dumps(body, ensure_ascii=False) if content_type == "application/json" else body,
@@ -2252,11 +2231,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
async def background_run():
try:
async for ans in canvas.run(
query="",
user_id=cvs.user_id,
webhook_payload=clean_request
):
async for ans in canvas.run(query="", user_id=cvs.user_id, webhook_payload=clean_request):
if is_test:
append_webhook_trace(agent_id, start_ts, ans)
@@ -2268,7 +2243,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "finished",
"elapsed_time": time.time() - start_ts,
"success": True,
}
},
)
cvs.dsl = json.loads(str(canvas))
@@ -2285,7 +2260,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "error",
"message": str(e),
"error_type": type(e).__name__,
}
},
)
append_webhook_trace(
agent_id,
@@ -2294,7 +2269,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "finished",
"elapsed_time": time.time() - start_ts,
"success": False,
}
},
)
except Exception:
logging.exception("Failed to append webhook trace")
@@ -2305,6 +2280,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
task.add_done_callback(_background_tasks.discard)
return resp
else:
async def sse():
nonlocal canvas
contents: list[str] = []
@@ -2326,11 +2302,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
if ans["event"] == "message_end":
status = int(ans["data"].get("status", status))
if is_test:
append_webhook_trace(
agent_id,
start_ts,
ans
)
append_webhook_trace(agent_id, start_ts, ans)
if is_test:
append_webhook_trace(
agent_id,
@@ -2339,13 +2311,13 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "finished",
"elapsed_time": time.time() - start_ts,
"success": True,
}
},
)
final_content = "".join(contents)
return {
"message": final_content,
"success": True,
"code": status,
"code": status,
}
except Exception as e:
@@ -2357,7 +2329,7 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "error",
"message": str(e),
"error_type": type(e).__name__,
}
},
)
append_webhook_trace(
agent_id,
@@ -2366,9 +2338,9 @@ async def _webhook_impl(agent_id: str, is_test: bool):
"event": "finished",
"elapsed_time": time.time() - start_ts,
"success": False,
}
},
)
return {"code": 400, "message": str(e),"success":False}
return {"code": 400, "message": str(e), "success": False}
result = await sse()
return Response(
@@ -2401,6 +2373,7 @@ async def webhook_trace(agent_id: str):
if encode_webhook_id(ts) == enc_id:
return ts
return None
since_ts = request.args.get("since_ts", type=float)
webhook_id = request.args.get("webhook_id")
@@ -2434,9 +2407,7 @@ async def webhook_trace(agent_id: str):
webhooks = obj.get("webhooks", {})
if webhook_id is None:
candidates = [
float(k) for k in webhooks.keys() if float(k) > since_ts
]
candidates = [float(k) for k in webhooks.keys() if float(k) > since_ts]
if not candidates:
return get_json_result(
@@ -2492,6 +2463,7 @@ async def webhook_trace(agent_id: str):
}
)
@manager.route("/agents/attachments/<attachment_id>/download", methods=["GET"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs

137
api/db/services/canvas_service.py
View File

@@ -34,10 +34,12 @@ from peewee import fn
class CanvasTemplateService(CommonService):
model = CanvasTemplate
class DataFlowTemplateService(CommonService):
"""
Alias of CanvasTemplateService
"""
model = CanvasTemplate
@@ -46,8 +48,7 @@ class UserCanvasService(CommonService):
@classmethod
@DB.connection_context()
def get_list(cls, tenant_id,
page_number, items_per_page, orderby, desc, id, title, canvas_category=CanvasCategory.Agent):
def get_list(cls, tenant_id, page_number, items_per_page, orderby, desc, id, title, canvas_category=CanvasCategory.Agent):
agents = cls.model.select()
if id:
agents = agents.where(cls.model.id == id)
@@ -68,20 +69,9 @@ class UserCanvasService(CommonService):
@DB.connection_context()
def get_all_agents_by_tenant_ids(cls, tenant_ids, user_id):
# will get all permitted agents, be cautious
fields = [
cls.model.id,
cls.model.avatar,
cls.model.title,
cls.model.permission,
cls.model.canvas_type,
cls.model.canvas_category
]
fields = [cls.model.id, cls.model.avatar, cls.model.title, cls.model.permission, cls.model.canvas_type, cls.model.canvas_category]
# find team agents and owned agents
agents = cls.model.select(*fields).where(
(cls.model.user_id.in_(tenant_ids) & (cls.model.permission == TenantPermission.TEAM.value)) | (
cls.model.user_id == user_id
)
)
agents = cls.model.select(*fields).where((cls.model.user_id.in_(tenant_ids) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id))
# sort by create_time, asc
agents.order_by(cls.model.create_time.asc())
# maybe cause slow query by deep paginate, optimize later
@@ -100,7 +90,6 @@ class UserCanvasService(CommonService):
@DB.connection_context()
def get_by_canvas_id(cls, pid):
try:
fields = [
cls.model.id,
cls.model.avatar,
@@ -115,11 +104,9 @@ class UserCanvasService(CommonService):
cls.model.update_date,
cls.model.canvas_category,
User.nickname,
User.avatar.alias('tenant_avatar'),
User.avatar.alias("tenant_avatar"),
]
agents = cls.model.select(*fields) \
.join(User, on=(cls.model.user_id == User.id)) \
.where(cls.model.id == pid)
agents = cls.model.select(*fields).join(User, on=(cls.model.user_id == User.id)).where(cls.model.id == pid)
# obj = cls.model.query(id=pid)[0]
return True, agents.dicts()[0]
except Exception as e:
@@ -129,14 +116,7 @@ class UserCanvasService(CommonService):
@classmethod
@DB.connection_context()
def get_basic_info_by_canvas_ids(cls, canvas_id):
fields = [
cls.model.id,
cls.model.avatar,
cls.model.user_id,
cls.model.title,
cls.model.permission,
cls.model.canvas_category
]
fields = [cls.model.id, cls.model.avatar, cls.model.user_id, cls.model.title, cls.model.permission, cls.model.canvas_category]
return cls.model.select(*fields).where(cls.model.id.in_(canvas_id)).dicts()
@classmethod
@@ -162,20 +142,26 @@ class UserCanvasService(CommonService):
cls.model.permission,
cls.model.user_id.alias("tenant_id"),
User.nickname,
User.avatar.alias('tenant_avatar'),
User.avatar.alias("tenant_avatar"),
cls.model.update_time,
cls.model.canvas_type,
cls.model.canvas_category,
cls.model.tags,
]
if keywords:
agents = cls.model.select(*fields).join(User, on=(cls.model.user_id == User.id)).where(
(((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id)),
(fn.LOWER(cls.model.title).contains(keywords.lower()))
agents = (
cls.model.select(*fields)
.join(User, on=(cls.model.user_id == User.id))
.where(
(((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id)),
(fn.LOWER(cls.model.title).contains(keywords.lower())),
)
)
else:
agents = cls.model.select(*fields).join(User, on=(cls.model.user_id == User.id)).where(
(((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id))
agents = (
cls.model.select(*fields)
.join(User, on=(cls.model.user_id == User.id))
.where((((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id)))
)
if canvas_category:
agents = agents.where(cls.model.canvas_category == canvas_category)
@@ -201,7 +187,7 @@ class UserCanvasService(CommonService):
# Get latest release time for each canvas
if agents_list:
canvas_ids = [a['id'] for a in agents_list]
canvas_ids = [a["id"] for a in agents_list]
release_times = (
UserCanvasVersion.select(UserCanvasVersion.user_canvas_id, fn.MAX(UserCanvasVersion.create_time).alias("release_time"))
.where((UserCanvasVersion.user_canvas_id.in_(canvas_ids)) & (UserCanvasVersion.release))
@@ -210,7 +196,7 @@ class UserCanvasService(CommonService):
release_time_map = {r.user_canvas_id: r.release_time for r in release_times}
for agent in agents_list:
agent['release_time'] = release_time_map.get(agent['id'])
agent["release_time"] = release_time_map.get(agent["id"])
return agents_list, count
@@ -218,9 +204,7 @@ class UserCanvasService(CommonService):
@DB.connection_context()
def list_tags(cls, joined_tenant_ids, user_id, canvas_category=None):
"""Return {tag: agent_count} aggregated across agents visible to the user."""
query = cls.model.select(cls.model.tags).where(
((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id)
)
query = cls.model.select(cls.model.tags).where(((cls.model.user_id.in_(joined_tenant_ids)) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.user_id == user_id))
if canvas_category:
query = query.where(cls.model.canvas_category == canvas_category)
@@ -281,6 +265,7 @@ class UserCanvasService(CommonService):
@DB.connection_context()
def accessible(cls, canvas_id, tenant_id):
from api.db.services.user_service import UserTenantService
e, c = UserCanvasService.get_by_canvas_id(canvas_id)
if not e:
return False
@@ -345,18 +330,15 @@ async def completion(tenant_id, agent_id, session_id=None, **kwargs):
conv = API4Conversation(**conv)
message_id = str(uuid4())
conv.message.append({
"role": "user",
"content": query,
"id": message_id,
"files": files
})
conv.message.append({"role": "user", "content": query, "id": message_id, "files": files})
txt = ""
run_kwargs = {
"query": query,
"files": files,
"user_id": user_id,
"inputs": inputs,
# Used by Canvas.run to correlate RAGFlow's Langfuse generations by session.
"session_id": session_id,
}
if chat_template_kwargs is not None:
run_kwargs["chat_template_kwargs"] = chat_template_kwargs
@@ -394,14 +376,7 @@ async def completion_openai(tenant_id, agent_id, question, session_id=None, stre
if stream:
completion_tokens = 0
try:
async for ans in completion(
tenant_id=tenant_id,
agent_id=agent_id,
session_id=session_id,
query=question,
user_id=user_id,
**kwargs
):
async for ans in completion(tenant_id=tenant_id, agent_id=agent_id, session_id=session_id, query=question, user_id=user_id, **kwargs):
if isinstance(ans, str):
try:
ans = json.loads(ans[5:]) # remove "data:"
@@ -417,14 +392,7 @@ async def completion_openai(tenant_id, agent_id, question, session_id=None, stre
completion_tokens += len(tiktoken_encoder.encode(content_piece))
openai_data = get_data_openai(
id=session_id or str(uuid4()),
model=agent_id,
content=content_piece,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
stream=True
)
openai_data = get_data_openai(id=session_id or str(uuid4()), model=agent_id, content=content_piece, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, stream=True)
if ans.get("data", {}).get("reference", None):
openai_data["choices"][0]["delta"]["reference"] = ans["data"]["reference"]
@@ -435,32 +403,29 @@ async def completion_openai(tenant_id, agent_id, question, session_id=None, stre
except Exception as e:
logging.exception(e)
yield "data: " + json.dumps(
get_data_openai(
id=session_id or str(uuid4()),
model=agent_id,
content=f"**ERROR**: {str(e)}",
finish_reason="stop",
prompt_tokens=prompt_tokens,
completion_tokens=len(tiktoken_encoder.encode(f"**ERROR**: {str(e)}")),
stream=True
),
ensure_ascii=False
) + "\n\n"
yield (
"data: "
+ json.dumps(
get_data_openai(
id=session_id or str(uuid4()),
model=agent_id,
content=f"**ERROR**: {str(e)}",
finish_reason="stop",
prompt_tokens=prompt_tokens,
completion_tokens=len(tiktoken_encoder.encode(f"**ERROR**: {str(e)}")),
stream=True,
),
ensure_ascii=False,
)
+ "\n\n"
)
yield "data: [DONE]\n\n"
else:
try:
all_content = ""
reference = {}
async for ans in completion(
tenant_id=tenant_id,
agent_id=agent_id,
session_id=session_id,
query=question,
user_id=user_id,
**kwargs
):
async for ans in completion(tenant_id=tenant_id, agent_id=agent_id, session_id=session_id, query=question, user_id=user_id, **kwargs):
if isinstance(ans, str):
ans = json.loads(ans[5:])
if ans.get("event") not in ["message", "message_end"]:
@@ -475,13 +440,7 @@ async def completion_openai(tenant_id, agent_id, question, session_id=None, stre
completion_tokens = len(tiktoken_encoder.encode(all_content))
openai_data = get_data_openai(
id=session_id or str(uuid4()),
model=agent_id,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
content=all_content,
finish_reason="stop",
param=None
id=session_id or str(uuid4()), model=agent_id, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, content=all_content, finish_reason="stop", param=None
)
if reference:
@@ -497,5 +456,5 @@ async def completion_openai(tenant_id, agent_id, question, session_id=None, stre
completion_tokens=len(tiktoken_encoder.encode(f"**ERROR**: {str(e)}")),
content=f"**ERROR**: {str(e)}",
finish_reason="stop",
param=None
param=None,
)

81
api/db/services/llm_service.py
View File

@@ -27,7 +27,7 @@ from langfuse import propagate_attributes
from api.db.db_models import LLM
from api.db.services.common_service import CommonService
from api.db.services.tenant_llm_service import LLM4Tenant
from common.token_utils import num_tokens_from_string
from common.token_utils import num_tokens_from_string, record_run_token_usage, langfuse_run_attrs
class LLMService(CommonService):
@@ -39,11 +39,49 @@ class LLMBundle(LLM4Tenant):
super().__init__(tenant_id, model_config, lang, **kwargs)
def _start_langfuse_observation(self, **kwargs):
# Correlating attributes (session_id/user_id) let Langfuse group all of a
# turn's generations. They may come from this bundle (chat/dialog path) or,
# for agent runs whose bundles are created without them, from the per-run
# context installed by Canvas.run.
attrs = {}
if self.langfuse_session_id:
with propagate_attributes(session_id=self.langfuse_session_id):
attrs["session_id"] = self.langfuse_session_id
run_attrs = langfuse_run_attrs.get()
if run_attrs:
for k in ("session_id", "user_id"):
if run_attrs.get(k) and k not in attrs:
attrs[k] = run_attrs[k]
if attrs:
with propagate_attributes(**attrs):
return self.langfuse.start_observation(**kwargs)
return self.langfuse.start_observation(**kwargs)
def _reset_last_usage(self) -> None:
"""Clear the model's per-call usage so a failed call that returns before
updating it cannot leak the previous call's usage into this run."""
if hasattr(self.mdl, "last_usage"):
self.mdl.last_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
def _report_usage(self, total_tokens: int) -> dict:
"""Record a chat call's usage to the active agent run and return the
prompt/completion/total split for Langfuse.
``total_tokens`` is the authoritative total from the call. The prompt/completion
split is taken from the provider response (``mdl.last_usage``) only when it is
consistent with ``total_tokens`` (i.e. produced by this same call); otherwise the
split is reported as 0 while the total still aggregates correctly.
"""
split = getattr(self.mdl, "last_usage", None) or {}
prompt = int(split.get("prompt_tokens", 0) or 0)
completion = int(split.get("completion_tokens", 0) or 0)
if not total_tokens:
total_tokens = int(split.get("total_tokens", 0) or 0)
if (prompt + completion) != total_tokens:
# Stale or inconsistent split — keep the total, drop the unreliable split.
prompt, completion = 0, 0
record_run_token_usage(prompt, completion, total_tokens)
return {"input": prompt, "output": completion, "total": total_tokens}
def close(self):
"""Release resources held by this LLMBundle instance."""
super().close()
@@ -139,7 +177,9 @@ class LLMBundle(LLM4Tenant):
def similarity(self, query: str, texts: list):
if self.langfuse:
generation = self._start_langfuse_observation(trace_context=self.trace_context, as_type="generation", name="similarity", model=self.model_config["llm_name"], input={"query": query, "texts": texts})
generation = self._start_langfuse_observation(
trace_context=self.trace_context, as_type="generation", name="similarity", model=self.model_config["llm_name"], input={"query": query, "texts": texts}
)
sim, used_tokens = self.mdl.similarity(query, texts)
logging.info("LLMBundle.similarity used_tokens: %d", used_tokens)
@@ -165,7 +205,9 @@ class LLMBundle(LLM4Tenant):
def describe_with_prompt(self, image, prompt):
if self.langfuse:
generation = self._start_langfuse_observation(trace_context=self.trace_context, as_type="generation", name="describe_with_prompt", metadata={"model": self.model_config["llm_name"], "prompt": prompt})
generation = self._start_langfuse_observation(
trace_context=self.trace_context, as_type="generation", name="describe_with_prompt", metadata={"model": self.model_config["llm_name"], "prompt": prompt}
)
txt, used_tokens = self.mdl.describe_with_prompt(image, prompt)
logging.info("LLMBundle.describe_with_prompt used_tokens: %d", used_tokens)
@@ -194,7 +236,8 @@ class LLMBundle(LLM4Tenant):
supports_stream = hasattr(mdl, "stream_transcription") and callable(getattr(mdl, "stream_transcription"))
if supports_stream:
if self.langfuse:
generation = self._start_langfuse_observation(as_type="generation",
generation = self._start_langfuse_observation(
as_type="generation",
trace_context=self.trace_context,
name="stream_transcription",
metadata={"model": self.model_config["llm_name"]},
@@ -228,7 +271,8 @@ class LLMBundle(LLM4Tenant):
return
if self.langfuse:
generation = self._start_langfuse_observation(as_type="generation",
generation = self._start_langfuse_observation(
as_type="generation",
trace_context=self.trace_context,
name="stream_transcription",
metadata={"model": self.model_config["llm_name"]},
@@ -377,11 +421,14 @@ class LLMBundle(LLM4Tenant):
generation = None
if self.langfuse:
generation = self._start_langfuse_observation(trace_context=self.trace_context, as_type="generation", name="chat", model=self.model_config["llm_name"], input={"system": system, "history": history})
generation = self._start_langfuse_observation(
trace_context=self.trace_context, as_type="generation", name="chat", model=self.model_config["llm_name"], input={"system": system, "history": history}
)
chat_partial = partial(base_fn, system, history, gen_conf)
use_kwargs = self._clean_param(chat_partial, **kwargs)
self._reset_last_usage()
try:
txt, used_tokens = await chat_partial(**use_kwargs)
except Exception as e:
@@ -397,8 +444,10 @@ class LLMBundle(LLM4Tenant):
if used_tokens:
logging.info("LLMBundle.async_chat used_tokens: %d", used_tokens)
usage_details = self._report_usage(used_tokens)
if generation:
generation.update(output={"output": txt}, usage_details={"total_tokens": used_tokens})
generation.update(output={"output": txt}, usage_details=usage_details)
generation.end()
return txt
@@ -418,11 +467,14 @@ class LLMBundle(LLM4Tenant):
generation = None
if self.langfuse:
generation = self._start_langfuse_observation(trace_context=self.trace_context, as_type="generation", name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history})
generation = self._start_langfuse_observation(
trace_context=self.trace_context, as_type="generation", name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history}
)
if stream_fn:
chat_partial = partial(stream_fn, system, history, gen_conf)
use_kwargs = self._clean_param(chat_partial, **kwargs)
self._reset_last_usage()
try:
async for txt in chat_partial(**use_kwargs):
if isinstance(txt, int):
@@ -444,8 +496,9 @@ class LLMBundle(LLM4Tenant):
raise
if total_tokens:
logging.info("LLMBundle.async_chat_streamly used_tokens: %d", total_tokens)
usage_details = self._report_usage(total_tokens)
if generation:
generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens})
generation.update(output={"output": ans}, usage_details=usage_details)
generation.end()
return
@@ -461,11 +514,14 @@ class LLMBundle(LLM4Tenant):
generation = None
if self.langfuse:
generation = self._start_langfuse_observation(trace_context=self.trace_context, as_type="generation", name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history})
generation = self._start_langfuse_observation(
trace_context=self.trace_context, as_type="generation", name="chat_streamly", model=self.model_config["llm_name"], input={"system": system, "history": history}
)
if stream_fn:
chat_partial = partial(stream_fn, system, history, gen_conf)
use_kwargs = self._clean_param(chat_partial, **kwargs)
self._reset_last_usage()
try:
async for txt in chat_partial(**use_kwargs):
if isinstance(txt, int):
@@ -487,7 +543,8 @@ class LLMBundle(LLM4Tenant):
raise
if total_tokens:
logging.info("LLMBundle.async_chat_streamly_delta used_tokens: %d", total_tokens)
usage_details = self._report_usage(total_tokens)
if generation:
generation.update(output={"output": ans}, usage_details={"total_tokens": total_tokens})
generation.update(output={"output": ans}, usage_details=usage_details)
generation.end()
return