sync load_prompt and aload_prompt

Author: Mahhheshh

py/src/braintrust/logger.py

@@ -786,7 +786,7 @@ def post_json(self, object_type: str, args: Mapping[str, Any] | None = None) ->
 
     async def aget_json(
         self, object_type: str, args: Optional[Mapping[str, Any]] = None, retries: int = 0
-    ) -> Mapping[str, Any]:
+    ) -> Mapping[str, Any] | None:
         """
         Async version of get_json. Makes a true async HTTP GET request and returns JSON response.
         """
@@ -833,15 +833,16 @@ async def _make_aiohttp_request(self, url: str) -> Mapping[str, Any]:
 
     async def _make_asyncio_request(self, url: str) -> Mapping[str, Any]:
         """Make async HTTP request using asyncio and urllib (fallback)"""
-        loop = asyncio.get_event_loop()
+        loop = asyncio.get_running_loop()
+        timeout_secs = parse_env_var_float("BRAINTRUST_HTTP_TIMEOUT", 60.0)
 
         def sync_request():
             request = Request(url)
             if self.token:
                 request.add_header("Authorization", f"Bearer {self.token}")
 
             try:
-                response_obj = urlopen(request)
+                response_obj = urlopen(request, timeout=timeout_secs)
                 response_data = response_obj.read()
                 return json.loads(response_data.decode("utf-8"))
             except HTTPError as e:
