pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-21 05:16:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: trace information can be returned by the agent completion API (#12019)

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

Trace information can be returned by the agent completion API.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Yongteng Lei
2025-12-18 15:52:11 +08:00
committed by GitHub
parent 2331b3a270
commit 151480dc85
2 changed files with 226 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
api/apps/sdk/session.py
View File

@ -14,6 +14,7 @@
# limitations under the License.
#
import json
import copy
import re
import time
@ -446,10 +447,12 @@ async def agents_completion_openai_compatibility(tenant_id, agent_id):
@token_required
async def agent_completions(tenant_id, agent_id):
req = await get_request_json()
return_trace = bool(req.get("return_trace", False))
if req.get("stream", True):
async def generate():
trace_items = []
async for answer in agent_completion(tenant_id=tenant_id, agent_id=agent_id, **req):
if isinstance(answer, str):
try:
@ -457,7 +460,21 @@ async def agent_completions(tenant_id, agent_id):
except Exception:
continue
if ans.get("event") not in ["message", "message_end"]:
event = ans.get("event")
if event == "node_finished":
if return_trace:
data = ans.get("data", {})
trace_items.append(
{
"component_id": data.get("component_id"),
"trace": [copy.deepcopy(data)],
}
)
ans.setdefault("data", {})["trace"] = trace_items
answer = "data:" + json.dumps(ans, ensure_ascii=False) + "\n\n"
yield answer
if event not in ["message", "message_end"]:
continue
yield answer
@ -474,6 +491,7 @@ async def agent_completions(tenant_id, agent_id):
full_content = ""
reference = {}
final_ans = ""
trace_items = []
async for answer in agent_completion(tenant_id=tenant_id, agent_id=agent_id, **req):
try:
ans = json.loads(answer[5:])
@ -484,11 +502,22 @@ async def agent_completions(tenant_id, agent_id):
if ans.get("data", {}).get("reference", None):
reference.update(ans["data"]["reference"])
if return_trace and ans.get("event") == "node_finished":
data = ans.get("data", {})
trace_items.append(
{
"component_id": data.get("component_id"),
"trace": [copy.deepcopy(data)],
}
)
final_ans = ans
except Exception as e:
return get_result(data=f"**ERROR**: {str(e)}")
final_ans["data"]["content"] = full_content
final_ans["data"]["reference"] = reference
if return_trace and final_ans:
final_ans["data"]["trace"] = trace_items
return get_result(data=final_ans)