diff --git a/promptlayer/utils.py b/promptlayer/utils.py index 4d0ddb0..6305e0d 100644 --- a/promptlayer/utils.py +++ b/promptlayer/utils.py @@ -9,6 +9,8 @@ import types from copy import deepcopy from enum import Enum +from typing import Any, AsyncGenerator, AsyncIterable, Callable, Dict, Generator, List, Optional, Union +from urllib.parse import quote from typing import Any, Dict, List, Optional, Union from uuid import uuid4 @@ -91,7 +93,7 @@ async def _resolve_workflow_id(workflow_id_or_name: Union[int, str], headers): # TODO(dmu) LOW: Should we warn user here to avoid using workflow names in favor of workflow id? async with _make_httpx_client() as client: # TODO(dmu) MEDIUM: Generalize the way we make async calls to PromptLayer API and reuse it everywhere - response = await client.get(f"{URL_API_PROMPTLAYER}/workflows/{workflow_id_or_name}", headers=headers) + response = await client.get(f"{URL_API_PROMPTLAYER}/workflows/{quote(str(workflow_id_or_name), safe='')}", headers=headers) if RAISE_FOR_STATUS: response.raise_for_status() elif response.status_code != 200: @@ -1085,7 +1087,7 @@ def get_prompt_template( if params: json_body = {**json_body, **params} response = requests.post( - f"{URL_API_PROMPTLAYER}/prompt-templates/{prompt_name}", + f"{URL_API_PROMPTLAYER}/prompt-templates/{quote(prompt_name, safe='')}", headers={"X-API-KEY": api_key}, json=json_body, ) @@ -1114,7 +1116,7 @@ async def aget_prompt_template( json_body.update(params) async with _make_httpx_client() as client: response = await client.post( - f"{URL_API_PROMPTLAYER}/prompt-templates/{prompt_name}", + f"{URL_API_PROMPTLAYER}/prompt-templates/{quote(prompt_name, safe='')}", headers={"X-API-KEY": api_key}, json=json_body, ) diff --git a/tests/test_prompt_name_encoding.py b/tests/test_prompt_name_encoding.py new file mode 100644 index 0000000..78851fd --- /dev/null +++ b/tests/test_prompt_name_encoding.py @@ -0,0 +1,35 @@ +from urllib.parse import quote + +def simulate_url_construction(prompt_name): + """Simulate how the URL is constructed in the get_prompt_template function.""" + URL_API_PROMPTLAYER = "https://api.promptlayer.com" + return f"{URL_API_PROMPTLAYER}/prompt-templates/{quote(prompt_name, safe='')}" + +def test_get_prompt_template_url_encoding(): + """Test that prompt_name with slashes is properly URL encoded in the request URL.""" + test_cases = [ + {"prompt_name": "feature1/resolve_problem_2", "expected_encoded": quote("feature1/resolve_problem_2", safe='')}, + {"prompt_name": "feature1:test", "expected_encoded": quote("feature1:test", safe='')}, + {"prompt_name": "feature1/sub/test?query", "expected_encoded": quote("feature1/sub/test?query", safe='')} + ] + + for test_case in test_cases: + prompt_name = test_case["prompt_name"] + expected_encoded = test_case["expected_encoded"] + + url = simulate_url_construction(prompt_name) + + assert expected_encoded in url, f"Expected {expected_encoded} in {url}" + +def test_resolve_workflow_id_url_encoding(): + """Test that workflow_id_or_name with slashes is properly URL encoded.""" + pass + +def main(): + """Run all verification tests.""" + test_get_prompt_template_url_encoding() + test_resolve_workflow_id_url_encoding() + print("All tests passed successfully!") + +if __name__ == "__main__": + main() \ No newline at end of file