@@ -2079,8 +2080,10 @@ async def aload_prompt(
     slug: Optional[str] = None,
     version: Optional[Union[str, int]] = None,
     project_id: Optional[str] = None,
+    prompt_id: str | None = None,
     defaults: Optional[Mapping[str, Any]] = None,
     no_trace: bool = False,
+    environment: str | None = None,
     app_url: Optional[str] = None,
     api_key: Optional[str] = None,
     org_name: Optional[str] = None,
@@ -2092,81 +2095,134 @@ async def aload_prompt(
     :param slug: The slug of the prompt to load.
     :param version: An optional version of the prompt (to read). If not specified, the latest version will be used.
     :param project_id: The id of the project to load the prompt from. This takes precedence over `project` if specified.
+    :param prompt_id: The id of a specific prompt to load. If specified, this takes precedence over all other parameters (project, slug, version).
     :param defaults: (Optional) A dictionary of default values to use when rendering the prompt. Prompt values will override these defaults.
     :param no_trace: If true, do not include logging metadata for this prompt when build() is called.
+    :param environment: The environment to load the prompt from. Cannot be used together with version.
     :param app_url: The URL of the Braintrust App. Defaults to https://www.braintrust.dev.
     :param api_key: The API key to use. If the parameter is not specified, will try to use the `BRAINTRUST_API_KEY` environment variable. If no API
     key is specified, will prompt the user to login.
     :param org_name: (Optional) The name of a specific organization to connect to. This is useful if you belong to multiple.
     :returns: The prompt object.
     """
 
-    if not project and not project_id:
+    if version is not None and environment is not None:
+        raise ValueError(
+            "Cannot specify both 'version' and 'environment' parameters. Please use only one (remove the other)."
+        )
+
+    if prompt_id:
+        pass
+    elif not project and not project_id:
         raise ValueError("Must specify at least one of project or project_id")
-    if not slug:
+    elif not slug:
         raise ValueError("Must specify slug")
 
-    loop = asyncio.get_event_loop()
+    loop = asyncio.get_running_loop()
+    response = None
 
     try:
         # Run login in thread pool since it's synchronous
         await loop.run_in_executor(HTTP_REQUEST_THREAD_POOL, login, app_url, api_key, org_name)
-
-        # Make async HTTP request
-        args = _populate_args(
-            {
-                "project_name": project,
-                "project_id": project_id,
-                "slug": slug,
+        if prompt_id:
+            args = _populate_args({
                 "version": version,
-            },
-        )
+                "environment": environment
+            })
+
+            response = await _state.api_conn().aget_json(f"/v1/prompt/{prompt_id}", args)
+
+            if response:
+                response = {"objects": [response]}
+
+        else:
+            args = _populate_args(
+                {
+                    "project_name": project,
+                    "project_id": project_id,
+                    "slug": slug,
+                    "version": version,
+                    "environment": environment
+                },
+            )
 
-        response = await _state.api_conn().aget_json("/v1/prompt", args)
+            response = await _state.api_conn().aget_json("/v1/prompt", args)
 
     except Exception as server_error:
+        # If environment was specified, don't fall back to cache
+        if environment is not None:
+            raise ValueError(f"Prompt not found for specified environment {environment}") from server_error
+
         eprint(f"Failed to load prompt, attempting to fall back to cache: {server_error}")
         try:
-            cache_result = await loop.run_in_executor(
-                HTTP_REQUEST_THREAD_POOL,
-                lambda: _state._prompt_cache.get(
-                    slug,
-                    version=str(version) if version else "latest",
-                    project_id=project_id,
-                    project_name=project,
-                ),
-            )
+            if prompt_id:
+                cache_result = await loop.run_in_executor(
+                    HTTP_REQUEST_THREAD_POOL,
+                    lambda: _state._prompt_cache.get(
+                        id=prompt_id
+                    ),
+                )
+            else:
+                cache_result = await loop.run_in_executor(
+                    HTTP_REQUEST_THREAD_POOL,
+                    lambda: _state._prompt_cache.get(
+                        slug,
+                        version=str(version) if version else "latest",
+                        project_id=project_id,
+                        project_name=project,
+                    ),
+                )
             # Return Prompt with pre-computed metadata from cache
             return Prompt(
                 lazy_metadata=LazyValue(lambda: cache_result, use_mutex=True),
                 defaults=defaults or {},
                 no_trace=no_trace,
             )
         except Exception as cache_error:
+            if prompt_id:
+                raise ValueError(
+                    f"Prompt with id {prompt_id} not found (not found on server or in local cache): {cache_error}"
+                ) from server_error
             raise ValueError(
                 f"Prompt {slug} (version {version or 'latest'}) not found in {project or project_id} (not found on server or in local cache): {cache_error}"
             ) from server_error
 
     if response is None or "objects" not in response or len(response["objects"]) == 0:
+        if prompt_id:
+            raise ValueError(f"Prompt with id {prompt_id} not found.")
+
         raise ValueError(f"Prompt {slug} not found in project {project or project_id}.")
     elif len(response["objects"]) > 1:
+        if prompt_id:
+            raise ValueError(f"Multiple prompts found with id {prompt_id}. This should never happen.")
+
         raise ValueError(
             f"Multiple prompts found with slug {slug} in project {project or project_id}. This should never happen."
         )
 
     resp_prompt = response["objects"][0]
     prompt_metadata = PromptSchema.from_dict_deep(resp_prompt)
     try:
-        await loop.run_in_executor(
-            HTTP_REQUEST_THREAD_POOL,
-            lambda: _state._prompt_cache.set(
-                slug,
-                str(version) if version else "latest",
-                prompt_metadata,
-                project_id=project_id,
-                project_name=project,
-            ),
-        )
+        # save prompt to cache
+        if prompt_id:
+            await loop.run_in_executor(
+                HTTP_REQUEST_THREAD_POOL,
+                lambda: _state._prompt_cache.set(
+                    prompt_metadata,
+                    id=prompt_id
+                ),
+            )
+        else:
+            await loop.run_in_executor(
+                HTTP_REQUEST_THREAD_POOL,
+                lambda: _state._prompt_cache.set(
+                    prompt_metadata,
+                    slug=slug,
+                    version=str(version) if version else "latest",
+                    project_id=project_id,
+                    project_name=project,
+                ),
+            )
     except Exception as e:
         eprint(f"Failed to store prompt in cache: {e}")