ragaai-catalyst 2.2.1__py3-none-any.whl → 2.2.2__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py

@@ -1,25 +1,28 @@
-import requests
 import json
+import logging
 import os
+import re
 import time
-import logging
-from datetime import datetime
 from urllib.parse import urlparse, urlunparse
-import re
+
+import requests
 
 logger = logging.getLogger(__name__)
 
+from ragaai_catalyst.ragaai_catalyst import RagaAICatalyst
+
 
 class UploadAgenticTraces:
-    def __init__(self, 
-                 json_file_path,
-                 project_name,
-                 project_id,
-                 dataset_name,
-                 user_detail,
-                 base_url,
-                 timeout=120,
-                 ):
+    def __init__(
+        self,
+        json_file_path,
+        project_name,
+        project_id,
+        dataset_name,
+        user_detail,
+        base_url,
+        timeout=120,
+    ):
         self.json_file_path = json_file_path
         self.project_name = project_name
         self.project_id = project_id
@@ -28,12 +31,13 @@ class UploadAgenticTraces:
         self.base_url = base_url
         self.timeout = timeout
 
-
     def _get_presigned_url(self):
-        payload = json.dumps({
+        payload = json.dumps(
+            {
                 "datasetName": self.dataset_name,
                 "numFiles": 1,
-            })
+            }
+        )
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
@@ -44,50 +48,82 @@ class UploadAgenticTraces:
             start_time = time.time()
             endpoint = f"{self.base_url}/v1/llm/presigned-url"
             # Changed to POST from GET
