diff --git a/js/src/cli/functions/upload.ts b/js/src/cli/functions/upload.ts index 891a2e66c..73e143b02 100644 --- a/js/src/cli/functions/upload.ts +++ b/js/src/cli/functions/upload.ts @@ -45,6 +45,7 @@ interface BundledFunctionSpec { function_schema?: FunctionObject["function_schema"]; if_exists?: IfExists; metadata?: Record; + environment?: string; } const pathInfoSchema = z @@ -114,6 +115,7 @@ export async function uploadHandleBundles({ : undefined, if_exists: fn.ifExists, metadata: fn.metadata, + environment: fn.environment, }); } @@ -363,6 +365,7 @@ async function uploadBundles({ function_schema: spec.function_schema, if_exists: spec.if_exists, metadata: spec.metadata, + environment: spec.environment, })), )) as FunctionEvent[]), ].map((fn) => ({ diff --git a/js/src/framework-types.ts b/js/src/framework-types.ts index f29f3b586..50708938c 100644 --- a/js/src/framework-types.ts +++ b/js/src/framework-types.ts @@ -21,6 +21,7 @@ interface BaseFnOpts { description: string; ifExists: IfExists; metadata?: Record; + environment?: string; } export type ToolOpts< diff --git a/js/src/framework.test.ts b/js/src/framework.test.ts index cbc5206da..414dfcc2a 100644 --- a/js/src/framework.test.ts +++ b/js/src/framework.test.ts @@ -880,6 +880,39 @@ describe("framework2 metadata support", () => { expect(prompts[0].metadata).toBeUndefined(); }); + test("prompt stores environment correctly", () => { + const project = projects.create({ name: "test-project" }); + const environment = "production"; + + project.prompts.create({ + name: "test-prompt", + prompt: "Hello {{name}}", + model: "gpt-4", + environment, + }); + + // The environment is stored on the CodePrompt in _publishablePrompts + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const prompts = (project as any)._publishablePrompts; + expect(prompts).toHaveLength(1); + expect(prompts[0].environment).toEqual(environment); + }); + + test("prompt works without environment", () => { + const project = projects.create({ name: "test-project" }); + + project.prompts.create({ + name: "test-prompt", + prompt: "Hello {{name}}", + model: "gpt-4", + }); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const prompts = (project as any)._publishablePrompts; + expect(prompts).toHaveLength(1); + expect(prompts[0].environment).toBeUndefined(); + }); + test("toFunctionDefinition includes metadata when present", async () => { const project = projects.create({ name: "test-project" }); const metadata = { version: "2.0", tag: "production" }; @@ -935,6 +968,62 @@ describe("framework2 metadata support", () => { expect(funcDef.metadata).toBeUndefined(); }); + + test("toFunctionDefinition includes environment when present", async () => { + const project = projects.create({ name: "test-project" }); + const environment = "staging"; + + const codePrompt = new CodePrompt( + project, + { + prompt: { type: "completion", content: "Hello {{name}}" }, + options: { model: "gpt-4" }, + }, + [], + { + name: "test-prompt", + slug: "test-prompt", + environment, + }, + ); + + const mockProjectMap = { + resolve: vi.fn().mockResolvedValue("project-123"), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any; + + const funcDef = await codePrompt.toFunctionDefinition(mockProjectMap); + + expect(funcDef.environment).toEqual(environment); + expect(funcDef.name).toBe("test-prompt"); + expect(funcDef.project_id).toBe("project-123"); + }); + + test("toFunctionDefinition excludes environment when undefined", async () => { + const project = projects.create({ name: "test-project" }); + + const codePrompt = new CodePrompt( + project, + { + prompt: { type: "completion", content: "Hello {{name}}" }, + options: { model: "gpt-4" }, + }, + [], + { + name: "test-prompt", + slug: "test-prompt", + }, + ); + + const mockProjectMap = { + resolve: vi.fn().mockResolvedValue("project-123"), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any; + + const funcDef = await codePrompt.toFunctionDefinition(mockProjectMap); + + expect(funcDef.environment).toBeUndefined(); + }); }); describe("Scorer metadata", () => { diff --git a/js/src/framework2.ts b/js/src/framework2.ts index 096f5284c..ce70d445b 100644 --- a/js/src/framework2.ts +++ b/js/src/framework2.ts @@ -36,6 +36,7 @@ interface BaseFnOpts { description: string; ifExists: IfExists; metadata?: Record; + environment?: string; } export { toolFunctionDefinitionSchema }; @@ -405,6 +406,7 @@ export class CodePrompt { public readonly functionType?: FunctionType; public readonly toolFunctions: (SavedFunctionId | GenericCodeFunction)[]; public readonly metadata?: Record; + public readonly environment?: string; constructor( project: Project, @@ -426,6 +428,7 @@ export class CodePrompt { this.id = opts.id; this.functionType = functionType; this.metadata = opts.metadata; + this.environment = opts.environment; } async toFunctionDefinition( @@ -467,6 +470,7 @@ export class CodePrompt { prompt_data, if_exists: this.ifExists, metadata: this.metadata, + environment: this.environment, }; } } @@ -624,6 +628,7 @@ export interface FunctionEvent { function_type?: FunctionType; if_exists?: IfExists; metadata?: Record; + environment?: string; } export class ProjectNameIdMap { diff --git a/py/src/braintrust/framework2.py b/py/src/braintrust/framework2.py index 8ea9d37c4..685d3b292 100644 --- a/py/src/braintrust/framework2.py +++ b/py/src/braintrust/framework2.py @@ -68,6 +68,7 @@ class CodePrompt: id: Optional[str] if_exists: Optional[IfExists] metadata: Optional[Dict[str, Any]] = None + environment: Optional[str] = None def to_function_definition(self, if_exists: Optional[IfExists], project_ids: ProjectIdCache) -> Dict[str, Any]: prompt_data = self.prompt @@ -101,6 +102,8 @@ def to_function_definition(self, if_exists: Optional[IfExists], project_ids: Pro j["function_type"] = self.function_type if self.metadata is not None: j["metadata"] = self.metadata + if self.environment is not None: + j["environment"] = self.environment return j @@ -179,13 +182,14 @@ def create( slug: Optional[str] = None, description: Optional[str] = None, id: Optional[str] = None, - prompt: str, - model: str, - params: Optional[ModelParams] = None, - tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None, - if_exists: Optional[IfExists] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> CodePrompt: ... + prompt: str, + model: str, + params: Optional[ModelParams] = None, + tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None, + if_exists: Optional[IfExists] = None, + metadata: Optional[Dict[str, Any]] = None, + environment: Optional[str] = None, + ) -> CodePrompt: ... @overload # messages only, no prompt def create( @@ -201,6 +205,7 @@ def create( tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None, if_exists: Optional[IfExists] = None, metadata: Optional[Dict[str, Any]] = None, + environment: Optional[str] = None, ) -> CodePrompt: ... def create( @@ -217,6 +222,7 @@ def create( tools: Optional[List[Union[CodeFunction, SavedFunctionId, ToolFunctionDefinition]]] = None, if_exists: Optional[IfExists] = None, metadata: Optional[Dict[str, Any]] = None, + environment: Optional[str] = None, ): """Creates a prompt. @@ -232,6 +238,7 @@ def create( tools: The tools to use for the prompt. if_exists: What to do if the prompt already exists. metadata: Custom metadata to attach to the prompt. + environment: The environment to assign the prompt to. """ self._task_counter += 1 if not name: @@ -281,6 +288,7 @@ def create( id=id, if_exists=if_exists, metadata=metadata, + environment=environment, ) self.project.add_prompt(p) return p