mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 02:15:06 +00:00
fix: minor parser and id hardening
This commit is contained in:
Alex committed
1 parent
223a7588aa
commit
9f19bc64b5
14 files changed
+308
-20
No files matched your search
@@ -81,6 +81,22 @@ def _redact_args_for_log(args: Any) -> Any:
|
||||
return redacted
|
||||
|
||||
|
||||
def _journal_key(call_id: str, message_id: Optional[str]) -> str:
|
||||
"""Namespace the durability-journal key by the per-turn ``message_id``.
|
||||
|
||||
``tool_call_attempts.call_id`` is a table-wide primary key, but providers
|
||||
reuse deterministic ids (e.g. ``functions.create_artifact:0``) across turns
|
||||
and users, so distinct calls collide on that PK and the later journal rows
|
||||
are silently dropped (``ON CONFLICT DO NOTHING``). Scoping the key by
|
||||
``message_id`` (unique per turn) gives each logical call its own row while a
|
||||
genuine retry of the same call within the same turn still dedupes. The raw
|
||||
``call_id`` is left untouched for LLM tool-call/tool-result pairing and the
|
||||
UI. Headless attempts with no ``message_id`` keep the raw key (unchanged
|
||||
pre-existing behaviour).
|
||||
"""
|
||||
return f"{message_id}:{call_id}" if message_id else call_id
|
||||
|
||||
|
||||
def _record_proposed(
|
||||
call_id: str,
|
||||
tool_name: str,
|
||||
@@ -104,7 +120,7 @@ def _record_proposed(
|
||||
try:
|
||||
with db_session() as conn:
|
||||
inserted = ToolCallAttemptsRepository(conn).record_proposed(
|
||||
call_id,
|
||||
_journal_key(call_id, message_id),
|
||||
tool_name,
|
||||
action_name,
|
||||
arguments,
|
||||
@@ -145,12 +161,13 @@ def _mark_executed(
|
||||
effect would be invisible to the journal. Both paths are scoped to
|
||||
the owning ``user_id`` so a reused ``call_id`` can't cross tenants.
|
||||
"""
|
||||
key = _journal_key(call_id, message_id)
|
||||
try:
|
||||
with db_session() as conn:
|
||||
repo = ToolCallAttemptsRepository(conn)
|
||||
if proposed_ok:
|
||||
updated = repo.mark_executed(
|
||||
call_id,
|
||||
key,
|
||||
result,
|
||||
message_id=message_id,
|
||||
artifact_id=artifact_id,
|
||||
@@ -160,7 +177,7 @@ def _mark_executed(
|
||||
return
|
||||
# Fallback synthesizes the row so the journal isn't lost.
|
||||
repo.upsert_executed(
|
||||
call_id,
|
||||
key,
|
||||
tool_name=tool_name or "unknown",
|
||||
action_name=action_name or "",
|
||||
arguments=arguments if arguments is not None else {},
|
||||
@@ -175,10 +192,18 @@ def _mark_executed(
|
||||
logger.exception("tool_call_attempts executed write failed for %s", call_id)
|
||||
|
||||
|
||||
def _mark_failed(call_id: str, error: str, *, user_id: Optional[str] = None) -> None:
|
||||
def _mark_failed(
|
||||
call_id: str,
|
||||
error: str,
|
||||
*,
|
||||
message_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
) -> None:
|
||||
try:
|
||||
with db_session() as conn:
|
||||
ToolCallAttemptsRepository(conn).mark_failed(call_id, error, user_id=user_id)
|
||||
ToolCallAttemptsRepository(conn).mark_failed(
|
||||
_journal_key(call_id, message_id), error, user_id=user_id
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("tool_call_attempts failed-write failed for %s", call_id)
|
||||
|
||||
@@ -676,7 +701,12 @@ class ToolExecutor:
|
||||
user_id=self.user,
|
||||
agent_id=self.agent_id,
|
||||
):
|
||||
_mark_failed(call_id, tool_call_data["result"], user_id=self.user)
|
||||
_mark_failed(
|
||||
call_id,
|
||||
tool_call_data["result"],
|
||||
message_id=self.message_id,
|
||||
user_id=self.user,
|
||||
)
|
||||
yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}}
|
||||
self.tool_calls.append(tool_call_data)
|
||||
return "Failed to parse tool call.", call_id
|
||||
@@ -714,6 +744,7 @@ class ToolExecutor:
|
||||
_mark_failed(
|
||||
call_id,
|
||||
f"Tool with ID {tool_id} not found.",
|
||||
message_id=self.message_id,
|
||||
user_id=self.user,
|
||||
)
|
||||
yield {"type": "tool_call", "data": {**tool_call_data, "status": "error"}}
|
||||
@@ -756,7 +787,9 @@ class ToolExecutor:
|
||||
tool_call_data["arguments"] = {}
|
||||
tool_call_data["status"] = "error"
|
||||
if proposed_ok:
|
||||
_mark_failed(call_id, error_message, user_id=self.user)
|
||||
_mark_failed(
|
||||
call_id, error_message, message_id=self.message_id, user_id=self.user
|
||||
)
|
||||
yield {
|
||||
"type": "tool_call",
|
||||
"data": {**tool_call_data, "status": "error"},
|
||||
@@ -813,7 +846,9 @@ class ToolExecutor:
|
||||
tool_call_data["result"] = error_message
|
||||
tool_call_data["status"] = "error"
|
||||
if proposed_ok:
|
||||
_mark_failed(call_id, error_message, user_id=self.user)
|
||||
_mark_failed(
|
||||
call_id, error_message, message_id=self.message_id, user_id=self.user
|
||||
)
|
||||
yield {"type": "tool_call", "data": {**tool_call_data}}
|
||||
self.tool_calls.append(tool_call_data)
|
||||
return error_message, call_id
|
||||
@@ -842,7 +877,9 @@ class ToolExecutor:
|
||||
result = tool.execute_action(action_name, **parameters)
|
||||
except Exception as exc:
|
||||
if proposed_ok:
|
||||
_mark_failed(call_id, str(exc), user_id=self.user)
|
||||
_mark_failed(
|
||||
call_id, str(exc), message_id=self.message_id, user_id=self.user
|
||||
)
|
||||
raise
|
||||
|
||||
get_artifact_id = getattr(tool, "get_artifact_id", None) if tool_data["name"] != "api_tool" else None
|
||||
|
||||
@@ -366,6 +366,14 @@ class WorkflowEngine:
|
||||
}
|
||||
|
||||
node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs)
|
||||
# Attribute this node's tool calls to the run's message and, crucially,
|
||||
# namespace their durability journal keys by it: node executors are built
|
||||
# without a message_id, so providers that reuse deterministic call ids
|
||||
# ("functions.create_artifact:0") would otherwise collide across runs on
|
||||
# the tool_call_attempts primary key and drop the later journal rows.
|
||||
node_agent.tool_executor.message_id = getattr(
|
||||
self.agent.tool_executor, "message_id", None
|
||||
)
|
||||
# Run-scope the node agent's tools so artifact_generator / code_executor
|
||||
# address artifacts by this workflow run: a short ref (A1) created by one
|
||||
# node resolves for edit_artifact in a later node within the same run. Only
|
||||
@@ -998,7 +1006,14 @@ class WorkflowEngine:
|
||||
output_template = str(config.get("output_template", ""))
|
||||
if output_template:
|
||||
formatted_output = self._format_template(output_template)
|
||||
# A prior streaming node's text otherwise runs straight into the
|
||||
# end-node output ("...enterprise growth.Sales analysis complete");
|
||||
# insert the same paragraph break streamed nodes use between each
|
||||
# other so the segments stay readable.
|
||||
if getattr(self, "_has_streamed", False):
|
||||
yield {"answer": "\n\n"}
|
||||
yield {"answer": formatted_output}
|
||||
self._has_streamed = True
|
||||
|
||||
def _parse_structured_output(self, raw_response: str) -> tuple[bool, Optional[Any]]:
|
||||
normalized_response = raw_response.strip()
|
||||
@@ -1031,13 +1046,20 @@ class WorkflowEngine:
|
||||
fence = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
|
||||
if fence:
|
||||
return fence.group(1).strip()
|
||||
# Fall back to the outermost {...} or [...] span.
|
||||
# Fall back to the outermost {...} or [...] span, choosing whichever
|
||||
# bracket opens first. Fixing the order to "{" before "[" would slice a
|
||||
# top-level array (``[{...},{...}]``) from its first "{" to its last "}",
|
||||
# dropping the array framing (invalid JSON) or extracting an inner object
|
||||
# that parses as silently-wrong structured data.
|
||||
best: Optional[str] = None
|
||||
best_start = -1
|
||||
for open_ch, close_ch in (("{", "}"), ("[", "]")):
|
||||
start = text.find(open_ch)
|
||||
end = text.rfind(close_ch)
|
||||
if start != -1 and end > start:
|
||||
return text[start : end + 1]
|
||||
return None
|
||||
if start != -1 and end > start and (best_start == -1 or start < best_start):
|
||||
best_start = start
|
||||
best = text[start : end + 1]
|
||||
return best
|
||||
|
||||
def _normalize_node_json_schema(
|
||||
self, schema: Optional[Dict[str, Any]], node_title: str
|
||||
|
||||
@@ -237,6 +237,10 @@ class BaseAnswerResource:
|
||||
Server-sent event strings
|
||||
"""
|
||||
response_full, thought, source_log_docs, tool_calls = "", "", [], []
|
||||
# Set when a workflow agent run emits its ``workflow_run`` event; persisted
|
||||
# onto the message metadata so the chat can render the run's produced
|
||||
# artifacts on reload.
|
||||
workflow_run_id: Optional[str] = None
|
||||
is_structured = False
|
||||
schema_info = None
|
||||
structured_chunks = []
|
||||
@@ -557,6 +561,15 @@ class BaseAnswerResource:
|
||||
# documents were dropped). Forwarded verbatim so the client can
|
||||
# surface it without failing the turn; never sanitized as an error.
|
||||
yield _emit({"type": "notice", "notice": line.get("notice", "")})
|
||||
elif line.get("type") == "workflow_run":
|
||||
# Stash the run id in the message metadata so every
|
||||
# persistence path (finalize / save / abort / error) records
|
||||
# it — the chat renders the run's produced artifacts from it
|
||||
# on reload. Still forwarded so the live client captures it.
|
||||
workflow_run_id = line.get("workflow_run_id")
|
||||
if workflow_run_id:
|
||||
query_metadata["workflow_run_id"] = workflow_run_id
|
||||
yield _emit(line)
|
||||
else:
|
||||
yield _emit(line)
|
||||
if is_structured and structured_chunks:
|
||||
|
||||
@@ -154,7 +154,12 @@ class ConversationService:
|
||||
completion = llm.gen(
|
||||
model=getattr(llm, "model_id", None) or model_id,
|
||||
messages=messages_summary,
|
||||
max_tokens=500,
|
||||
# Reasoning-capable default models spend the whole budget inside
|
||||
# reasoning_content before emitting any title, so 500 came back
|
||||
# empty (finish_reason=length). Give enough room to finish
|
||||
# thinking and still produce the 3-word title; non-reasoning
|
||||
# models stop far short of this cap.
|
||||
max_tokens=2000,
|
||||
)
|
||||
|
||||
if not completion or not completion.strip():
|
||||
@@ -418,7 +423,11 @@ class ConversationService:
|
||||
completion = llm.gen(
|
||||
model=getattr(llm, "model_id", None) or title_inputs.get("model_id"),
|
||||
messages=messages_summary,
|
||||
max_tokens=500,
|
||||
# Reasoning-capable default models spend the whole budget inside
|
||||
# reasoning_content before emitting any title, so 500 came back empty
|
||||
# (finish_reason=length). Give room to finish and still produce the
|
||||
# 3-word title; non-reasoning models stop far short of this cap.
|
||||
max_tokens=2000,
|
||||
)
|
||||
if not completion or not completion.strip():
|
||||
completion = fallback_name or "New Conversation"
|
||||
|
||||
@@ -233,6 +233,9 @@ class GetSingleConversation(Resource):
|
||||
"status": msg.get("status"),
|
||||
"request_id": msg.get("request_id"),
|
||||
"last_heartbeat_at": metadata.get("last_heartbeat_at"),
|
||||
# Surfaced from metadata so the chat can render a
|
||||
# workflow run's produced artifacts on reload.
|
||||
"workflow_run_id": metadata.get("workflow_run_id"),
|
||||
}
|
||||
if metadata:
|
||||
query["metadata"] = metadata
|
||||
|
||||
@@ -963,6 +963,7 @@ class LLMHandler(ABC):
|
||||
_mark_failed(
|
||||
pause_info["call_id"],
|
||||
f"headless: {deny_reason}",
|
||||
message_id=agent.tool_executor.message_id,
|
||||
user_id=agent.tool_executor.user,
|
||||
)
|
||||
yield {
|
||||
@@ -1047,7 +1048,18 @@ class LLMHandler(ABC):
|
||||
updated_messages.append(assistant_msg)
|
||||
assistant_appended = True
|
||||
|
||||
updated_messages.append(self.create_tool_message(call, tool_response))
|
||||
# The tool result's tool_call_id must match the id put on the
|
||||
# assistant tool_call above (``call_id`` — a synthesized UUID
|
||||
# when the provider omitted an id), not the raw ``call.id`` which
|
||||
# may be empty. A mismatch orphans the tool message and 400s the
|
||||
# next completion ("'tool' must be a response to a preceding
|
||||
# message with 'tool_calls'").
|
||||
resolved_call = ToolCall(
|
||||
id=call_id, name=call.name, arguments=call.arguments
|
||||
)
|
||||
updated_messages.append(
|
||||
self.create_tool_message(resolved_call, tool_response)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing tool: {str(e)}", exc_info=True)
|
||||
# The error tool message's tool_call_id must match the assistant
|
||||
|
||||
@@ -14,6 +14,7 @@ import remarkGfm from 'remark-gfm';
|
||||
import remarkMath from 'remark-math';
|
||||
|
||||
import SchedulerToolCallCard from '../agents/schedules/SchedulerToolCallCard';
|
||||
import WorkflowRunArtifacts from '../agents/workflow/WorkflowRunArtifacts';
|
||||
import ChevronDown from '../assets/chevron-down.svg';
|
||||
import Cloud from '../assets/cloud.svg';
|
||||
import DocsGPT3 from '../assets/cute_docsgpt3.svg';
|
||||
@@ -68,6 +69,9 @@ const ConversationBubble = forwardRef<
|
||||
thought?: string;
|
||||
sources?: { title: string; text: string; link: string }[];
|
||||
toolCalls?: ToolCallsType[];
|
||||
/** Set when this answer came from a workflow agent run; renders the
|
||||
* run's produced artifacts as click-through chips/previews. */
|
||||
workflowRunId?: string;
|
||||
research?: ResearchState;
|
||||
retryBtn?: React.ReactElement;
|
||||
questionNumber?: number;
|
||||
@@ -97,6 +101,7 @@ const ConversationBubble = forwardRef<
|
||||
thought,
|
||||
sources,
|
||||
toolCalls,
|
||||
workflowRunId,
|
||||
research,
|
||||
retryBtn,
|
||||
questionNumber,
|
||||
@@ -491,6 +496,14 @@ const ConversationBubble = forwardRef<
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
{workflowRunId && (
|
||||
<div className="my-2 mr-5 ml-2">
|
||||
<WorkflowRunArtifacts
|
||||
workflowRunId={workflowRunId}
|
||||
inProgress={isStreaming}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
{thought && (
|
||||
<Thought thought={thought} preprocessLaTeX={preprocessLaTeX} />
|
||||
)}
|
||||
@@ -523,7 +536,10 @@ const ConversationBubble = forwardRef<
|
||||
{segment.type === 'text' ? (
|
||||
<ReactMarkdown
|
||||
className="fade-in flex flex-col gap-3 leading-normal wrap-break-word whitespace-pre-wrap"
|
||||
remarkPlugins={[remarkGfm, remarkMath]}
|
||||
remarkPlugins={[
|
||||
remarkGfm,
|
||||
[remarkMath, { singleDollarTextMath: false }],
|
||||
]}
|
||||
rehypePlugins={[rehypeKatex]}
|
||||
components={{
|
||||
a({ href, children }) {
|
||||
@@ -1335,7 +1351,10 @@ function Thought({
|
||||
<div className="bg-muted dark:bg-answer-bubble rounded-3xl px-7 py-4.5">
|
||||
<ReactMarkdown
|
||||
className="fade-in leading-normal wrap-break-word whitespace-pre-wrap"
|
||||
remarkPlugins={[remarkGfm, remarkMath]}
|
||||
remarkPlugins={[
|
||||
remarkGfm,
|
||||
[remarkMath, { singleDollarTextMath: false }],
|
||||
]}
|
||||
rehypePlugins={[rehypeKatex]}
|
||||
components={{
|
||||
code(props) {
|
||||
|
||||
@@ -287,7 +287,10 @@ export default function ConversationMessages({
|
||||
query.response ||
|
||||
(query.tool_calls && query.tool_calls.length > 0) ||
|
||||
query.research ||
|
||||
query.notice;
|
||||
query.notice ||
|
||||
// An artifact-only workflow run may produce no textual answer / tool_calls;
|
||||
// the run id alone is enough to render its produced artifacts.
|
||||
query.workflow_run_id;
|
||||
if (hasContent) {
|
||||
const isCurrentlyStreaming =
|
||||
status === 'loading' && index === queries.length - 1;
|
||||
@@ -308,6 +311,7 @@ export default function ConversationMessages({
|
||||
thought={query.thought}
|
||||
sources={query.sources}
|
||||
toolCalls={query.tool_calls}
|
||||
workflowRunId={query.workflow_run_id}
|
||||
research={query.research}
|
||||
onOpenArtifact={onOpenArtifact}
|
||||
onToolAction={onToolAction}
|
||||
|
||||
@@ -62,6 +62,9 @@ export interface Query {
|
||||
thought?: string;
|
||||
sources?: { title: string; text: string; link: string }[];
|
||||
tool_calls?: ToolCallsType[];
|
||||
// Set when this answer came from a workflow agent run; lets the chat render
|
||||
// the run's produced artifacts via WorkflowRunArtifacts.
|
||||
workflow_run_id?: string;
|
||||
error?: string;
|
||||
// Non-fatal notice (e.g. some workflow input documents were dropped). Shown
|
||||
// alongside the answer; unlike ``error`` it does not fail the turn or end the stream.
|
||||
|
||||
@@ -55,6 +55,7 @@ export function mapServerQueryToClient(raw: any): Query {
|
||||
thought: raw?.thought ?? undefined,
|
||||
sources: sources && sources.length > 0 ? sources : undefined,
|
||||
tool_calls: toolCalls && toolCalls.length > 0 ? toolCalls : undefined,
|
||||
workflow_run_id: raw?.workflow_run_id ?? undefined,
|
||||
attachments: raw?.attachments ?? undefined,
|
||||
messageId: raw?.message_id ?? undefined,
|
||||
messageStatus: status,
|
||||
@@ -435,6 +436,13 @@ export const fetchAnswer = createAsyncThunk<
|
||||
dispatch(
|
||||
conversationSlice.actions.setStatus('awaiting_tool_actions'),
|
||||
);
|
||||
} else if (data.type === 'workflow_run') {
|
||||
dispatch(
|
||||
conversationSlice.actions.setWorkflowRunId({
|
||||
index: targetIndex,
|
||||
workflowRunId: data.workflow_run_id,
|
||||
}),
|
||||
);
|
||||
} else if (data.type === 'research_plan') {
|
||||
dispatch(
|
||||
updateResearchPlan({
|
||||
@@ -877,6 +885,17 @@ export const conversationSlice = createSlice({
|
||||
};
|
||||
} else state.queries[index].tool_calls.push(tool_call);
|
||||
},
|
||||
// Records the workflow run id so the answer bubble can render the run's
|
||||
// produced artifacts (WorkflowRunArtifacts fetches by this id).
|
||||
setWorkflowRunId(
|
||||
state,
|
||||
action: PayloadAction<{ index: number; workflowRunId: string }>,
|
||||
) {
|
||||
const { index, workflowRunId } = action.payload;
|
||||
const query = state.queries[index];
|
||||
if (!query || !workflowRunId) return;
|
||||
query.workflow_run_id = workflowRunId;
|
||||
},
|
||||
updateResearchPlan(
|
||||
state,
|
||||
action: PayloadAction<{
|
||||
@@ -1090,6 +1109,7 @@ export const {
|
||||
updateThought,
|
||||
updateStreamingSource,
|
||||
updateToolCall,
|
||||
setWorkflowRunId,
|
||||
updateResearchPlan,
|
||||
updateResearchProgress,
|
||||
setConversation,
|
||||
|
||||
@@ -147,11 +147,60 @@ class TestExecuteJournaling:
|
||||
|
||||
_drain(executor.execute(_TOOLS_DICT, _make_call(call_id="cm1"), "MockLLM"))
|
||||
|
||||
row = _select_attempt(pg_conn, "cm1")
|
||||
# The durability key is namespaced by message_id so providers that reuse
|
||||
# deterministic call ids across turns don't collide on the table-wide PK;
|
||||
# the raw id alone therefore matches nothing.
|
||||
assert _select_attempt(pg_conn, "cm1") is None
|
||||
row = _select_attempt(pg_conn, f"{message_uuid}:cm1")
|
||||
assert row is not None
|
||||
assert row["status"] == "executed"
|
||||
assert str(row["message_id"]) == message_uuid
|
||||
|
||||
def test_reused_call_id_across_turns_does_not_collide(
|
||||
self, pg_conn, mock_tool_manager, monkeypatch
|
||||
):
|
||||
"""Regression for the tool_call_attempts PK collision (#9): a provider
|
||||
that reuses a deterministic call id (``functions.create_artifact:0``)
|
||||
across turns must journal a distinct row per turn, not silently drop the
|
||||
later one on the table-wide ``call_id`` primary key."""
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
conv = repo.create("u-collide", "collision-test")
|
||||
msg1 = repo.reserve_message(
|
||||
str(conv["id"]), prompt="q1", placeholder_response="...",
|
||||
request_id="req-c1", status="pending",
|
||||
)
|
||||
msg2 = repo.reserve_message(
|
||||
str(conv["id"]), prompt="q2", placeholder_response="...",
|
||||
request_id="req-c2", status="pending",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.agents.tool_executor.ToolActionParser",
|
||||
lambda _cls, **kw: Mock(
|
||||
parse_args=Mock(return_value=("t1", "test_action", {}))
|
||||
),
|
||||
)
|
||||
_patch_db(monkeypatch, pg_conn)
|
||||
|
||||
reused = "functions.create_artifact:0"
|
||||
for msg in (msg1, msg2):
|
||||
executor = ToolExecutor(user="u-collide")
|
||||
executor.message_id = str(msg["id"])
|
||||
_drain(
|
||||
executor.execute(_TOOLS_DICT, _make_call(call_id=reused), "MockLLM")
|
||||
)
|
||||
|
||||
# Both turns journaled their own row; before the fix the second call hit
|
||||
# ON CONFLICT DO NOTHING on the shared PK and was dropped.
|
||||
assert _select_attempt(pg_conn, reused) is None
|
||||
r1 = _select_attempt(pg_conn, f"{msg1['id']}:{reused}")
|
||||
r2 = _select_attempt(pg_conn, f"{msg2['id']}:{reused}")
|
||||
assert r1 is not None and str(r1["message_id"]) == str(msg1["id"])
|
||||
assert r2 is not None and str(r2["message_id"]) == str(msg2["id"])
|
||||
|
||||
def test_tool_raises_marks_failed_and_reraises(
|
||||
self, pg_conn, mock_tool_manager, monkeypatch
|
||||
):
|
||||
|
||||
@@ -234,6 +234,29 @@ class TestExecuteStateNode:
|
||||
assert "result" not in engine.state
|
||||
|
||||
|
||||
class TestEndNodeSeparator:
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_separates_from_prior_streamed_output(self):
|
||||
# A prior streaming node ran, so the end template must be preceded by a
|
||||
# paragraph break instead of mashing ("...growth.Sales analysis...").
|
||||
node = _make_node("n1", NodeType.END, config={"config": {"output_template": "Final line."}})
|
||||
engine = WorkflowEngine(_make_graph([node], []), _make_agent())
|
||||
engine.state = {}
|
||||
engine._has_streamed = True
|
||||
answers = [e["answer"] for e in engine._execute_end_node(node)]
|
||||
assert answers[0] == "\n\n"
|
||||
assert answers[-1].strip() == "Final line."
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_no_separator_when_nothing_streamed(self):
|
||||
node = _make_node("n1", NodeType.END, config={"config": {"output_template": "Final line."}})
|
||||
engine = WorkflowEngine(_make_graph([node], []), _make_agent())
|
||||
engine.state = {}
|
||||
answers = [e["answer"] for e in engine._execute_end_node(node)]
|
||||
assert answers == ["Final line."]
|
||||
|
||||
|
||||
class TestExecuteConditionNode:
|
||||
|
||||
@pytest.mark.unit
|
||||
|
||||
@@ -42,3 +42,21 @@ def test_parse_structured_output_empty():
|
||||
def test_strip_json_fence_prefers_fence_over_braces():
|
||||
out = WorkflowEngine._strip_json_fence('lead {"outer": 1} ```json\n{"in": 2}\n```')
|
||||
assert out == '{"in": 2}'
|
||||
|
||||
|
||||
def test_parse_structured_output_prose_wrapped_array():
|
||||
ok, val = _engine()._parse_structured_output(
|
||||
'Here are the results: [{"name": "a"}, {"name": "b"}]'
|
||||
)
|
||||
assert ok and val == [{"name": "a"}, {"name": "b"}]
|
||||
|
||||
|
||||
def test_strip_json_fence_earliest_bracket_wins_for_array():
|
||||
# A top-level array must not be mis-sliced from its first "{" to its last
|
||||
# "}" (which would drop the array framing or extract an inner object).
|
||||
assert WorkflowEngine._strip_json_fence('x [1, 2, {"k": 3}] y') == '[1, 2, {"k": 3}]'
|
||||
|
||||
|
||||
def test_strip_json_fence_earliest_bracket_wins_for_object():
|
||||
# An object that opens before any "[" still wins.
|
||||
assert WorkflowEngine._strip_json_fence('x {"a": [1, 2]} y') == '{"a": [1, 2]}'
|
||||
@@ -225,6 +225,62 @@ class TestLLMHandler:
|
||||
assert chunks == ["chunk1", "chunk2"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# handle_tool_calls: assistant tool_call id <-> tool result id pairing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _StubToolExecutor:
|
||||
def __init__(self):
|
||||
self._name_to_tool = {}
|
||||
|
||||
def check_pause(self, tools_dict, call, llm_class_name):
|
||||
return None
|
||||
|
||||
|
||||
class _StubAgent:
|
||||
"""Minimal agent driving handle_tool_calls' success path."""
|
||||
|
||||
def __init__(self, resolved_call_id):
|
||||
self.llm = Mock()
|
||||
self.tool_executor = _StubToolExecutor()
|
||||
self._resolved_call_id = resolved_call_id
|
||||
|
||||
def _execute_tool_action(self, tools_dict, call):
|
||||
# Mirror tool_executor.execute: returns (result, call_id) where call_id
|
||||
# is a synthesized UUID when the provider omitted the tool-call id.
|
||||
yield from ()
|
||||
return {"ok": True}, self._resolved_call_id
|
||||
|
||||
|
||||
class TestHandleToolCallsIdPairing:
|
||||
def _drain(self, gen):
|
||||
try:
|
||||
while True:
|
||||
next(gen)
|
||||
except StopIteration as e:
|
||||
return e.value
|
||||
|
||||
def test_tool_result_id_matches_assistant_when_provider_omits_id(self):
|
||||
# Provider returned an empty tool-call id; the executor synthesizes one.
|
||||
handler = ConcreteHandler()
|
||||
call = ToolCall(id="", name="run_code", arguments={"code": "1"})
|
||||
resolved = "resolved-uuid-123"
|
||||
agent = _StubAgent(resolved)
|
||||
|
||||
updated_messages, pending = self._drain(
|
||||
handler.handle_tool_calls(agent, [call], tools_dict={}, messages=[])
|
||||
)
|
||||
|
||||
assistant = next(m for m in updated_messages if m["role"] == "assistant")
|
||||
tool_msg = next(m for m in updated_messages if m["role"] == "tool")
|
||||
assert assistant["tool_calls"][0]["id"] == resolved
|
||||
# The tool result must carry the SAME id as the assistant tool_call,
|
||||
# not the empty provider id — otherwise the next completion 400s.
|
||||
assert tool_msg["tool_call_id"] == resolved
|
||||
assert pending is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _append_unsupported_attachments
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in new issue
Block a user