Skip to content

Commit

Permalink
[python] Created get_parameters API (#668)
Browse files Browse the repository at this point in the history
[python] Created `get_parameters` API


**Disclaimer:** I know this is 400+ lines changed, but ~300 of those are
just adding/refactoring automated tests!!!

We already had `get_prompt_parameters`, but this isn't good enough if
parameters aren't defined locally and we want to bubble up to the
aiconfig defaults. See comment in
#661 (comment).

This is important so that we can ensure we have valid params if not
defined in the correct section of our AIConfig schema. This is similar
in theme to #600

The only callsite we have today for this is in python in the `params.py`
function, so I updated that callsite too. I don't know why it doesn't
exist in Typescript (cc @rholinshead) but I also added this
functionality in typescript in
#669 (I'm not as familiar
with automated testing there so will ask Ryan for a code pointer to help
get me started)

## Test plan
Added a bunch of automated tests which you can run by doing going to the
`aiconfig/python` dir and running `pytest`

---
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/lastmile-ai/aiconfig/pull/668).
* #693
* #692
* #691
* #690
* #688
* #670
* __->__ #668
  • Loading branch information
rossdanlm authored Jan 2, 2024
2 parents 526d45e + 5a75379 commit 6bb4f7c
Show file tree
Hide file tree
Showing 4 changed files with 407 additions and 99 deletions.
87 changes: 79 additions & 8 deletions python/src/aiconfig/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,85 @@ def get_metadata(self, prompt_name: Optional[str] = None):
else:
return self.metadata


def get_parameters(
self,
prompt_or_prompt_name: Optional[str | Prompt] = None,
) -> JSONObject:
"""
Get the parameters for a prompt, using the global parameters if
needed.
Args:
prompt_or_prompt_name Optional[str | Prompt]: The name of the
prompt or the prompt object. If not specified, use the
global parameters.
"""
prompt = prompt_or_prompt_name
if isinstance(prompt_or_prompt_name, str):
if prompt_or_prompt_name not in self.prompt_index:
raise IndexError(f"Prompt '{prompt_or_prompt_name}' not found in config, available prompts are:\n {list(self.prompt_index.keys())}")
prompt = self.prompt_index[prompt_or_prompt_name]

assert prompt is None or isinstance(prompt, Prompt)
if prompt is None or not prompt.metadata or not prompt.metadata.parameters:
return self.get_global_parameters()

return self.get_prompt_parameters(prompt)

# pylint: disable=W0102
def get_global_parameters(
self,
default_return_value: JSONObject = {},
) -> JSONObject:
"""
Get the global parameters for the AIConfig. If they're not defined,
return a default value ({} unless overridden)
Args:
default_return_value JSONObject - Default value to return if
global parameters are not defined.
"""
return self._get_global_parameters_exact() or default_return_value
# pylint: enable=W0102

def _get_global_parameters_exact(self) -> JSONObject | None:
"""
Get the global parameters for the AIConfig. This should be the
the explicit value (ie: if parameters is None, return None, not {})
"""
return self.metadata.parameters

# pylint: disable=W0102
def get_prompt_parameters(
self,
prompt: Prompt,
default_return_value: JSONObject = {},
) -> JSONObject:
"""
Get the prompt's local parameters. If they're not defined,
return a default value ({} unless overridden)
Args:
default_return_value JSONObject - Default value to return if
prompt parameters are not defined.
"""
return self._get_prompt_parameters_exact(prompt) \
or default_return_value
# pylint: enable=W0102

def _get_prompt_parameters_exact(
self,
prompt: Prompt,
) -> JSONObject | None:
"""
Get the global parameters for the AIConfig. This should be the
the explicit value (ie: if parameters is None, return None, not {})
"""
if not prompt.metadata:
return prompt.metadata
return prompt.metadata.parameters

def set_parameter(self, parameter_name: str, parameter_value, prompt_name: Optional[str] = None):
"""
Sets a parameter in the AI configuration metadata. If a prompt_name is specified, it adds the parameter to
Expand Down Expand Up @@ -733,14 +812,6 @@ def get_output_text(self, prompt: str | Prompt):
prompt (str|Prompt): The name of the prompt or the prompt object.
"""

def get_prompt_parameters(self, prompt: Prompt):
"""
Gets the prompt's local parameters for a prompt.
"""
if not prompt.metadata:
return {}
return prompt.metadata.parameters

"""
Library Helpers
"""
Expand Down
2 changes: 1 addition & 1 deletion python/src/aiconfig/util/params.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ def resolve_prompt_string(
augmented_params = collect_prompt_references(current_prompt, ai_config)

# augment params with config-level params
augmented_params.update(ai_config.metadata.parameters)
augmented_params.update(ai_config.get_global_parameters())

# augment params with prompt level params
augmented_params.update(ai_config.get_prompt_parameters(current_prompt))
Expand Down
Loading

0 comments on commit 6bb4f7c

Please sign in to comment.