feat: Add support for setting environment during prompt creation

js/src/cli/functions/upload.ts

@@ -45,6 +45,7 @@ interface BundledFunctionSpec {
   function_schema?: FunctionObject["function_schema"];
   if_exists?: IfExists;
   metadata?: Record<string, unknown>;
+  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) => ({

js/src/framework-types.ts

@@ -21,6 +21,7 @@ interface BaseFnOpts {
   description: string;
   ifExists: IfExists;
   metadata?: Record<string, unknown>;
+  environment?: string;
 }
 
 export type ToolOpts<

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", () => {

js/src/framework2.ts

@@ -36,6 +36,7 @@ interface BaseFnOpts {
   description: string;
   ifExists: IfExists;
   metadata?: Record<string, unknown>;
+  environment?: string;
 }
 
 export { toolFunctionDefinitionSchema };
@@ -405,6 +406,7 @@ export class CodePrompt {
   public readonly functionType?: FunctionType;
   public readonly toolFunctions: (SavedFunctionId | GenericCodeFunction)[];
   public readonly metadata?: Record<string, unknown>;
+  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<string, unknown>;
+  environment?: string;
 }
 
 export class ProjectNameIdMap {

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