Allow setting default LLM args on instance level

Closes #20

Author: yamalight

packages/litlytics/engine/runPrompt.ts

@@ -1,14 +1,14 @@
-import type { CoreMessage, CoreTool } from 'ai';
+import type { CoreMessage } from 'ai';
 import type { LLMProviders } from '../litlytics';
 import { executeOnLLM } from '../llm/llm';
-import type { LLMModel, LLMProvider, LLMRequest } from '../llm/types';
+import type { LLMArgs, LLMModel, LLMProvider, LLMRequest } from '../llm/types';
 
 export interface RunPromptFromMessagesArgs {
   provider: LLMProviders;
   key: string;
   model: LLMModel;
   messages: CoreMessage[];
-  args?: Record<string, CoreTool>;
+  args?: LLMArgs;
 }
 export const runPromptFromMessages = async ({
   provider,
@@ -36,7 +36,7 @@ export interface RunPromptArgs {
   model: LLMModel;
   system: string;
   user: string;
-  args?: Record<string, CoreTool>;
+  args?: LLMArgs;
 }
 export const runPrompt = async ({
   provider,

packages/litlytics/litlytics.ts

@@ -9,7 +9,7 @@ import {
 import { runStep, type RunStepArgs } from './engine/runStep';
 import { runLLMStep, type RunLLMStepArgs } from './engine/step/runLLMStep';
 import { testPipelineStep } from './engine/testStep';
-import type { LLMModel, LLMProvider } from './llm/types';
+import type { LLMArgs, LLMModel, LLMProvider } from './llm/types';
 import { OUTPUT_ID } from './output/Output';
 import {
   pipelineFromText,
@@ -38,6 +38,7 @@ export { modelCosts } from './llm/costs';
 export {
   LLMModelsList,
   LLMProvidersList,
+  type LLMArgs,
   type LLMModel,
   type LLMProvider,
 } from './llm/types';
@@ -70,6 +71,7 @@ export class LitLytics {
   // model config
   provider?: LLMProviders;
   model?: LLMModel;
+  llmArgs?: LLMArgs;
   #llmKey?: string;
 
   // pipeline
@@ -82,14 +84,17 @@ export class LitLytics {
     provider,
     model,
     key,
+    llmArgs,
   }: {
     provider: LLMProviders;
     model: LLMModel;
     key: string;
+    llmArgs?: LLMArgs;
   }) {
     this.provider = provider;
     this.model = model;
     this.#llmKey = key;
+    this.llmArgs = llmArgs;
   }
 
   /**
@@ -179,7 +184,10 @@ export class LitLytics {
       key: this.#llmKey,
       model: this.model,
       messages,
-      args,
+      args: {
+        ...args,
+        ...this.llmArgs,
+      },
     });
   };
 
@@ -202,7 +210,10 @@ export class LitLytics {
       model: this.model,
       system,
       user,
-      args,
+      args: {
+        ...args,
+        ...this.llmArgs,
+      },
     });
   };
 

packages/litlytics/llm/types.ts

@@ -1,4 +1,4 @@
-import type { CoreMessage, CoreTool } from 'ai';
+import type { CoreMessage, generateText } from 'ai';
 
 export const LLMProvidersList = [
   'openai',
@@ -40,10 +40,12 @@ export const LLMModelsList = {
 export type LLMModel =
   (typeof LLMModelsList)[keyof typeof LLMModelsList][number];
 
+export type LLMArgs = Partial<Parameters<typeof generateText>[0]>;
+
 export interface LLMRequest {
   provider: LLMProvider;
   key: string;
   model: LLMModel;
   messages: CoreMessage[];
-  modelArgs?: Record<string, CoreTool>;
+  modelArgs?: LLMArgs;
 }

packages/litlytics/step/Step.ts

@@ -1,5 +1,6 @@
-import type { CoreTool, LanguageModelUsage } from 'ai';
+import type { LanguageModelUsage } from 'ai';
 import type { Doc } from '../doc/Document';
+import type { LLMArgs } from '../llm/types';
 
 export interface StepResult {
   stepId: string;
@@ -46,7 +47,7 @@ export interface ProcessingStep extends BaseStep {
   input?: StepInput;
   // llm
   prompt?: string;
-  llmArgs?: Record<string, CoreTool>;
+  llmArgs?: LLMArgs;
   // code
   code?: string;
   codeExplanation?: string;

packages/litlytics/test/litlytics.test.ts

@@ -1,9 +1,11 @@
 import type { LanguageModelUsage } from 'ai';
-import { expect, test } from 'vitest';
+import { expect, test, vi } from 'vitest';
+import * as run from '../engine/runPrompt';
 import {
   LitLytics,
   OUTPUT_ID,
   type Doc,
+  type LLMArgs,
   type Pipeline,
   type StepInput,
 } from '../litlytics';
@@ -229,3 +231,54 @@ test('should generate suggested tasks for current pipeline', async () => {
   const newNonTestDoc = litlytics.docs.find((d) => d.id === docNonTest.id);
   expect(newNonTestDoc?.summary).toBeUndefined();
 });
+
+test('should pass llm args when running prompt', async () => {
+  const testArgs: LLMArgs = {
+    temperature: 0.5,
+    maxTokens: 1000,
+  };
+  const litlytics = new LitLytics({
+    provider: 'openai',
+    model: 'test',
+    key: 'test',
+    llmArgs: testArgs,
+  });
+  litlytics.pipeline.pipelineDescription = 'test description';
+
+  const testResult = `Step name: Generate Title and Description
+Step type: llm
+Step input: doc
+Step description: Generate an Etsy product title and description based on the provided document describing the product.
+
+---
+
+Step name: Check for Copyrighted Terms
+Step type: llm
+Step input: result
+Step description: Analyze the generated title and description for possible copyrighted terms and suggest edits.
+`;
+
+  // mock prompt replies
+  const spy = vi
+    .spyOn(run, 'runPrompt')
+    .mockImplementation(
+      async ({
+        user,
+        args,
+      }: {
+        system: string;
+        user: string;
+        args?: LLMArgs;
+      }) => {
+        expect(args).toEqual(testArgs);
+        expect(user).toEqual('test description');
+        return { result: testResult, usage: {} as LanguageModelUsage };
+      }
+    );
+  // run generation
+  await litlytics.generatePipeline();
+  // check that spy was called
+  expect(spy).toHaveBeenCalled();
+  // cleanup
+  spy.mockClear();
+});