Clean up in response to PR feedback

Author: AlePouroullis

src/humanloop/cli/__main__.py

@@ -2,7 +2,7 @@
 import logging
 from typing import Optional, Callable
 from functools import wraps
-from dotenv import load_dotenv, find_dotenv
+from dotenv import load_dotenv
 import os
 import sys
 from humanloop import Humanloop
@@ -26,10 +26,9 @@
 
 
 def load_api_key(env_file: Optional[str] = None) -> str:
-    """Load API key from provided value, .env file, or environment variable.
+    """Load API key from .env file or environment variable.
 
     Args:
-        api_key: Optional API key provided directly
         env_file: Optional path to .env file
 
     Returns:
@@ -38,18 +37,17 @@ def load_api_key(env_file: Optional[str] = None) -> str:
     Raises:
         click.ClickException: If no API key is found
     """
-    # Try loading from .env file
+    # Try specific .env file if provided, otherwise default to .env in current directory
     if env_file:
-        load_dotenv(env_file)
+        if not load_dotenv(env_file):  # load_dotenv returns False if file not found/invalid
+            raise click.ClickException(
+                click.style(
+                    f"Failed to load environment file: {env_file} (file not found or invalid format)",
+                    fg=ERROR_COLOR,
+                )
+            )
     else:
-        # Try to find .env file in current directory or parent directories
-        env_path = find_dotenv()
-        if env_path:
-            load_dotenv(env_path)
-        elif os.path.exists(".env"):
-            load_dotenv(".env")
-        else:
-            load_dotenv()
+        load_dotenv()  # Attempt to load from default .env in current directory
 
     # Get API key from environment
     api_key = os.getenv("HUMANLOOP_API_KEY")
@@ -151,8 +149,9 @@ def cli():  # Does nothing because used as a group for other subcommands (pull,
 @click.option(
     "--path",
     "-p",
-    help="Path in the Humanloop workspace to pull from (file or directory). You can pull an entire directory (e.g. 'my/directory/') "
-    "or a specific file (e.g. 'my/directory/my_prompt.prompt'). When pulling a directory, all files within that directory and its subdirectories will be included.",
+    help="Path in the Humanloop workspace to pull from (file or directory). You can pull an entire directory (e.g. 'my/directory') "
+    "or a specific file (e.g. 'my/directory/my_prompt.prompt'). When pulling a directory, all files within that directory and its subdirectories will be included. "
+    "If not specified, pulls from the root of the workspace.",
     default=None,
 )
 @click.option(
@@ -190,7 +189,7 @@ def pull(
     \b
     This command will:
     1. Fetch Prompt and Agent files from your Humanloop workspace
-    2. Save them to your local filesystem (default: humanloop/)
+    2. Save them to your local filesystem (default directory: humanloop/)
     3. Maintain the same directory structure as in Humanloop
     4. Add appropriate file extensions (.prompt or .agent)
 

src/humanloop/overload.py

@@ -22,15 +22,20 @@
 from humanloop.types.create_prompt_log_response import CreatePromptLogResponse
 from humanloop.types.create_tool_log_response import CreateToolLogResponse
 from humanloop.types.prompt_call_response import PromptCallResponse
+from humanloop.types.agent_call_response import AgentCallResponse
 
 logger = logging.getLogger("humanloop.sdk")
 
-ResponseType = Union[
+LogResponseType = Union[
     CreatePromptLogResponse,
     CreateToolLogResponse,
     CreateFlowLogResponse,
     CreateEvaluatorLogResponse,
+]
+
+CallResponseType = Union[
     PromptCallResponse,
+    AgentCallResponse,
 ]
 
 
@@ -128,7 +133,7 @@ def _handle_evaluation_context(kwargs: Dict[str, Any]) -> tuple[Dict[str, Any],
     return kwargs, None
 
 
-def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> ResponseType:
+def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> LogResponseType:
     try:
         # Special handling for flows - prevent direct log usage
         if type(self) is FlowsClient and get_trace_id() is not None:
@@ -162,7 +167,7 @@ def _overload_log(self: Any, sync_client: Optional[SyncClient], use_local_files:
         raise HumanloopRuntimeError from e
 
 
-def _overload_call(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> PromptCallResponse:
+def _overload_call(self: Any, sync_client: Optional[SyncClient], use_local_files: bool, **kwargs) -> CallResponseType:
     try:
         kwargs = _handle_tracing_context(kwargs, self)
         kwargs = _handle_local_files(kwargs, self, sync_client, use_local_files)
@@ -186,7 +191,7 @@ def overload_client(
         client._log = client.log  # type: ignore [attr-defined]
 
         # Create a closure to capture sync_client and use_local_files
-        def log_wrapper(self: Any, **kwargs) -> ResponseType:
+        def log_wrapper(self: Any, **kwargs) -> LogResponseType:
             return _overload_log(self, sync_client, use_local_files, **kwargs)
 
         client.log = types.MethodType(log_wrapper, client)
@@ -200,7 +205,7 @@ def log_wrapper(self: Any, **kwargs) -> ResponseType:
             client._call = client.call  # type: ignore [attr-defined]
 
             # Create a closure to capture sync_client and use_local_files
-            def call_wrapper(self: Any, **kwargs) -> PromptCallResponse:
+            def call_wrapper(self: Any, **kwargs) -> CallResponseType:
                 return _overload_call(self, sync_client, use_local_files, **kwargs)
 
             client.call = types.MethodType(call_wrapper, client)

src/humanloop/sync/sync_client.py

@@ -33,8 +33,10 @@ def format_api_error(error: Exception) -> str:
     try:
         # Extract the body part and parse as JSON
         body_str = error_msg.split("body: ")[1]
-        # Convert Python dict string to valid JSON by replacing single quotes with double quotes
-        body_str = body_str.replace("'", '"')
+        # Convert Python dict string to valid JSON by:
+        # 1. Escaping double quotes
+        # 2. Replacing single quotes with double quotes
+        body_str = body_str.replace('"', '\\"').replace("'", '"')
         body = json.loads(body_str)
 
         # Get the detail from the body
@@ -52,7 +54,7 @@ def format_api_error(error: Exception) -> str:
         return error_msg
 
 
-_SERIALIZABLE_FILE_TYPES_TYPE = typing.Literal["prompt", "agent"]
+SerializableFileType = typing.Literal["prompt", "agent"]
 
 
 class SyncClient:
@@ -68,7 +70,7 @@ class SyncClient:
     """
 
     # File types that can be serialized to/from the filesystem