-            response = requests.request("POST", 
-                                        endpoint, 
-                                        headers=headers, 
-                                        data=payload,
-                                        timeout=self.timeout)
+            response = requests.request(
+                "POST", endpoint, headers=headers, data=payload, timeout=self.timeout
+            )
             elapsed_ms = (time.time() - start_time) * 1000
             logger.debug(
-                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
-            
+                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
+
             if response.status_code == 200:
                 presignedURLs = response.json()["data"]["presignedUrls"][0]
-                presignedurl = self.update_presigned_url(presignedURLs,self.base_url)
+                presignedurl = self.update_presigned_url(presignedURLs, self.base_url)
                 return presignedurl
             else:
                 # If POST fails, try GET
-                response = requests.request("GET", 
-                                        endpoint, 
-                                        headers=headers, 
-                                        data=payload,
-                                        timeout=self.timeout)
+                response = requests.request(
+                    "GET", endpoint, headers=headers, data=payload, timeout=self.timeout
+                )
                 elapsed_ms = (time.time() - start_time) * 1000
                 logger.debug(
-                    f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+                    f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+                )
                 if response.status_code == 200:
                     presignedURLs = response.json()["data"]["presignedUrls"][0]
-                    presignedurl = self.update_presigned_url(presignedURLs,self.base_url)
+                    presignedurl = self.update_presigned_url(
+                        presignedURLs, self.base_url
+                    )
                     return presignedurl
+                elif response.status_code == 401:
+                    logger.warning("Received 401 error. Attempting to refresh token.")
+                    token = RagaAICatalyst.get_token(force_refresh=True)
+                    headers = {
+                        "Content-Type": "application/json",
+                        "Authorization": f"Bearer {token}",
+                        "X-Project-Name": self.project_name,
+                    }
+                    response = requests.request(
+                        "POST",
+                        endpoint,
+                        headers=headers,
+                        data=payload,
+                        timeout=self.timeout,
+                    )
+                    elapsed_ms = (time.time() - start_time) * 1000
+                    logger.debug(
+                        f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+                    )
+                    if response.status_code == 200:
+                        presignedURLs = response.json()["data"]["presignedUrls"][0]
+                        presignedurl = self.update_presigned_url(
+                            presignedURLs, self.base_url
+                        )
+                        return presignedurl
+                    else:
+                        logger.error(
+                            f"Error while getting presigned url: {response.json()['message']}"
+                        )
+                        return None
+                else:
+                    logger.error(
+                        f"Error while getting presigned url: {response.json()['message']}"
+                    )
+                    return None
 
-                logger.error(f"Error while getting presigned url: {response.json()['message']}")
-                return None
-            
         except requests.exceptions.RequestException as e:
             logger.error(f"Error while getting presigned url: {e}")
             return None
-        
+
     def update_presigned_url(self, presigned_url, base_url):
-        """Replaces the domain (and port, if applicable) of the presigned URL 
+        """Replaces the domain (and port, if applicable) of the presigned URL
         with that of the base URL only if the base URL contains 'localhost' or an IP address."""
-        #To Do: If Proxy URL has domain name how do we handle such cases
-        
+        # To Do: If Proxy URL has domain name how do we handle such cases
+
         presigned_parts = urlparse(presigned_url)
         base_parts = urlparse(base_url)
         # Check if base_url contains localhost or an IP address
-        if re.match(r'^(localhost|\d{1,3}(\.\d{1,3}){3})$', base_parts.hostname):
+        if re.match(r"^(localhost|\d{1,3}(\.\d{1,3}){3})$", base_parts.hostname):
             new_netloc = base_parts.hostname  # Extract domain from base_url
             if base_parts.port:  # Add port if present in base_url
                 new_netloc += f":{base_parts.port}"
@@ -97,12 +133,12 @@ class UploadAgenticTraces:
 
     def _put_presigned_url(self, presignedUrl, filename):
         headers = {
-                "Content-Type": "application/json",
-            }
+            "Content-Type": "application/json",
+        }
 
         if "blob.core.windows.net" in presignedUrl:  # Azure
             headers["x-ms-blob-type"] = "BlockBlob"
-        print(f"Uploading agentic traces...")
+        print("Uploading agentic traces...")
         try:
             with open(filename) as f:
                 payload = f.read().replace("\n", "").replace("\r", "").encode()
@@ -111,14 +147,13 @@ class UploadAgenticTraces:
             return None
         try:
             start_time = time.time()
-            response = requests.request("PUT", 
-                                        presignedUrl, 
-                                        headers=headers, 
-                                        data=payload,
-                                        timeout=self.timeout)
+            response = requests.request(
+                "PUT", presignedUrl, headers=headers, data=payload, timeout=self.timeout
+            )
             elapsed_ms = (time.time() - start_time) * 1000
             logger.debug(
-                f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+                f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
             if response.status_code != 200 or response.status_code != 201:
                 return response, response.status_code
         except requests.exceptions.RequestException as e:
@@ -127,29 +162,55 @@ class UploadAgenticTraces:
 
     def insert_traces(self, presignedUrl):
         headers = {
-                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
-                "Content-Type": "application/json",
-                "X-Project-Name": self.project_name,
-            }
-        payload = json.dumps({
+            "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+            "Content-Type": "application/json",
+            "X-Project-Name": self.project_name,
+        }
+        payload = json.dumps(
+            {
                 "datasetName": self.dataset_name,
                 "presignedUrl": presignedUrl,
-                "datasetSpans": self._get_dataset_spans(), #Extra key for agentic traces
-            })
+                "datasetSpans": self._get_dataset_spans(),  # Extra key for agentic traces
+            }
+        )
         try:
             start_time = time.time()
             endpoint = f"{self.base_url}/v1/llm/insert/trace"
-            response = requests.request("POST", 
-                                        endpoint, 
-                                        headers=headers, 
-                                        data=payload,
-                                        timeout=self.timeout)
+            response = requests.request(
+                "POST", endpoint, headers=headers, data=payload, timeout=self.timeout
+            )
             elapsed_ms = (time.time() - start_time) * 1000
             logger.debug(
-                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
             if response.status_code != 200:
                 print(f"Error inserting traces: {response.json()['message']}")
                 return None
+            elif response.status_code == 401:
+                logger.warning("Received 401 error. Attempting to refresh token.")
+                token = RagaAICatalyst.get_token(force_refresh=True)
+                headers = {
+                    "Authorization": f"Bearer {token}",
+                    "Content-Type": "application/json",
+                    "X-Project-Name": self.project_name,
+                }
+                response = requests.request(
+                    "POST",
+                    endpoint,
+                    headers=headers,
+                    data=payload,
+                    timeout=self.timeout,
+                )
+                elapsed_ms = (time.time() - start_time) * 1000
+                logger.debug(
+                    f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+                )
+                if response.status_code != 200:
+                    print(f"Error inserting traces: {response.json()['message']}")
+                    return None
+                else:
+                    print("Error while inserting traces")
+                    return None
         except requests.exceptions.RequestException as e:
             print(f"Error while inserting traces: {e}")
             return None
@@ -166,12 +227,16 @@ class UploadAgenticTraces:
             dataset_spans = []
             for span in spans:
                 try:
-                    dataset_spans.append({
-                        "spanId": span.get("context", {}).get("span_id", ""),
-                        "spanName": span.get("name", ""),
-                        "spanHash": span.get("hash_id", ""),
-                        "spanType": span.get("attributes", {}).get("openinference.span.kind", ""),
-                    })
+                    dataset_spans.append(
+                        {
+                            "spanId": span.get("context", {}).get("span_id", ""),
+                            "spanName": span.get("name", ""),
+                            "spanHash": span.get("hash_id", ""),
+                            "spanType": span.get("attributes", {}).get(
+                                "openinference.span.kind", ""
+                            ),
+                        }
+                    )
                 except Exception as e:
                     logger.warning(f"Error processing span: {e}")
                     continue

ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py

@@ -1,26 +1,43 @@
-from aiohttp import payload
-import requests
 import json
+import logging
 import os
 import time
-import logging
+
+import requests
+
 from ragaai_catalyst.ragaai_catalyst import RagaAICatalyst
+
 logger = logging.getLogger(__name__)
-from urllib.parse import urlparse, urlunparse
 import re
+from urllib.parse import urlparse, urlunparse
+
 
-def upload_code(hash_id, zip_path, project_name, dataset_name, base_url=None, timeout=120):
-    code_hashes_list = _fetch_dataset_code_hashes(project_name, dataset_name, base_url, timeout=timeout)
+def upload_code(
+    hash_id, zip_path, project_name, dataset_name, base_url=None, timeout=120
+):
+    code_hashes_list = _fetch_dataset_code_hashes(
+        project_name, dataset_name, base_url, timeout=timeout
+    )
 
     if hash_id not in code_hashes_list:
-        presigned_url = _fetch_presigned_url(project_name, dataset_name, base_url, timeout=timeout)
+        presigned_url = _fetch_presigned_url(
+            project_name, dataset_name, base_url, timeout=timeout
+        )
         _put_zip_presigned_url(project_name, presigned_url, zip_path, timeout=timeout)
 
-        response = _insert_code(dataset_name, hash_id, presigned_url, project_name, base_url, timeout=timeout)
+        response = _insert_code(
+            dataset_name,
+            hash_id,
+            presigned_url,
+            project_name,
+            base_url,
+            timeout=timeout,
+        )
         return response
     else:
         return "Code already exists"
 
+
 def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None, timeout=120):
     payload = {}
     headers = {
@@ -32,32 +49,49 @@ def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None, timeou
         url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
         start_time = time.time()
         endpoint = f"{url_base}/v2/llm/dataset/code?datasetName={dataset_name}"
-        response = requests.request("GET", 
-                                    endpoint, 
-                                    headers=headers, 
-                                    data=payload,
-                                    timeout=timeout)
+        response = requests.request(
+            "GET", endpoint, headers=headers, data=payload, timeout=timeout
+        )
         elapsed_ms = (time.time() - start_time) * 1000
         logger.debug(
-            f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+            f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+        )
 
         if response.status_code == 200:
             return response.json()["data"]["codeHashes"]
-        else:
-            raise Exception(f"Failed to fetch code hashes: {response.json()['message']}")
+        elif response.status_code == 401:
+            logger.warning("Received 401 error. Attempting to refresh token.")
+            RagaAICatalyst.get_token(force_refresh=True)
+            headers = {
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Name": project_name,
+            }
+            response = requests.request(
+                "GET", endpoint, headers=headers, data=payload, timeout=timeout
+            )
+            elapsed_ms = (time.time() - start_time) * 1000
+            logger.debug(
+                f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
+            if response.status_code == 200:
+                return response.json()["data"]["codeHashes"]
+            else:
+                raise Exception(
+                    f"Failed to fetch code hashes: {response.json()['message']}"
+                )
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to list datasets: {e}")
-        raise 
+        raise
 
 
 def update_presigned_url(presigned_url, base_url):
     """Replaces the domain (and port, if applicable) of the presigned URL with that of the base URL."""
-    #To Do: If Proxy URL has domain name how do we handle such cases? Engineering Dependency.
-    
+    # To Do: If Proxy URL has domain name how do we handle such cases? Engineering Dependency.
+
     presigned_parts = urlparse(presigned_url)
     base_parts = urlparse(base_url)
     # Check if base_url contains localhost or an IP address
-    if re.match(r'^(localhost|\d{1,3}(\.\d{1,3}){3})$', base_parts.hostname):
+    if re.match(r"^(localhost|\d{1,3}(\.\d{1,3}){3})$", base_parts.hostname):
         new_netloc = base_parts.hostname  # Extract domain from base_url
         if base_parts.port:  # Add port if present in base_url
             new_netloc += f":{base_parts.port}"
@@ -67,11 +101,9 @@ def update_presigned_url(presigned_url, base_url):
 
 
 def _fetch_presigned_url(project_name, dataset_name, base_url=None, timeout=120):
-    payload = json.dumps({
-            "datasetName": dataset_name,
-            "numFiles": 1,
-            "contentType": "application/zip"
-            })
+    payload = json.dumps(
+        {"datasetName": dataset_name, "numFiles": 1, "contentType": "application/zip"}
+    )
 
     headers = {
         "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
@@ -84,93 +116,145 @@ def _fetch_presigned_url(project_name, dataset_name, base_url=None, timeout=120)
         start_time = time.time()
         # Changed to POST from GET
         endpoint = f"{url_base}/v1/llm/presigned-url"
-        response = requests.request("POST", 
-                                    endpoint, 
-                                    headers=headers, 
-                                    data=payload,
-                                    timeout=timeout)
+        response = requests.request(
+            "POST", endpoint, headers=headers, data=payload, timeout=timeout
+        )
         elapsed_ms = (time.time() - start_time) * 1000
         logger.debug(
-            f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+            f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+        )
 
         if response.status_code == 200:
             presigned_url = response.json()["data"]["presignedUrls"][0]
-            presigned_url = update_presigned_url(presigned_url,url_base)
+            presigned_url = update_presigned_url(presigned_url, url_base)
             return presigned_url
         else:
             # If POST fails, try GET
-            response = requests.request("GET", 
-                                    endpoint, 
-                                    headers=headers, 
-                                    data=payload,
-                                    timeout=timeout)
+            response = requests.request(
+                "POST", endpoint, headers=headers, data=payload, timeout=timeout
+            )
             elapsed_ms = (time.time() - start_time) * 1000
             logger.debug(
-                f"API Call: [GET] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
             if response.status_code == 200:
                 presigned_url = response.json()["data"]["presignedUrls"][0]
-                presigned_url = update_presigned_url(presigned_url,url_base)
+                presigned_url = update_presigned_url(presigned_url, url_base)
                 return presigned_url
-
-            logger.error(f"Failed to fetch code hashes: {response.json()['message']}")
-            raise Exception(f"Failed to fetch code hashes: {response.json()['message']}")
+            elif response.status_code == 401:
+                logger.warning("Received 401 error. Attempting to refresh token.")
+                RagaAICatalyst.get_token(force_refresh=True)
+                headers = {
+                    "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                    "Content-Type": "application/json",
+                    "X-Project-Name": project_name,
+                }
+                response = requests.request(
+                    "POST", endpoint, headers=headers, data=payload, timeout=timeout
+                )
+                elapsed_ms = (time.time() - start_time) * 1000
+                logger.debug(
+                    f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+                )
+                if response.status_code == 200:
+                    presigned_url = response.json()["data"]["presignedUrls"][0]
+                    presigned_url = update_presigned_url(presigned_url, url_base)
+                    return presigned_url
+                else:
+                    logger.error(
+                        f"Failed to fetch code hashes: {response.json()['message']}"
+                    )
+                    raise Exception(
+                        f"Failed to fetch code hashes: {response.json()['message']}"
+                    )
+            else:
+                logger.error(
+                    f"Failed to fetch code hashes: {response.json()['message']}"
+                )
+                raise Exception(
+                    f"Failed to fetch code hashes: {response.json()['message']}"
+                )
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to list datasets: {e}")
         raise
 
+
 def _put_zip_presigned_url(project_name, presignedUrl, filename, timeout=120):
     headers = {
-            "X-Project-Name": project_name,
-            "Content-Type": "application/zip",
-        }
+        "X-Project-Name": project_name,
+        "Content-Type": "application/zip",
+    }
 
     if "blob.core.windows.net" in presignedUrl:  # Azure
         headers["x-ms-blob-type"] = "BlockBlob"
-    print(f"Uploading code...")
-    with open(filename, 'rb') as f:
+    print("Uploading code...")
+    with open(filename, "rb") as f:
         payload = f.read()
 
     start_time = time.time()
-    response = requests.request("PUT", 
-                                presignedUrl, 
-                                headers=headers, 
-                                data=payload,
-                                timeout=timeout)
+    response = requests.request(
+        "PUT", presignedUrl, headers=headers, data=payload, timeout=timeout
+    )
     elapsed_ms = (time.time() - start_time) * 1000
     logger.debug(
-        f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+        f"API Call: [PUT] {presignedUrl} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+    )
     if response.status_code != 200 or response.status_code != 201:
         return response, response.status_code
 
-def _insert_code(dataset_name, hash_id, presigned_url, project_name, base_url=None, timeout=120):
-    payload = json.dumps({
-        "datasetName": dataset_name,
-        "codeHash": hash_id,
-        "presignedUrl": presigned_url
-        })
-    
-    headers = {
-        'X-Project-Name': project_name,
-        'Content-Type': 'application/json',
-        'Authorization': f'Bearer {os.getenv("RAGAAI_CATALYST_TOKEN")}'
+
+def _insert_code(
+    dataset_name, hash_id, presigned_url, project_name, base_url=None, timeout=120
+):
+    payload = json.dumps(
+        {
+            "datasetName": dataset_name,
+            "codeHash": hash_id,
+            "presignedUrl": presigned_url,
         }
-    
+    )
+
+    headers = {
+        "X-Project-Name": project_name,
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+    }
+
     try:
         url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
         start_time = time.time()
         endpoint = f"{url_base}/v2/llm/dataset/code"
-        response = requests.request("POST", 
-                                    endpoint, 
-                                    headers=headers, 
-                                    data=payload,
-                                    timeout=timeout)
+        response = requests.request(
+            "POST", endpoint, headers=headers, data=payload, timeout=timeout
+        )
         elapsed_ms = (time.time() - start_time) * 1000
         logger.debug(
-            f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
+            f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+        )
         if response.status_code == 200:
             return response.json()["message"]
+
+        elif response.status_code == 401:
+            logger.warning("Received 401 error. Attempting to refresh token.")
+            token = RagaAICatalyst.get_token(force_refresh=True)
+            headers = {
+                "X-Project-Name": project_name,
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {token}",
+            }
+            response = requests.request(
+                "POST", endpoint, headers=headers, data=payload, timeout=timeout
+            )
+            elapsed_ms = (time.time() - start_time) * 1000
+            logger.debug(
+                f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms"
+            )
+            if response.status_code == 200:
+                return response.json()["message"]
+            else:
+                raise Exception(f"Failed to insert code: {response.json()['message']}")
         else:
             raise Exception(f"Failed to insert code: {response.json()['message']}")
     except requests.exceptions.RequestException as e:
         logger.error(f"Failed to insert code: {e}")
-        raise
+        raise