34ff326 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(agent): timer should reset after each chunk

34ff326kujtimiihoxha1 year ago

internal/llm/agent/agent.go+7−1

compare with main

unifiedsplit

@@ -555,6 +555,9 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg

555 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

556

557 // Process each event in the stream.

558+ timer := time.NewTimer(streamChunkTimeout)

559+ defer timer.Stop()

560+

558loop:

561loop:

559 for {

562 for {

560 select {

563 select {

@@ -562,6 +565,9 @@ loop:

562 if !ok {

565 if !ok {

563 break loop

566 break loop

564 }

567 }

568+ // Reset the timeout timer since we received a chunk

569+ timer.Reset(streamChunkTimeout)

570+

565 if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {

571 if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {

566 if errors.Is(processErr, context.Canceled) {

572 if errors.Is(processErr, context.Canceled) {

567 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")

573 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")

@@ -570,7 +576,7 @@ loop:

570 }

576 }

571 return assistantMsg, nil, processErr

577 return assistantMsg, nil, processErr

572 }

578 }

573− case <-time.After(streamChunkTimeout):

579+ case <-timer.C:

574 a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "Stream timeout", "No chunk received within timeout")

580 a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "Stream timeout", "No chunk received within timeout")

575 return assistantMsg, nil, fmt.Errorf("stream chunk timeout")

581 return assistantMsg, nil, fmt.Errorf("stream chunk timeout")

576 case <-ctx.Done():

582 case <-ctx.Done():