@forge/llm 0.1.0 → 0.2.0-next.1

package/README.md

@@ -67,6 +67,11 @@ The response is a `ChatResponse` object:
 
 ```json
 {
+  "usage": {
+    "input_tokens": 150,
+    "output_tokens": 50,
+    "total_tokens": 200
+  },
   "choices": [
     {
       "finish_reason": "tool_use",
@@ -75,7 +80,7 @@ The response is a `ChatResponse` object:
         "content": [
           {
             "type": "text",
-            "text": "I'll help you check the weather in Melbourne using the WeatherTool." 
+            "text": "I'll help you check the weather in Melbourne using the WeatherTool."
           }
         ],
         "tool_calls": [

package/out/__test__/llm-api.test.js

@@ -1,9 +1,11 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
+const tslib_1 = require("tslib");
 const fetch_wrapper_1 = require("../fetch-wrapper");
 const llm_api_1 = require("../llm-api");
 const errors_1 = require("../errors");
 const test_helpers_1 = require("./test-helpers");
+const responseMapper = tslib_1.__importStar(require("../response-mapper"));
 function setupEnvironment(response) {
     const mockedFetch = jest.fn().mockResolvedValue(response || new Response(null, { status: 200 }));
     global.__forge_fetch__ = jest.fn((_ctx, path, options) => mockedFetch(path, options));
@@ -84,21 +86,11 @@ describe('llm api', () => {
         setupEnvironment(apiResponse);
         const llmAPI = new llm_api_1.LlmApiImpl((0, fetch_wrapper_1.getFetchWrapper)());
         const prompt = (0, test_helpers_1.constructPrompt)();
+        const mapResponseSpy = jest.spyOn(responseMapper, 'mapForgeLLMResponse');
         const response = await llmAPI.chat(prompt);
-        expect(response.choices[0]).toEqual({
-            finish_reason: 'tool_calls',
-            message: {
-                content: [
-                    {
-                        type: 'text',
-                        text: "It's 20 degrees celsius"
-                    }
-                ],
-                role: 'assistant',
-                tool_calls: undefined
-            },
-            index: 19
-        });
+        expect(mapResponseSpy).toHaveBeenCalledWith(apiResponse);
+        expect(response.choices).toEqual(llmResponse.choices);
+        expect(response.usage).toEqual(llmResponse.usage);
     });
     it('should throw an error for a non-200 response', async () => {
         const apiResponse = new Response(JSON.stringify({}), {

package/out/__test__/response-mapper.test.js

@@ -7,7 +7,7 @@ describe('response mapper', () => {
         const llmResponseBody = { choices: [] };
         const apiResponse = (0, test_helpers_1.makeApiResponse)({ json: () => Promise.resolve(llmResponseBody) });
         const mappedResponse = await (0, response_mapper_1.mapForgeLLMResponse)(apiResponse);
-        expect(mappedResponse).toEqual({ choices: [] });
+        expect(mappedResponse).toEqual({ choices: [], usage: undefined });
     });
     it('should map the Forge LLM response', async () => {
         const llmResponseBody = {
@@ -40,6 +40,34 @@ describe('response mapper', () => {
         };
         const apiResponse = (0, test_helpers_1.makeApiResponse)({ json: () => Promise.resolve(llmResponseBody) });
         const mappedResponse = await (0, response_mapper_1.mapForgeLLMResponse)(apiResponse);
-        expect(mappedResponse).toEqual({ choices: llmResponseBody.choices });
+        expect(mappedResponse).toEqual({ choices: llmResponseBody.choices, usage: undefined });
+    });
+    it('should map the Forge LLM response with usage attribute', async () => {
+        const llmResponseBody = {
+            choices: [
+                {
+                    finish_reason: 'stop',
+                    message: {
+                        role: 'assistant',
+                        content: [{ type: 'text', text: 'Hello!' }]
+                    }
+                }
+            ],
+            usage: {
+                input_token: 10,
+                output_token: 20,
+                total_token: 30
+            }
+        };
+        const apiResponse = (0, test_helpers_1.makeApiResponse)({ json: () => Promise.resolve(llmResponseBody) });
+        const mappedResponse = await (0, response_mapper_1.mapForgeLLMResponse)(apiResponse);
+        expect(mappedResponse).toEqual({
+            choices: llmResponseBody.choices,
+            usage: {
+                input_token: 10,
+                output_token: 20,
+                total_token: 30
+            }
+        });
     });
 });

package/out/__test__/test-helpers.d.ts

@@ -15,5 +15,6 @@ export declare const constructLlmResponse: (overrides?: Partial<Choice>) => {
             }[];
         };
     }[];
+    usage: {};
 };
 //# sourceMappingURL=test-helpers.d.ts.map

package/out/__test__/test-helpers.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"test-helpers.d.ts","sourceRoot":"","sources":["../../src/__test__/test-helpers.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,MAAM,qBAAqB,CAAC;AAC7C,OAAO,EAAE,WAAW,EAAE,MAAM,YAAY,CAAC;AACzC,OAAO,EAAE,MAAM,EAAE,MAAM,uBAAuB,CAAC;AAE/C,eAAO,MAAM,eAAe,sBAAqB,MAU/C,CAAC;AAEH,eAAO,MAAM,eAAe,cAAe,QAAQ,WAAW,CAAC,KAAG,WAWjE,CAAC;AAEF,eAAO,MAAM,oBAAoB,eAAe,QAAQ,MAAM,CAAC;;;;;;;;;;;;CAiB7D,CAAC"}
+{"version":3,"file":"test-helpers.d.ts","sourceRoot":"","sources":["../../src/__test__/test-helpers.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,MAAM,EAAE,MAAM,qBAAqB,CAAC;AAC7C,OAAO,EAAE,WAAW,EAAE,MAAM,YAAY,CAAC;AACzC,OAAO,EAAE,MAAM,EAAE,MAAM,uBAAuB,CAAC;AAE/C,eAAO,MAAM,eAAe,sBAAqB,MAU/C,CAAC;AAEH,eAAO,MAAM,eAAe,cAAe,QAAQ,WAAW,CAAC,KAAG,WAWjE,CAAC;AAEF,eAAO,MAAM,oBAAoB,eAAe,QAAQ,MAAM,CAAC;;;;;;;;;;;;;CAkB7D,CAAC"}

package/out/__test__/test-helpers.js

@@ -42,6 +42,7 @@ const constructLlmResponse = (overrides = {}) => ({
             index: 19,
             ...overrides
         }
-    ]
+    ],
+    usage: {}
 });
 exports.constructLlmResponse = constructLlmResponse;

package/out/interfaces/llm-api.d.ts

@@ -25,14 +25,6 @@ export declare type ToolChoice = 'auto' | 'none' | 'required' | {
         name: string;
     };
 };
-export interface LlmResponseTool {
-    id: string;
-    type: 'function';
-    function: {
-        name: string;
-        arguments: Record<string, unknown>;
-    };
-}
 export interface SystemMessage {
     content: Content;
     role: 'system';
@@ -63,12 +55,18 @@ export interface ToolCall {
 }
 export interface LlmResponse {
     choices: Choice[];
+    usage?: Usage;
 }
 export interface Choice {
     finish_reason: string;
     index?: number;
     message: AssistantMessage;
 }
+export interface Usage {
+    input_token?: number;
+    output_token?: number;
+    total_token?: number;
+}
 export declare type Content = string | ContentPart[];
 export declare type ContentPart = TextPart;
 export interface TextPart {

package/out/interfaces/llm-api.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"llm-api.d.ts","sourceRoot":"","sources":["../../src/interfaces/llm-api.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,YAAY,EAAE,MAAM,EAAE,MAAM,SAAS,CAAC;AAE/C,MAAM,WAAW,MAAM;IACrB,IAAI,CAAC,MAAM,EAAE,MAAM,GAAG,OAAO,CAAC,YAAY,CAAC,CAAC;CAC7C;AAED,MAAM,WAAW,UAAU;IACzB,QAAQ,EAAE,OAAO,EAAE,CAAC;IACpB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,qBAAqB,CAAC,EAAE,MAAM,CAAC;IAC/B,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,KAAK,CAAC,EAAE,IAAI,EAAE,CAAC;IACf,WAAW,CAAC,EAAE,UAAU,CAAC;CAC1B;AAED,oBAAY,OAAO,GAAG,aAAa,GAAG,WAAW,GAAG,gBAAgB,GAAG,WAAW,CAAC;AAEnF,MAAM,WAAW,IAAI;IACnB,IAAI,EAAE,UAAU,CAAC;IACjB,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,WAAW,EAAE,MAAM,CAAC;QACpB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;KACrC,CAAC;CACH;AAED,oBAAY,UAAU,GAClB,MAAM,GACN,MAAM,GACN,UAAU,GACV;IACE,IAAI,EAAE,UAAU,CAAC;IACjB,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;KACd,CAAC;CACH,CAAC;AAEN,MAAM,WAAW,eAAe;IAC9B,EAAE,EAAE,MAAM,CAAC;IACX,IAAI,EAAE,UAAU,CAAC;IACjB,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,SAAS,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;KACpC,CAAC;CACH;AAED,MAAM,WAAW,aAAa;IAC5B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,QAAQ,CAAC;CAChB;AAED,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,MAAM,CAAC;CACd;AACD,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,MAAM,CAAC;IACb,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,IAAI,CAAC,EAAE,MAAM,CAAC;CACf;AAED,MAAM,WAAW,gBAAgB;IAC/B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,WAAW,CAAC;IAClB,UAAU,CAAC,EAAE,QAAQ,EAAE,CAAC;CACzB;AAED,MAAM,WAAW,QAAQ;IACvB,EAAE,EAAE,MAAM,CAAC;IACX,IAAI,EAAE,UAAU,CAAC;IACjB,KAAK,EAAE,MAAM,CAAC;IACd,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,SAAS,EAAE,MAAM,CAAC;KACnB,CAAC;CACH;AAED,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,MAAM,EAAE,CAAC;CACnB;AAED,MAAM,WAAW,MAAM;IACrB,aAAa,EAAE,MAAM,CAAC;IACtB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,OAAO,EAAE,gBAAgB,CAAC;CAC3B;AAED,oBAAY,OAAO,GAAG,MAAM,GAAG,WAAW,EAAE,CAAC;AAE7C,oBAAY,WAAW,GAAG,QAAQ,CAAC;AAEnC,MAAM,WAAW,QAAQ;IACvB,IAAI,EAAE,MAAM,CAAC;IACb,IAAI,EAAE,MAAM,CAAC;CACd"}
+{"version":3,"file":"llm-api.d.ts","sourceRoot":"","sources":["../../src/interfaces/llm-api.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,YAAY,EAAE,MAAM,EAAE,MAAM,SAAS,CAAC;AAE/C,MAAM,WAAW,MAAM;IACrB,IAAI,CAAC,MAAM,EAAE,MAAM,GAAG,OAAO,CAAC,YAAY,CAAC,CAAC;CAC7C;AAED,MAAM,WAAW,UAAU;IACzB,QAAQ,EAAE,OAAO,EAAE,CAAC;IACpB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,qBAAqB,CAAC,EAAE,MAAM,CAAC;IAC/B,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,KAAK,CAAC,EAAE,IAAI,EAAE,CAAC;IACf,WAAW,CAAC,EAAE,UAAU,CAAC;CAC1B;AAED,oBAAY,OAAO,GAAG,aAAa,GAAG,WAAW,GAAG,gBAAgB,GAAG,WAAW,CAAC;AAEnF,MAAM,WAAW,IAAI;IACnB,IAAI,EAAE,UAAU,CAAC;IACjB,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,WAAW,EAAE,MAAM,CAAC;QACpB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;KACrC,CAAC;CACH;AAED,oBAAY,UAAU,GAClB,MAAM,GACN,MAAM,GACN,UAAU,GACV;IACE,IAAI,EAAE,UAAU,CAAC;IACjB,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;KACd,CAAC;CACH,CAAC;AAEN,MAAM,WAAW,aAAa;IAC5B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,QAAQ,CAAC;CAChB;AAED,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,MAAM,CAAC;CACd;AACD,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,MAAM,CAAC;IACb,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,IAAI,CAAC,EAAE,MAAM,CAAC;CACf;AAED,MAAM,WAAW,gBAAgB;IAC/B,OAAO,EAAE,OAAO,CAAC;IACjB,IAAI,EAAE,WAAW,CAAC;IAClB,UAAU,CAAC,EAAE,QAAQ,EAAE,CAAC;CACzB;AAED,MAAM,WAAW,QAAQ;IACvB,EAAE,EAAE,MAAM,CAAC;IACX,IAAI,EAAE,UAAU,CAAC;IACjB,KAAK,EAAE,MAAM,CAAC;IACd,QAAQ,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,SAAS,EAAE,MAAM,CAAC;KACnB,CAAC;CACH;AAED,MAAM,WAAW,WAAW;IAC1B,OAAO,EAAE,MAAM,EAAE,CAAC;IAClB,KAAK,CAAC,EAAE,KAAK,CAAC;CACf;AAED,MAAM,WAAW,MAAM;IACrB,aAAa,EAAE,MAAM,CAAC;IACtB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,OAAO,EAAE,gBAAgB,CAAC;CAC3B;AAED,MAAM,WAAW,KAAK;IACpB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,WAAW,CAAC,EAAE,MAAM,CAAC;CACtB;AAED,oBAAY,OAAO,GAAG,MAAM,GAAG,WAAW,EAAE,CAAC;AAE7C,oBAAY,WAAW,GAAG,QAAQ,CAAC;AAEnC,MAAM,WAAW,QAAQ;IACvB,IAAI,EAAE,MAAM,CAAC;IACb,IAAI,EAAE,MAAM,CAAC;CACd"}

package/out/interfaces/types.d.ts

@@ -1,10 +1,10 @@
 import { ForgeLlmError, PromptValidationError } from '../errors';
-import { LlmRequest, LlmResponse, Message, Tool, ToolCall } from './llm-api';
+import { LlmRequest, LlmResponse, Message, Tool, ToolCall, Usage } from './llm-api';
 export declare type Prompt = LlmRequest & {
     model: string;
 };
 export declare type ChatResponse = LlmResponse;
 export declare type ToolSchema = Tool;
-export type { Message, ToolCall };
+export type { Message, ToolCall, Usage };
 export { ForgeLlmError, PromptValidationError };
 //# sourceMappingURL=types.d.ts.map

package/out/interfaces/types.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../../src/interfaces/types.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,aAAa,EAAE,qBAAqB,EAAE,MAAM,WAAW,CAAC;AACjE,OAAO,EAAE,UAAU,EAAE,WAAW,EAAE,OAAO,EAAE,IAAI,EAAE,QAAQ,EAAE,MAAM,WAAW,CAAC;AAE7E,oBAAY,MAAM,GAAG,UAAU,GAAG;IAChC,KAAK,EAAE,MAAM,CAAC;CACf,CAAC;AAEF,oBAAY,YAAY,GAAG,WAAW,CAAC;AAEvC,oBAAY,UAAU,GAAG,IAAI,CAAC;AAE9B,YAAY,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAC;AAElC,OAAO,EAAE,aAAa,EAAE,qBAAqB,EAAE,CAAC"}
+{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../../src/interfaces/types.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,aAAa,EAAE,qBAAqB,EAAE,MAAM,WAAW,CAAC;AACjE,OAAO,EAAE,UAAU,EAAE,WAAW,EAAE,OAAO,EAAE,IAAI,EAAE,QAAQ,EAAE,KAAK,EAAE,MAAM,WAAW,CAAC;AAEpF,oBAAY,MAAM,GAAG,UAAU,GAAG;IAChC,KAAK,EAAE,MAAM,CAAC;CACf,CAAC;AAEF,oBAAY,YAAY,GAAG,WAAW,CAAC;AAEvC,oBAAY,UAAU,GAAG,IAAI,CAAC;AAE9B,YAAY,EAAE,OAAO,EAAE,QAAQ,EAAE,KAAK,EAAE,CAAC;AAEzC,OAAO,EAAE,aAAa,EAAE,qBAAqB,EAAE,CAAC"}

package/out/response-mapper.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"response-mapper.d.ts","sourceRoot":"","sources":["../src/response-mapper.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,WAAW,EAAE,MAAM,YAAY,CAAC;AACzC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAGlD,eAAO,MAAM,mBAAmB,aAAoB,WAAW,KAAG,QAAQ,YAAY,CAIrF,CAAC"}
+{"version":3,"file":"response-mapper.d.ts","sourceRoot":"","sources":["../src/response-mapper.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,WAAW,EAAE,MAAM,YAAY,CAAC;AACzC,OAAO,EAAE,YAAY,EAAE,MAAM,oBAAoB,CAAC;AAElD,eAAO,MAAM,mBAAmB,aAAoB,WAAW,KAAG,QAAQ,YAAY,CAGrF,CAAC"}

package/out/response-mapper.js

@@ -2,8 +2,7 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.mapForgeLLMResponse = void 0;
 const mapForgeLLMResponse = async (response) => {
-    const jsonBody = await response.json();
-    const choices = jsonBody.choices ?? [];
-    return { choices };
+    const chatResponse = await response.json();
+    return chatResponse;
 };
 exports.mapForgeLLMResponse = mapForgeLLMResponse;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@forge/llm",
-  "version": "0.1.0",
+  "version": "0.2.0-next.1",
   "description": "Forge LLM SDK",
   "main": "out/index.js",
   "types": "out/index.d.ts",
@@ -10,7 +10,7 @@
   "author": "Atlassian",
   "license": "SEE LICENSE IN LICENSE.txt",
   "dependencies": {
-    "@forge/api": "^6.2.0"
+    "@forge/api": "^6.2.1-next.0"
   },
   "devDependencies": {
     "@types/node": "20.19.1",