fix: add span for prompt template rendering

[MichaelDoyle]

Fixes #2064

Checklist (if applicable):

• PR title is following https://www.conventionalcommits.org/en/v1.0.0/
• Tested (manually, unit tested, etc.)
• Docs updated (updated docs or a docs bug required)

[pavelgj]

Can you share the code snippet that produces the trace in the screenshot?

[MichaelDoyle]

hello.history.prompt

---
model: googleai/gemini-1.5-flash
config:
  maxOutputTokens: 2048
  temperature: 0.6
  topK: 16
  topP: 0.95
  safetySettings:
    - category: HARM_CATEGORY_HATE_SPEECH
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_DANGEROUS_CONTENT
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_HARASSMENT
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_SEXUALLY_EXPLICIT
      threshold: BLOCK_ONLY_HIGH
input:
  schema:
    name: string
    persona?: string
  default:
    persona: Space Pirate
---

{{role "system"}}
You are a helpful AI assistant that really loves to make impressions.
{{role "user"}}
Say hello to Michael in the voice of a Space Pirate.
{{role "model"}}
Shiver me timbers, matey! We be sailing the solar winds!
{{role "user"}}
Say hello to {{name}} in the voice of a {{persona}}.

flowDotPromptHistory

export const HelloSchema = z.object({
  name: z.string(),
  persona: z.string().optional(),
});

ai.defineFlow(
  {
    name: 'flowDotPromptHistory',
    inputSchema: HelloSchema,
    outputSchema: z.any(),
  },
  async (input) => {
    const hello = ai.prompt('hello', {
      variant: 'history',
    });
    return (await hello(input)).text;
  }
);

[schnecle]

This will be hugely helpful for understanding when rendering issues are causing problems. Thanks Mike!

[pavelgj]

hello.history.prompt

---
model: googleai/gemini-1.5-flash
config:
  maxOutputTokens: 2048
  temperature: 0.6
  topK: 16
  topP: 0.95
  safetySettings:
    - category: HARM_CATEGORY_HATE_SPEECH
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_DANGEROUS_CONTENT
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_HARASSMENT
      threshold: BLOCK_ONLY_HIGH
    - category: HARM_CATEGORY_SEXUALLY_EXPLICIT
      threshold: BLOCK_ONLY_HIGH
input:
  schema:
    name: string
    persona?: string
  default:
    persona: Space Pirate
---

{{role "system"}}
You are a helpful AI assistant that really loves to make impressions.
{{role "user"}}
Say hello to Michael in the voice of a Space Pirate.
{{role "model"}}
Shiver me timbers, matey! We be sailing the solar winds!
{{role "user"}}
Say hello to {{name}} in the voice of a {{persona}}.

flowDotPromptHistory

export const HelloSchema = z.object({
  name: z.string(),
  persona: z.string().optional(),
});

ai.defineFlow(
  {
    name: 'flowDotPromptHistory',
    inputSchema: HelloSchema,
    outputSchema: z.any(),
  },
  async (input) => {
    const hello = ai.prompt('hello', {
      variant: 'history',
    });
    return (await hello(input)).text;
  }
);

So, what I'm thinking is that perhaps the traces should look something like this:

hello.history
  |- render
  |- generate
    |- googleai/gemini-1.5-flash

os do you think it would get too verbose? extra span with redundant info?

Semi-related: I was also chatting with Shruti yesterday... I think we should change render to return GenereateActionOptions instead of GenerateRequest....

[MichaelDoyle]

I think that's the right approach. Been swept up in some other things, but will try and knock this out today.

[MichaelDoyle]

@pavelgj I got this going for the generate case. I wasn't exactly sure what to do for prompt.stream(), though, since generateStreaming is sync while runInNewSpan returns a Promise.

[pavelgj]

@pavelgj I got this going for the generate case. I wasn't exactly sure what to do for prompt.stream(), though, since generateStreaming is sync while runInNewSpan returns a Promise.

I swear I replied to this.... sigh....

One thing I can think of is to use the Channel approach used in generateStream to wrap it manually:

genkit/js/ai/src/generate.ts

Line 466 in 238b656

let channel = new Channel<GenerateResponseChunk>();

[schnecle]

We might need to add to our action types in the trace and log handling in the plugin. We might still omit the input/output for this one. That can definitely be a followup item though.

[MichaelDoyle]

@schnecle same applies to the generate util as well, correct? maybe we should do a separate pass on metrics. I think theres a couple things there.