Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 47 additions & 19 deletions packages/core/src/llm-core/platform/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -239,23 +239,34 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
const latestTokenUsage = this._createTokenUsageTracker()
let stream: AsyncGenerator<ChatGenerationChunk> | null = null
let hasChunk = false
let hasResponse = false
let hasToolCallChunk = false

try {
stream = await this._createStream(streamParams)

for await (const chunk of stream) {
hasToolCallChunk =
this._handleStreamChunk(
chunk,
runManager,
latestTokenUsage
) || hasToolCallChunk
const hasTool = this._handleStreamChunk(
chunk,
runManager,
latestTokenUsage
)
hasToolCallChunk = hasTool || hasToolCallChunk
hasChunk = true
hasResponse =
hasResponse ||
this._hasResponse(
chunk.message as AIMessage | AIMessageChunk
)
yield chunk
}

this._ensureChunksReceived(hasChunk)
if (!hasResponse) {
throw new ChatLunaError(
ChatLunaErrorCode.API_REQUEST_FAILED
)
}

this._finalizeStream(
hasToolCallChunk,
latestTokenUsage,
Expand Down Expand Up @@ -284,7 +295,7 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
logger.debug(
`Stream failed before first chunk (attempt ${attempt + 1}/${maxRetries}), retrying...`
)
await sleep(2000)
await sleep(2000 * 2 ** attempt)
}
}
}
Expand Down Expand Up @@ -322,14 +333,29 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
return hasToolCallChunk
}

private _hasToolCallChunk(message?: AIMessageChunk): boolean {
private _hasToolCallChunk(message?: AIMessage | AIMessageChunk): boolean {
return (
(message?.tool_calls?.length ?? 0) > 0 ||
(message?.tool_call_chunks?.length ?? 0) > 0 ||
((message as AIMessageChunk | undefined)?.tool_call_chunks
?.length ?? 0) > 0 ||
(message?.invalid_tool_calls?.length ?? 0) > 0
)
}

private _hasResponse(message?: AIMessage | AIMessageChunk): boolean {
const content = message?.content

return (
(typeof content === 'string'
? content.trim().length > 0
: Array.isArray(content) && content.length > 0) ||
this._hasToolCallChunk(message) ||
((message?.additional_kwargs?.tool_calls as unknown[] | undefined)
?.length ?? 0) > 0 ||
message?.additional_kwargs?.function_call != null
)
}

private _updateTokenUsageFromChunk(
chunk: ChatGenerationChunk,
latestTokenUsage: TokenUsageTracker
Expand All @@ -347,14 +373,6 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
latestTokenUsage.output_token_details = usage.output_token_details
}

private _ensureChunksReceived(hasChunk: boolean) {
if (hasChunk) {
return
}

throw new ChatLunaError(ChatLunaErrorCode.API_REQUEST_FAILED)
}

private _finalizeStream(
hasToolCallChunk: boolean,
latestTokenUsage: TokenUsageTracker,
Expand Down Expand Up @@ -495,6 +513,16 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
})
}

if (
!this._hasResponse(
response.message as AIMessage | AIMessageChunk
)
) {
throw new ChatLunaError(
ChatLunaErrorCode.API_REQUEST_FAILED
)
}

return response
} catch (error) {
if (
Expand All @@ -505,7 +533,7 @@ export class ChatLunaChatModel extends BaseChatModel<ChatLunaModelCallOptions> {
throw error
}

await sleep(2000)
await sleep(2000 * 2 ** attempt)
}
}

Expand Down
Loading