dhisana 0.0.1.dev85__py3-none-any.whl → 0.0.1.dev236__py3-none-any.whl

dhisana/utils/built_with_api_tools.py

@@ -23,7 +23,7 @@ def get_builtwith_api_key(tool_config: Optional[List[Dict]] = None) -> str:
         str: The BUILTWITH_API_KEY access token.
 
     Raises:
-        ValueError: If the access token is not found in the tool configuration or environment variable.
+        ValueError: If the BuiltWith integration has not been configured.
     """
     if tool_config:
         builtwith_config = next(
@@ -43,7 +43,9 @@ def get_builtwith_api_key(tool_config: Optional[List[Dict]] = None) -> str:
 
     BUILTWITH_API_KEY = BUILTWITH_API_KEY or os.getenv("BUILTWITH_API_KEY")
     if not BUILTWITH_API_KEY:
-        raise ValueError("BUILTWITH_API_KEY access token not found in tool_config or environment variable")
+        raise ValueError(
+            "BuiltWith integration is not configured. Please configure the connection to BuiltWith in Integrations."
+        )
     return BUILTWITH_API_KEY
 
 # Use BuiltWith API to find tech stack and financials of a company

dhisana/utils/cache_output_tools.py

@@ -5,28 +5,31 @@ import logging
 
 from azure.storage.blob import BlobServiceClient
 from azure.core.exceptions import ResourceNotFoundError, AzureError
-logger = logging.getLogger(__name__)
 
-# Retrieve the connection string from the environment variable
-AZURE_BLOB_CONNECTION_STRING = os.environ.get("AZURE_BLOB_CONNECTION_STRING")
-if not AZURE_BLOB_CONNECTION_STRING:
-    raise ValueError("AZURE_BLOB_CONNECTION_STRING environment variable is not set")
+logger = logging.getLogger(__name__)
+logging.getLogger("azure").setLevel(logging.CRITICAL)
 
-# Define the container name for caching outputs
 CONTAINER_NAME = "cacheoutputs"
 
-# Initialize the BlobServiceClient
-blob_service_client = BlobServiceClient.from_connection_string(AZURE_BLOB_CONNECTION_STRING)
-logging.getLogger("azure").setLevel(logging.CRITICAL)
+def _get_container_client():
+    """
+    Returns the container client for the cache container.
+    Ensures that the container is created if it doesn't exist.
+    """
+    connection_string = os.environ.get("AZURE_BLOB_CONNECTION_STRING")
+    if not connection_string:
+        raise ValueError("AZURE_BLOB_CONNECTION_STRING environment variable is not set")
 
-# Ensure the container exists (if it already exists, an error is caught and ignored)
-try:
+    blob_service_client = BlobServiceClient.from_connection_string(connection_string)
     container_client = blob_service_client.get_container_client(CONTAINER_NAME)
-    container_client.create_container()
-except AzureError as e:
-    # The container might already exist so we ignore the exception
-    pass
 
+    # Ensure the container exists; if already created, the AzureError is ignored.
+    try:
+        container_client.create_container()
+    except AzureError:
+        pass
+
+    return container_client
 
 def cache_output(tool_name: str, key: str, value, ttl: int = None) -> bool:
     """
@@ -43,9 +46,7 @@ def cache_output(tool_name: str, key: str, value, ttl: int = None) -> bool:
     """
     # Create a hash of the key for a consistent blob name
     key_hash = hashlib.sha256(key.encode()).hexdigest()
-
-    # Construct the blob name using a virtual folder for the tool name.
-    # For example: "my_tool/abcdef123456.json"
+    # Construct the blob name using a virtual folder for the tool name
     blob_name = f"{tool_name}/{key_hash}.json"
 
     # Prepare the cache data
@@ -56,8 +57,8 @@ def cache_output(tool_name: str, key: str, value, ttl: int = None) -> bool:
     data = json.dumps(cache_data)
 
     try:
-        # Get a blob client for the specific blob
-        blob_client = blob_service_client.get_blob_client(container=CONTAINER_NAME, blob=blob_name)
+        container_client = _get_container_client()
+        blob_client = container_client.get_blob_client(blob=blob_name)
         # Upload the blob content (overwrite if the blob already exists)
         blob_client.upload_blob(data, overwrite=True)
         return True
@@ -65,7 +66,6 @@ def cache_output(tool_name: str, key: str, value, ttl: int = None) -> bool:
         logger.error(f"Error uploading blob '{blob_name}': {e}")
         return False
 
-
 def retrieve_output(tool_name: str, key: str):
     """
     Retrieve the cached output for a given tool and cache key from Azure Blob Storage.
@@ -79,12 +79,12 @@ def retrieve_output(tool_name: str, key: str):
     """
     # Create a hash of the key to locate the blob
     key_hash = hashlib.sha256(key.encode()).hexdigest()
-
     # Construct the blob name using the tool name folder
     blob_name = f"{tool_name}/{key_hash}.json"
 
     try:
-        blob_client = blob_service_client.get_blob_client(container=CONTAINER_NAME, blob=blob_name)
+        container_client = _get_container_client()
+        blob_client = container_client.get_blob_client(blob=blob_name)
         download_stream = blob_client.download_blob()
         content = download_stream.readall()  # content is in bytes
         cache_data = json.loads(content.decode("utf-8"))