22// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
33// See the LICENSE file in the project root for more information
44
5+ using System . Diagnostics ;
56using Microsoft . Extensions . Logging ;
67
78namespace Elastic . Documentation . Api . Core . AskAi ;
@@ -11,11 +12,46 @@ public class AskAiUsecase(
1112 IStreamTransformer streamTransformer ,
1213 ILogger < AskAiUsecase > logger )
1314{
15+ private static readonly ActivitySource AskAiActivitySource = new ( "Elastic.Documentation.Api.AskAi" ) ;
16+
1417 public async Task < Stream > AskAi ( AskAiRequest askAiRequest , Cancel ctx )
1518 {
19+ // Start activity for the chat request - DO NOT use 'using' because the stream is consumed later
20+ // The activity will be passed to the transformer which will dispose it when the stream completes
21+ var activity = AskAiActivitySource . StartActivity ( "chat" , ActivityKind . Client ) ;
22+
23+ // Generate a correlation ID for tracking if this is a new conversation
24+ // For first messages (no ThreadId), generate a temporary ID that will be updated when the provider responds
25+ var correlationId = askAiRequest . ThreadId ?? $ "temp-{ Guid . NewGuid ( ) } ";
26+
27+ // Set GenAI semantic convention attributes
28+ _ = ( activity ? . SetTag ( "gen_ai.operation.name" , "chat" ) ) ;
29+ _ = ( activity ? . SetTag ( "gen_ai.conversation.id" , correlationId ) ) ; // Will be updated when we receive ConversationStart with actual ID
30+ _ = ( activity ? . SetTag ( "gen_ai.usage.input_tokens" , askAiRequest . Message . Length ) ) ; // Approximate token count
31+
32+ // Custom attributes for tracking our abstraction layer
33+ // We use custom attributes because we don't know the actual GenAI provider (OpenAI, Anthropic, etc.)
34+ // or model (gpt-4, claude, etc.) - those are abstracted by AgentBuilder/LlmGateway
35+ _ = ( activity ? . SetTag ( "docs.ai.gateway" , streamTransformer . AgentProvider ) ) ; // agent-builder or llm-gateway
36+ _ = ( activity ? . SetTag ( "docs.ai.agent_name" , streamTransformer . AgentId ) ) ; // docs-agent or docs_assistant
37+
38+ // Add GenAI prompt event
39+ _ = ( activity ? . AddEvent ( new ActivityEvent ( "gen_ai.content.prompt" ,
40+ timestamp : DateTimeOffset . UtcNow ,
41+ tags :
42+ [
43+ new KeyValuePair < string , object ? > ( "gen_ai.prompt" , askAiRequest . Message ) ,
44+ new KeyValuePair < string , object ? > ( "gen_ai.system_instructions" , AskAiRequest . SystemPrompt )
45+ ] ) ) ) ;
46+
1647 logger . LogDebug ( "Processing AskAiRequest: {Request}" , askAiRequest ) ;
48+
1749 var rawStream = await askAiGateway . AskAi ( askAiRequest , ctx ) ;
18- return await streamTransformer . TransformAsync ( rawStream , ctx ) ;
50+
51+ // The stream transformer will handle disposing the activity when streaming completes
52+ var transformedStream = await streamTransformer . TransformAsync ( rawStream , activity , ctx ) ;
53+
54+ return transformedStream ;
1955 }
2056}
2157
0 commit comments