Project import generated by Copybara.

GitOrigin-RevId: 589948059c626b0f68f5633ecc36505c90c7dea7

src/gretel_client/navigator/__init__.py

@@ -1,2 +1,2 @@
 from gretel_client.navigator.data_designer.interface import DataDesigner
-from gretel_client.navigator.workflow import NavigatorWorkflow
+from gretel_client.navigator.workflow import DataDesignerWorkflow

src/gretel_client/navigator/client/remote.py

@@ -219,7 +219,6 @@ def submit_batch_workflow(
         logger.info(f"▶️ Starting your workflow run to generate {num_records} records:")
         logger.info(f"  |-- project_name: {project.name}")
         logger.info(f"  |-- project_id: {project.project_guid}")
-        logger.info(f"  |-- workflow_id: {batch_response.workflow_id}")
         logger.info(f"  |-- workflow_run_id: {batch_response.workflow_run_id}")
         logger.info(f"🔗 -> {workflow_run_url}")
 

src/gretel_client/navigator/data_designer/data_column.py

@@ -1,17 +1,19 @@
+from enum import Enum
 from typing import Optional
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 from gretel_client.navigator.client.interface import Client
 from gretel_client.navigator.data_designer.prompt_templates import (
     COLUMN_GENERATION_PROMPT,
+    get_prompt_template_keywords,
     system_prompt_dict,
 )
 from gretel_client.navigator.tasks.generate.generate_column_from_template import (
     GenerateColumnFromTemplate,
     TextParserType,
 )
-from gretel_client.navigator.tasks.types import LLMType, OutputColumnType
+from gretel_client.navigator.tasks.types import LLMType, OutputColumnType, RecordsT
 
 parser_instructions_map = {
     TextParserType.PASS_THROUGH: (
@@ -41,6 +43,20 @@
 }
 
 
+class PromptControlLevel(str, Enum):
+    ASSISTED = "assisted"
+    FULL_USER_CONTROL = "full_user_control"
+
+
+class UpsampledDataColumn(BaseModel):
+    name: str
+
+
+class UpsampledData(BaseModel):
+    records: RecordsT
+    columns: list[UpsampledDataColumn]
+
+
 class GeneratedDataColumn(BaseModel):
     """Generated data column to be used with DataDesigner.
 
@@ -54,13 +70,38 @@ class GeneratedDataColumn(BaseModel):
         llm_type: LLM type for generation of the column. Must be one of ["nl", "code", "judge"].
         output_type: Output type for the column. Must be one of ["auto", "text", "dict", "list", "code"].
             If "auto", the output type will be "code" when llm_type is "code" and "text" otherwise.
+        prompt_control_level: The level of control the user has over the prompt. If "assisted", the
+            `generation_prompt` will be injected into a pre-defined template that has been designed to
+            work well for a wide range of use cases, and the system prompt is set automatically.
+            This is the recommended setting for most users. If "full_user_control", `generation_prompt`
+            will be the exact prompt that is sent to the LLM (if `columns_to_list_in_prompt` is set, the specified
+            columns will also be listed in the prompt as "relevant context"), and the user now has full control
+            of the system prompt via the `system_prompt` parameter, which can use `special_system_instructions`
+            as a template keyword.
+        system_prompt:  The system prompt to be used for the LLM. This is only used when `prompt_control_level`
+            is set to "full_user_control". Otherwise, this parameter is ignored.
     """
 
     name: str
     generation_prompt: str
     columns_to_list_in_prompt: list[str] = Field(default_factory=list)
     llm_type: LLMType = LLMType.NATURAL_LANGUAGE
     output_type: OutputColumnType = OutputColumnType.AUTO
+    prompt_control_level: PromptControlLevel = PromptControlLevel.ASSISTED
+    system_prompt: Optional[str] = None
+
+    @field_validator("system_prompt")
+    def system_prompt_validator(cls, v, values):
+        if v is None:
+            return v
+        if (
+            values.data.get("prompt_control_level")
+            == PromptControlLevel.FULL_USER_CONTROL
+        ):
+            return v
+        raise ValueError(
+            "system_prompt can only be set when prompt_control_level is 'full_user_control'"
+        )
 
     @property
     def _resolved_output_type(self) -> OutputColumnType:
@@ -83,6 +124,8 @@ def prompt_template(self) -> str:
         Returns:
             Full prompt template string.
         """
+        if self.prompt_control_level == PromptControlLevel.FULL_USER_CONTROL:
+            return self.generation_prompt
         return COLUMN_GENERATION_PROMPT.format(
             name=self.name,
             generation_prompt=self.generation_prompt,
@@ -117,6 +160,35 @@ def generate_context_column_string(self, exclude: Optional[set[str]] = None) ->
             + "\n"
         )
 
+    def get_system_prompt(
+        self, special_system_instructions: Optional[str] = None
+    ) -> str:
+        """Get the system prompt for the column generation task.
+
+        Args:
+            special_instructions: Special instructions to be added to the system prompt.
+
+        Returns:
+            System prompt string.
+        """
+        if self.prompt_control_level == PromptControlLevel.FULL_USER_CONTROL:
+            system_prompt = self.system_prompt or ""
+            if "special_system_instructions" in get_prompt_template_keywords(
+                system_prompt
+            ):
+                system_prompt = system_prompt.format(
+                    special_system_instructions=special_system_instructions
+                )
+            return system_prompt
+
+        return system_prompt_dict[self.llm_type].format(
+            special_instructions=(
+                ""
+                if special_system_instructions is None
+                else f"\n{special_system_instructions}\n"
+            )
+        )
+
     def to_generation_task(
         self,
         special_system_instructions: Optional[str] = None,
@@ -137,13 +209,7 @@ def to_generation_task(
             response_column_name=self.name,
             workflow_label=f"generating {self.name}",
             llm_type=self.llm_type,
-            system_prompt=system_prompt_dict[self.llm_type].format(
-                special_instructions=(
-                    ""
-                    if special_system_instructions is None
-                    else f"\n{special_system_instructions}\n"
-                )
-            ),
+            system_prompt=self.get_system_prompt(special_system_instructions),
             client=client,
         )