-    SERIALIZABLE_FILE_TYPES = ["prompt", "agent"]
+    SERIALIZABLE_FILE_TYPES = frozenset(typing.get_args(SerializableFileType))
 
     def __init__(
         self,
@@ -92,19 +94,18 @@ def __init__(
         logger.setLevel(log_level)
 
         # Create a new cached version of get_file_content with the specified cache size
-        # @TODO: @ale, maybe move the cache to the client?
         self.get_file_content = lru_cache(maxsize=cache_size)(  # type: ignore [assignment]
             self._get_file_content_implementation,
         )
 
-    def _get_file_content_implementation(self, path: str, file_type: FileType) -> str:
+    def _get_file_content_implementation(self, path: str, file_type: SerializableFileType) -> str:
         """Implementation of get_file_content without the cache decorator.
 
         This is the actual implementation that gets wrapped by lru_cache.
 
         Args:
             path: The normalized path to the file (without extension)
-            file_type: The type of file (Prompt or Agent)
+            file_type: The type of file to get the content of (SerializableFileType)
 
         Returns:
             The raw file content
@@ -132,7 +133,7 @@ def _get_file_content_implementation(self, path: str, file_type: FileType) -> st
         except Exception as e:
             raise HumanloopRuntimeError(f"Error reading local file {local_path}: {str(e)}")
 
-    def get_file_content(self, path: str, file_type: FileType) -> str:
+    def get_file_content(self, path: str, file_type: SerializableFileType) -> str:
         """Get the raw file content of a file from cache or filesystem.
 
         This method uses an LRU cache to store file contents. When the cache is full,
@@ -152,7 +153,6 @@ def get_file_content(self, path: str, file_type: FileType) -> str:
 
     def clear_cache(self) -> None:
         """Clear the LRU cache."""
-        # @TODO: @ale, why not put the cache on the client? This is a bit of a hack.
         self.get_file_content.cache_clear()  # type: ignore [attr-defined]
 
     def _normalize_path(self, path: str) -> str:
@@ -170,14 +170,14 @@ def _normalize_path(self, path: str) -> str:
         return "/".join(part for part in normalized.replace("\\", "/").split("/") if part)
 
     def is_file(self, path: str) -> bool:
-        """Check if the path is a file by checking for .prompt or .agent extension."""
-        return path.endswith(".prompt") or path.endswith(".agent")
+        """Check if the path is a file by checking for .{file_type} extension for serializable file types."""
+        return path.endswith(tuple(f".{file_type}" for file_type in self.SERIALIZABLE_FILE_TYPES))
 
     def _save_serialized_file(
         self,
         serialized_content: str,
         file_path: str,
-        file_type: typing.Literal["prompt", "agent"],
+        file_type: SerializableFileType,
     ) -> None:
         """Save serialized file to local filesystem."""
         try:
@@ -192,53 +192,39 @@ def _save_serialized_file(
             # Write raw file content to file
             with open(new_path, "w") as f:
                 f.write(serialized_content)
-
-            # Clear the cache when a file is saved
-            self.clear_cache()
         except Exception as e:
-            logger.error(f"Failed to sync {file_type} {file_path}: {str(e)}")
+            logger.error(f"Failed to write {file_type} {file_path} to disk: {str(e)}")
             raise
 
-    def _pull_file(
-        self,
-        path: str | None = None,
-        environment: str | None = None,
-    ) -> bool:
+    def _pull_file(self, path: str | None = None, environment: str | None = None) -> bool:
         """Pull a specific file from Humanloop to local filesystem.
 
         Returns:
             True if the file was successfully pulled, False otherwise
-
-        Raises:
-            HumanloopRuntimeError: If there's an error communicating with the API
         """
         try:
             file = self.client.files.retrieve_by_path(
                 path=path,
                 environment=environment,
                 include_raw_file_content=True,
             )
-        except Exception as e:
-            logger.error(f"Failed to pull file {path}: {format_api_error(e)}")
-            return False
+            
+            if file.type not in self.SERIALIZABLE_FILE_TYPES:
+                logger.error(f"Unsupported file type: {file.type}")
+                return False
 
-        if file.type not in self.SERIALIZABLE_FILE_TYPES:
-            raise ValueError(f"Unsupported file type: {file.type}")
+            if not file.raw_file_content:  # type: ignore [union-attr]
+                logger.error(f"No content found for {file.type} {path}")
+                return False
 
-        file_type: _SERIALIZABLE_FILE_TYPES_TYPE = typing.cast(
-            _SERIALIZABLE_FILE_TYPES_TYPE,
-            file.type,
-        )
-
-        try:
             self._save_serialized_file(
-                serialized_content=file.raw_file_content,  # type: ignore [union-attr, arg-type]
+                serialized_content=file.raw_file_content,  # type: ignore [union-attr]
                 file_path=file.path,
-                file_type=file_type,
+                file_type=typing.cast(SerializableFileType, file.type),
             )
             return True
         except Exception as e:
-            logger.error(f"Failed to save file {path}: {str(e)}")
+            logger.error(f"Failed to pull file {path}: {str(e)}")
             return False
 
     def _pull_directory(
@@ -251,7 +237,8 @@ def _pull_directory(
         Returns:
             Tuple of two lists:
             - First list contains paths of successfully synced files
-            - Second list contains paths of files that failed to sync
+            - Second list contains paths of files that failed to sync. 
+              Failures can occur due to missing content in the response or errors during local file writing.
 
         Raises:
             HumanloopRuntimeError: If there's an error communicating with the API
@@ -266,7 +253,7 @@ def _pull_directory(
             try:
                 logger.debug(f"`{path}`: Requesting page {page} of files")
                 response = self.client.files.list_files(
-                    type=["prompt", "agent"],
+                    type=list(self.SERIALIZABLE_FILE_TYPES),
                     page=page,
                     include_raw_file_content=True,
                     environment=environment,
@@ -281,20 +268,20 @@ def _pull_directory(
 
                 # Process each file
                 for file in response.records:
-                    # Skip if not a Prompt or Agent
+                    # Skip if not a serializable file type
                     if file.type not in self.SERIALIZABLE_FILE_TYPES:
                         logger.warning(f"Skipping unsupported file type: {file.type}")
                         continue
 
-                    file_type: _SERIALIZABLE_FILE_TYPES_TYPE = typing.cast(
-                        _SERIALIZABLE_FILE_TYPES_TYPE,
+                    file_type: SerializableFileType = typing.cast(
+                        SerializableFileType,
                         file.type,
                     )
 
                     # Skip if no raw file content
-                    # @TODO: @ale, inconsistent, other places we are throwing if file is not serialisable
                     if not getattr(file, "raw_file_content", None) or not file.raw_file_content:  # type: ignore [union-attr]
-                        logger.warning(f"No content found for {file.type} {getattr(file, 'id', '<unknown>')}")
+                        logger.warning(f"No content found for {file.type} {file.path}")
+                        failed_files.append(file.path)
                         continue
 
                     try:
@@ -368,6 +355,9 @@ def pull(self, path: str | None = None, environment: str | None = None) -> Tuple
                     logger.debug(f"Pulling directory: {normalized_path}")
                     successful_files, failed_files = self._pull_directory(normalized_path, environment)
 
+            # Clear the cache at the end of each pull operation
+            self.clear_cache()
+
             duration_ms = int((time.time() - start_time) * 1000)
             logger.info(f"Pull completed in {duration_ms}ms: {len(successful_files)} files succeeded")