How to use LanguageModelMiddleware to add a message annotation, OR access updated params by the caller?

[Und3rf10w]

Let's say you have a hypothetical LanguageModelMiddleware that performs RAG and as part of the transformParams function, appends/injects the context at the end of the user message.

It reads in configuration options via a ragMiddleware providerOption.

How can I either:

1. Access the param (e.g. params.providerMetadata that have been updated after running this transform?
2. Alternatively, how can I use wrapStream and wrapGenerate to add a source, message annotation, or data chunk to the response?

Here's a very basic example of what that middleware might do:

export const documentRagMiddleware: LanguageModelV1Middleware = {
  transformParams: async ({ params }: { type: 'generate' | 'stream', params: LanguageModelV1CallOptions }): Promise<LanguageModelV1CallOptions> => {
    const docRagParams = params?.providerMetadata?.docRag 
    
    if (docRag.enabled) {
        docRagParams.startTime = Date.now()

        let ragCandidates = []
        // ... do whatever RAG logic here, let's pretend it's in ragCandidates
        docRagParams.ragCandidates = ragCandidates
        
        
        // We can return information about the rag config and rag candiates in the now updated providerMetadata, but there's no real way to access it from the calling `streamText` that wrapped this middleware?
        
       return {
         ...params,
         providerMetadata: {
             ...params.providerMetadata,
             docRag: {
               ...docRagParams
             }
         }
       }
    } else {
        return { ...params }
    },
    // Alternatively, we try to use wrapStream and wrapGenerate to add either a source part but I have no idea how to do this right
    wrapStream: async ({ doStream, params }) => {
      const ragMetadata = params.providerMetadata?.docRag
  
      const transformStream = new TransformStream<
        LanguageModelV1StreamPart,
        LanguageModelV1StreamPart
      >({
        start(controller) {
          if (ragMetadata?.ragCandidates && ragMetadata?.ragCandidates?.length > 0) {
            for (const metadata of ragMetadata?.ragCandidates) {
              controller.enqueue({
                type: 'source',
                source: {
                  id: metadata?.chunkId ?? `chunk_${generateId()}`,
                  sourceType: 'url',
                  url: metadata?.docId ?? '',
                  title: metadata?.docName ?? '',
                  providerMetadata: {
                    documentMiddleware: {
                      ...ragMetadata
                    }
                  }
                }
              })
            }
          }
        },
        transform(chunk, controller) {
          controller.enqueue(chunk)
        },
      })
      const result = await doStream()
  
      return {
        ...result,
        stream: result.stream.pipeThrough(transformStream)
      }
    },
    wrapGenerate: async ({ doGenerate, params }) => {
      const { providerMetadata, ...rest } = await doGenerate()
      
      // This just updates the provider metadata for the request, but same problem, can't access it
      // Alternatively, if we could inject sources or annotations here, that would work as well.
  
      return {
        ...rest,
        providerMetadata
      }
    }
 }

I've also tried using formatDataStreamPart in the transform stream, e.g.:

controller.enqueue(formatDataStreamPart('message_annotations', {
 id: metadata?.chunkId ?? `chunk_${generateId()}`,
 sourceType: 'url',
 url: metadata?.docId ?? '',
 title: metadata?.docName ?? '',
 providerMetadata: {
   documentMiddleware: {
     ...ragMetadata
   }
 }
}) as any

We could attempt to access the updated providerMetadata in the caller's side, but this is using the provider metadata passed for the configuration, not the results:

   // ... initial providerMetadata was set here, languageModel is now wrappedLanguageModel with the docRag middleware
   return createDataStreamResponse({
     async execute(dataStream) {
      const stream = streamText({
        model: wrappedLanguageModel, // This has the docRag middleware applied
        async onFinish(event) {
            // event.providerMetadata contains old metadata, likely since it was stripped in the provider, so this will not work.
            console.log(event.providerMetadata)
            if (event?.providerMetadata?.documentMiddleware) {
             dataStream.writeMessageAnnotation({
               docRag: {
                 ...event?.providerMetadata?.docRag
               }
             })
           }
        }
      })
      stream.mergeIntoDataStream(dataStream)
    }
})

But that doesn't seem to work either.