pacnpal/Roo-Code
1
0
Fork 0
mirror of https://github.com/pacnpal/Roo-Code.git synced 2025-12-21 12:51:17 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add typing for usageData and adds some unit tests

Browse Source
This commit is contained in:
Vignesh Subbiah
2025-01-29 22:53:43 +05:30
parent 6ccb061d33
commit d6433591b2
2 changed files with 54 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/api/providers/unbound.ts
View File

@@ -3,7 +3,12 @@ import OpenAI from "openai"
import { ApiHandler, SingleCompletionHandler } from "../"
import { ApiHandlerOptions, ModelInfo, UnboundModelId, unboundDefaultModelId, unboundModels } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
interface UnboundUsage extends OpenAI.CompletionUsage {
cache_creation_input_tokens?: number
cache_read_input_tokens?: number
}
export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
private options: ApiHandlerOptions
@@ -96,7 +101,7 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
for await (const chunk of completion) {
const delta = chunk.choices[0]?.delta
const usage = chunk.usage
const usage = chunk.usage as UnboundUsage
if (delta?.content) {
yield {
@@ -106,18 +111,18 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
}
if (usage) {
const usageData: any = {
const usageData: ApiStreamUsageChunk = {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
outputTokens: usage?.completion_tokens || 0,
inputTokens: usage.prompt_tokens || 0,
outputTokens: usage.completion_tokens || 0,
}
// Only add cache tokens if they exist
if ((usage as any)?.cache_creation_input_tokens) {
usageData.cacheWriteTokens = (usage as any).cache_creation_input_tokens
if (usage.cache_creation_input_tokens) {
usageData.cacheWriteTokens = usage.cache_creation_input_tokens
}
if ((usage as any)?.cache_read_input_tokens) {
usageData.cacheReadTokens = (usage as any).cache_read_input_tokens
if (usage.cache_read_input_tokens) {
usageData.cacheReadTokens = usage.cache_read_input_tokens
}
yield usageData