11import type { NextRequest } from "next/server" ;
2+ import z from "zod" ;
23import { handler as corsHandler } from "~/app/api/_cors/cors" ;
34import { getSession } from "~/auth/server" ;
45import { env } from "~/env" ;
6+ import { api } from "~/trpc/server" ;
57
6- // import { api } from "~/trpc/server";
8+ const OPENAI_MESSAGE_PREFIX = "data: " ;
9+ const OPENAI_MESSAGE_DONE_TEXT = "[DONE]" ;
10+ const OPENAI_MODEL_PROVIDER = "openai" ;
711
8- const OPENAI_API_URL = "https://api.openai.com/v1/" ;
12+ const zChatCompletionMessage = z . object ( {
13+ choices : z . array ( z . unknown ( ) ) ,
14+ created : z . number ( ) ,
15+ id : z . string ( ) ,
16+ model : z . string ( ) ,
17+ obfuscation : z . string ( ) ,
18+ object : z . literal ( "chat.completion.chunk" ) ,
19+ service_tier : z . string ( ) ,
20+ system_fingerprint : z . string ( ) ,
21+ usage : z
22+ . object ( {
23+ completion_tokens : z . number ( ) ,
24+ completion_tokens_details : z . object ( {
25+ accepted_prediction_tokens : z . number ( ) ,
26+ audio_tokens : z . number ( ) ,
27+ reasoning_tokens : z . number ( ) ,
28+ rejected_prediction_tokens : z . number ( ) ,
29+ } ) ,
30+ prompt_tokens : z . number ( ) ,
31+ prompt_tokens_details : z . object ( {
32+ audio_tokens : z . number ( ) ,
33+ cached_tokens : z . number ( ) ,
34+ } ) ,
35+ total_tokens : z . number ( ) ,
36+ } )
37+ . nullable ( ) ,
38+ } ) ;
939
1040async function handler ( req : NextRequest ) {
1141 const session = await getSession ( ) ;
@@ -22,7 +52,7 @@ async function handler(req: NextRequest) {
2252 return new Response ( "Not found" , { status : 404 } ) ;
2353 }
2454
25- const openAIResponse = await fetch ( new URL ( url , OPENAI_API_URL ) , {
55+ const openAIResponse = await fetch ( new URL ( url , env . OPENAI_API_URL ) , {
2656 body : JSON . stringify ( {
2757 ...requestBody ,
2858 stream : true ,
@@ -40,22 +70,46 @@ async function handler(req: NextRequest) {
4070
4171 const transformStream = new TransformStream ( {
4272 async transform ( chunk , controller ) {
43- const text = new TextDecoder ( ) . decode ( chunk ) ;
44-
45- // TODO: Parse chunk and determine if it's a usage chunk.
46- console . log ( "Intercepted chunk:" , text ) ;
47-
48- // ! TODO: Track usage.
49- // api.usage.trackModelUsage({
50- // model: requestBody.model,
51- // provider: requestBody.provider,
52- // user_id: session.user.id,
53- // quantity: requestBody.stream_options.include_usage ? 1 : 0,
54- // unit: "output_tokens",
55- // });
56-
57- // TODO: Do not forward the usage chunk.
58- controller . enqueue ( chunk ) ;
73+ try {
74+ const text = new TextDecoder ( ) . decode ( chunk ) ;
75+
76+ if ( text === `${ OPENAI_MESSAGE_PREFIX } ${ OPENAI_MESSAGE_DONE_TEXT } ) {
77+ return controller . enqueue ( chunk ) ;
78+ }
79+
80+ // Parsing the messages from the openAI response by removing the newline characters and the prefix.
81+ const data = text
82+ . replace ( / [ \r \n ] + / g, "" )
83+ . trim ( )
84+ . split ( OPENAI_MESSAGE_PREFIX )
85+ . filter ( ( t ) => t !== "" && t !== OPENAI_MESSAGE_DONE_TEXT ) ;
86+
87+ for ( const json of data ) {
88+ const message = zChatCompletionMessage . parse ( JSON . parse ( json ) ) ;
89+
90+ if ( ! message . usage ) continue ;
91+
92+ await api . usage . trackModelUsage ( {
93+ metrics : [
94+ {
95+ quantity : message . usage . prompt_tokens ,
96+ unit : "input_tokens" ,
97+ } ,
98+ {
99+ quantity : message . usage . completion_tokens ,
100+ unit : "output_tokens" ,
101+ } ,
102+ ] ,
103+ model_id : message . model ,
104+ model_provider : OPENAI_MODEL_PROVIDER ,
105+ user_id : session . user . id ,
106+ } ) ;
107+ }
108+
109+ controller . enqueue ( chunk ) ;
110+ } catch ( error ) {
111+ console . error ( "Error tracking model usage" , error ) ;
112+ }
59113 } ,
60114 } ) ;
61115
0 commit comments