ragaai-catalyst 2.1.5b32__py3-none-any.whl → 2.1.5b33__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py

@@ -1,7 +1,13 @@
 import requests
 import json
 import os
+import time
+import logging
 from datetime import datetime
+from urllib.parse import urlparse, urlunparse
+import re
+
+logger = logging.getLogger(__name__)
 
 
 class UploadAgenticTraces:
@@ -33,17 +39,39 @@ class UploadAgenticTraces:
         }
 
         try:
+            start_time = time.time()
+            endpoint = f"{self.base_url}/v1/llm/presigned-url"
             response = requests.request("GET", 
-                                        f"{self.base_url}/v1/llm/presigned-url", 
+                                        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")
+            
             if response.status_code == 200:
-                presignedUrls = response.json()["data"]["presignedUrls"][0]
-                return presignedUrls
+                presignedURLs = response.json()["data"]["presignedUrls"][0]
+                presignedurl = self.update_presigned_url(presignedURLs,self.base_url)
+                return presignedurl
+            
         except requests.exceptions.RequestException as e:
             print(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 
+        with that of the base URL only if the base URL contains 'localhost' or an IP address."""
+        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):
+            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}"
+            updated_parts = presigned_parts._replace(netloc=new_netloc)
+            return urlunparse(updated_parts)
+        return presigned_url
 
     def _put_presigned_url(self, presignedUrl, filename):
         headers = {
@@ -60,11 +88,15 @@ class UploadAgenticTraces:
             print(f"Error while reading file: {e}")
             return None
         try:
+            start_time = time.time()
             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")
             if response.status_code != 200 or response.status_code != 201:
                 return response, response.status_code
         except requests.exceptions.RequestException as e:
@@ -83,11 +115,16 @@ class UploadAgenticTraces:
                 "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", 
-                                        f"{self.base_url}/v1/llm/insert/trace", 
+                                        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
@@ -116,27 +153,43 @@ class UploadAgenticTraces:
                             "spanType": span["type"],
                         })
                 else:
-                    datasetSpans.append({
+                    datasetSpans.extend(self._get_agent_dataset_spans(span, datasetSpans))
+            datasetSpans = [dict(t) for t in set(tuple(sorted(d.items())) for d in datasetSpans)]
+            
+            return datasetSpans
+        except Exception as e:
+            print(f"Error while reading dataset spans: {e}")
+            return None
+
+    def _get_agent_dataset_spans(self, span, datasetSpans):
+        datasetSpans.append({
                                 "spanId": span["id"],
                                 "spanName": span["name"],
                                 "spanHash": span["hash_id"],
                                 "spanType": span["type"],
                             })
-                    children = span["data"]["children"]
-                    for child in children:
-                        existing_span = next((s for s in datasetSpans if s["spanHash"] == child["hash_id"]), None)
-                        if existing_span is None:
-                            datasetSpans.append({
-                                "spanId": child["id"],
-                                "spanName": child["name"],
-                                "spanHash": child["hash_id"],
-                                "spanType": child["type"],
-                            })
-            return datasetSpans
-        except Exception as e:
-            print(f"Error while reading dataset spans: {e}")
-            return None
-    
+        children = span["data"]["children"]
+        for child in children:
+            if child["type"] != "agent":
+                existing_span = next((s for s in datasetSpans if s["spanHash"] == child["hash_id"]), None)
+                if existing_span is None:
+                    datasetSpans.append({
+                        "spanId": child["id"],
+                        "spanName": child["name"],
+                        "spanHash": child["hash_id"],
+                        "spanType": child["type"],
+                    })
+            else:
+                datasetSpans.append({
+                            "spanId": child["id"],
+                            "spanName": child["name"],
+                            "spanHash": child["hash_id"],
+                            "spanType": child["type"],
+                        })
+                self._get_agent_dataset_spans(child, datasetSpans)
+        return datasetSpans
+        
+
     def upload_agentic_traces(self):
         try:
             presignedUrl = self._get_presigned_url()

ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py

@@ -2,23 +2,26 @@ from aiohttp import payload
 import requests
 import json
 import os
+import time
 import logging
 from ragaai_catalyst.ragaai_catalyst import RagaAICatalyst
 logger = logging.getLogger(__name__)
+from urllib.parse import urlparse, urlunparse
+import re
 
-def upload_code(hash_id, zip_path, project_name, dataset_name):
-    code_hashes_list = _fetch_dataset_code_hashes(project_name, dataset_name)
+def upload_code(hash_id, zip_path, project_name, dataset_name, base_url=None):
+    code_hashes_list = _fetch_dataset_code_hashes(project_name, dataset_name, base_url)
 
     if hash_id not in code_hashes_list:
-        presigned_url = _fetch_presigned_url(project_name, dataset_name)
+        presigned_url = _fetch_presigned_url(project_name, dataset_name, base_url)
         _put_zip_presigned_url(project_name, presigned_url, zip_path)
 
-        response = _insert_code(dataset_name, hash_id, presigned_url, project_name)
+        response = _insert_code(dataset_name, hash_id, presigned_url, project_name, base_url)
         return response
     else:
         return "Code already exists"
 
-def _fetch_dataset_code_hashes(project_name, dataset_name):
+def _fetch_dataset_code_hashes(project_name, dataset_name, base_url=None):
     payload = {}
     headers = {
         "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
@@ -26,11 +29,17 @@ def _fetch_dataset_code_hashes(project_name, dataset_name):
     }
 
     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?datasetName={dataset_name}"
         response = requests.request("GET", 
-                                    f"{RagaAICatalyst.BASE_URL}/v2/llm/dataset/code?datasetName={dataset_name}", 
+                                    endpoint, 
                                     headers=headers, 
                                     data=payload,
                                     timeout=99999)
+        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"]
@@ -40,7 +49,22 @@ def _fetch_dataset_code_hashes(project_name, dataset_name):
         logger.error(f"Failed to list datasets: {e}")
         raise 
 
-def _fetch_presigned_url(project_name, dataset_name):
+
+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."""
+    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):
+        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}"
+        updated_parts = presigned_parts._replace(netloc=new_netloc)
+        return urlunparse(updated_parts)
+    return presigned_url
+
+
+def _fetch_presigned_url(project_name, dataset_name, base_url=None):
     payload = json.dumps({
             "datasetName": dataset_name,
             "numFiles": 1,
@@ -54,14 +78,22 @@ def _fetch_presigned_url(project_name, dataset_name):
     }
 
     try:
+        url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
+        start_time = time.time()
+        endpoint = f"{url_base}/v1/llm/presigned-url"
         response = requests.request("GET", 
-                                    f"{RagaAICatalyst.BASE_URL}/v1/llm/presigned-url", 
+                                    endpoint, 
                                     headers=headers, 
                                     data=payload,
                                     timeout=99999)
+        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"]["presignedUrls"][0]
+            presigned_url = response.json()["data"]["presignedUrls"][0]
+            presigned_url = update_presigned_url(presigned_url,RagaAICatalyst.BASE_URL)
+            return presigned_url
         else:
             raise Exception(f"Failed to fetch code hashes: {response.json()['message']}")
     except requests.exceptions.RequestException as e:
@@ -80,15 +112,19 @@ def _put_zip_presigned_url(project_name, presignedUrl, filename):
     with open(filename, 'rb') as f:
         payload = f.read()
 
+    start_time = time.time()
     response = requests.request("PUT", 
                                 presignedUrl, 
                                 headers=headers, 
                                 data=payload,
                                 timeout=99999)
+    elapsed_ms = (time.time() - start_time) * 1000
+    logger.debug(
+        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):
+def _insert_code(dataset_name, hash_id, presigned_url, project_name, base_url=None):
     payload = json.dumps({
         "datasetName": dataset_name,
         "codeHash": hash_id,
@@ -102,11 +138,17 @@ def _insert_code(dataset_name, hash_id, presigned_url, project_name):
         }
     
     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", 
-                                    f"{RagaAICatalyst.BASE_URL}/v2/llm/dataset/code", 
+                                    endpoint, 
                                     headers=headers, 
                                     data=payload,
                                     timeout=99999)
+        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:

ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py

@@ -3,6 +3,7 @@ import logging
 import requests
 import os
 import json
+import time
 from ....ragaai_catalyst import RagaAICatalyst
 from ..utils.get_user_trace_metrics import get_user_trace_metrics
 
@@ -14,7 +15,7 @@ logging_level = (
 )
 
 
-def upload_trace_metric(json_file_path, dataset_name, project_name):
+def upload_trace_metric(json_file_path, dataset_name, project_name, base_url=None):
     try:
         with open(json_file_path, "r") as f:
             traces = json.load(f)
@@ -43,11 +44,17 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
             "datasetName": dataset_name,
             "metrics": metrics
         })
