Repository: google/adk-go
File: agent/llm_agent.go (Go)

Code before:
	RequestProcessors  []func(ctx context.Context, parentCtx *adk.InvocationContext, req *adk.LLMRequest) error
	ResponseProcessors []func(ctx context.Context, parentCtx *adk.InvocationContext, req *adk.LLMRequest, resp *adk.LLMResponse) error
}

func (f *baseFlow) Run(ctx context.Context, parentCtx *adk.InvocationContext) adk.EventStream {
	return func(yield func(*adk.Event, error) bool) {
		for {
			var lastEvent *adk.Event
			for ev, err := range f.runOneStep(ctx, parentCtx) {
				if err != nil {
					yield(nil, err)
					return
				}
				// forward the event first.
				if !yield(ev, nil) {
					return
				}
				lastEvent = ev
			}
			if lastEvent == nil || lastEvent.IsFinalResponse() {
				return
			}
			if lastEvent.LLMResponse.Partial {
				// We may reached max token limit during streaming mode.
				// TODO: handle Partial response in model level.
				if !yield(nil, fmt.Errorf("TODO: last event is not final")) {
					return
				}
			}
		}
	}
}

func (f *baseFlow) runOneStep(ctx context.Context, parentCtx *adk.InvocationContext) adk.EventStream {
	return func(yield func(*adk.Event, error) bool) {
		req := &adk.LLMRequest{Model: f.Model}

		// Preprocess before calling the LLM.
		if err := f.preprocess(ctx, parentCtx, req); err != nil {
			yield(nil, err)
			return
		}

		// Calls the LLM.
		for resp, err := range f.callLLM(ctx, parentCtx, req) {
			if err != nil {
				yield(nil, err)
				return
			}
			if err := f.postprocess(ctx, parentCtx, req, resp); err != nil {
				yield(nil, err)

Review comment (suggestion):
Sometimes an error stops the iteration (L135 above). Sometimes it doesn't (here). That's confusing. At a minimum, document this behavior. (Unless this is really a placeholder and the work is incomplete without it.)

Code after:
	RequestProcessors  []func(ctx context.Context, parentCtx *adk.InvocationContext, req *adk.LLMRequest) error
	ResponseProcessors []func(ctx context.Context, parentCtx *adk.InvocationContext, req *adk.LLMRequest, resp *adk.LLMResponse) error
}

func (f *baseFlow) Run(ctx context.Context, parentCtx *adk.InvocationContext) adk.EventStream {
	return func(yield func(*adk.Event, error) bool) {
		for {
			var lastEvent *adk.Event
			for ev, err := range f.runOneStep(ctx, parentCtx) {
				if err != nil {
					yield(nil, err)
					return
				}
				// forward the event first.
				if !yield(ev, nil) {
					return
				}
				lastEvent = ev
			}
			if lastEvent == nil || lastEvent.IsFinalResponse() {
				return
			}
			if lastEvent.LLMResponse.Partial {
				// We may have reached max token limit during streaming mode.
				// TODO: handle Partial response in model level. CL 781377328
				yield(nil, fmt.Errorf("TODO: last event is not final"))
				return
			}
		}
	}
}

func (f *baseFlow) runOneStep(ctx context.Context, parentCtx *adk.InvocationContext) adk.EventStream {
	return func(yield func(*adk.Event, error) bool) {
		req := &adk.LLMRequest{Model: f.Model}

		// Preprocess before calling the LLM.
		if err := f.preprocess(ctx, parentCtx, req); err != nil {
			yield(nil, err)
			return
		}

		// Calls the LLM.
		for resp, err := range f.callLLM(ctx, parentCtx, req) {
			if err != nil {
				yield(nil, err)
				return
			}
			if err := f.postprocess(ctx, parentCtx, req, resp); err != nil {
				yield(nil, err)
				return