Merge pull request #18 from AmbireTech/feat/add-token-count-to-llm-ouput

LLM output: add input and output token count

Author: cmihaylov

lib/types/types.ts

@@ -55,6 +55,8 @@ export type LlmProcessOutput = {
         model: string
     }
     response: Strategy[] | null
+    inputTokens: number
+    outputTokens: number
     error?: string | null
 }
 

lib/utils/llm/gemini.ts

@@ -14,6 +14,8 @@ export async function callGemini(llmInput: LlmProcessProps): Promise<LlmProcessO
     let output = null
     let error = null
     const model = llmInput.model || GEMINI_MODELS.gemini20flashExp
+    let inputTokens: number | undefined
+    let outputTokens: number | undefined
 
     try {
         const aiModel = genAI.getGenerativeModel({
@@ -28,6 +30,8 @@ export async function callGemini(llmInput: LlmProcessProps): Promise<LlmProcessO
         // console.log(JSON.stringify(result))
 
         const content = result.response.text()
+        inputTokens = result.response.usageMetadata?.promptTokenCount || 0
+        outputTokens = result.response.usageMetadata?.candidatesTokenCount || 0
 
         try {
             output = JSON.parse(content || '[]') as Strategy[]
@@ -46,6 +50,8 @@ export async function callGemini(llmInput: LlmProcessProps): Promise<LlmProcessO
             model
         },
         response: output,
+        inputTokens: inputTokens || 0,
+        outputTokens: outputTokens || 0,
         error
     }
 }

lib/utils/llm/grok.ts

@@ -18,6 +18,8 @@ export async function callGrok(llmInput: LlmProcessProps): Promise<LlmProcessOut
     let output = null
     let error = null
     const model = llmInput.model || XAI_MODELS.grok3latest
+    let inputTokens: number | undefined
+    let outputTokens: number | undefined
 
     try {
         const completion = await apiClient.chat.completions.create({
@@ -36,6 +38,8 @@ export async function callGrok(llmInput: LlmProcessProps): Promise<LlmProcessOut
         })
 
         const outputContent = completion.choices[0].message.content || '{}'
+        inputTokens = completion.usage?.prompt_tokens || 0
+        outputTokens = completion.usage?.completion_tokens || 0
 
         try {
             const parsed = JSON.parse(outputContent) as { strategies: Strategy[] }
@@ -55,6 +59,8 @@ export async function callGrok(llmInput: LlmProcessProps): Promise<LlmProcessOut
             model
         },
         response: output,
+        inputTokens: inputTokens || 0,
+        outputTokens: outputTokens || 0,
         error
     }
 }

lib/utils/llm/mockedAI.ts

@@ -43,6 +43,8 @@ export async function llmMockProcess({ prompt }: LlmProcessProps): Promise<LlmPr
                 name: 'L2 Yield Farming',
                 risk: StrategyRisk.HIGH
             }
-        ]
+        ],
+        inputTokens: 0,
+        outputTokens: 0
     }
 }

lib/utils/portfolio.ts

@@ -94,7 +94,9 @@ export const processAddress = async (
                         provider: 'local',
                         model: 'local'
                     },
-                    response: EMPTY_PORTFOLIO_STRATEGIES
+                    response: EMPTY_PORTFOLIO_STRATEGIES,
+                    inputTokens: 0,
+                    outputTokens: 0
                 }
             ]
         }

test/unit/portfolio.test.ts

@@ -47,7 +47,9 @@ const mockedLlmOutput: LlmProcessOutput = {
             name: 'Example USDC strategy name',
             risk: StrategyRisk.LOW
         }
-    ]
+    ],
+    inputTokens: 0,
+    outputTokens: 0
 }
 
 jest.mock('ambire-common/dist/src/consts/networks', () => {