+        url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
+        start_time = time.time()
+        endpoint = f"{url_base}/v1/llm/trace/metrics"
         response = requests.request("POST",
-                                    f"{RagaAICatalyst.BASE_URL}/v1/llm/trace/metrics",
+                                    endpoint,
                                     headers=headers,
                                     data=payload,
                                     timeout=10)
+        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:
             raise ValueError(f"Error inserting agentic trace metrics")
     except requests.exceptions.RequestException as e:

ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py

@@ -4,7 +4,7 @@ import re
 import requests
 from ragaai_catalyst.tracers.agentic_tracing.tracers.base import RagaAICatalyst
 
-def create_dataset_schema_with_trace(project_name, dataset_name):
+def create_dataset_schema_with_trace(project_name, dataset_name, base_url=None):
     def make_request():
         headers = {
             "Content-Type": "application/json",
@@ -15,8 +15,10 @@ def create_dataset_schema_with_trace(project_name, dataset_name):
             "datasetName": dataset_name,
             "traceFolderUrl": None,
         })
+        # Use provided base_url or fall back to default
+        url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
         response = requests.request("POST",
-            f"{RagaAICatalyst.BASE_URL}/v1/llm/dataset/logs",
+            f"{url_base}/v1/llm/dataset/logs",
             headers=headers,
             data=payload,
             timeout=10

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -4,7 +4,7 @@ from .trace_utils import (
     convert_usage_to_dict,
 )
 from importlib import resources
-from litellm import model_cost
+#from litellm import model_cost
 import json
 import os
 import asyncio
@@ -14,6 +14,15 @@ import logging
 
 logger = logging.getLogger(__name__)
 
+def get_model_cost():
+    file="model_prices_and_context_window_backup.json"
+    d={}
+    with resources.open_text("ragaai_catalyst.tracers.utils", file) as f:
+        d= json.load(f)
+    return d 
+
+model_cost = get_model_cost()
+
 def extract_model_name(args, kwargs, result):
     """Extract model name from kwargs or result"""
     # First try direct model parameter