ragaai-catalyst 2.2b0__py3-none-any.whl → 2.2.1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -1,16 +1,14 @@
-from ..data.data_structure import LLMCall
-from .trace_utils import (
-    calculate_cost,
-    convert_usage_to_dict,
-)
-from importlib import resources
+import asyncio
+
 #from litellm import model_cost
 import json
+import logging
 import os
-import asyncio
-import psutil
+from importlib import resources
+
 import tiktoken
-import logging
+
+from ..data.data_structure import LLMCall
 
 logger = logging.getLogger(__name__)
 
@@ -61,7 +59,7 @@ def extract_model_name(args, kwargs, result):
     try:
         if not model:
             model = result.raw.model
-    except Exception as e:
+    except Exception:
         pass
     
     
@@ -604,3 +602,18 @@ def extract_llm_data(args, kwargs, result):
         tool_call=parsed_tool_call,
     )
     return llm_data
+
+
+def count_tokens(input_str: str) -> int:
+    # Use tiktoken to count tokens
+    try:
+        import tiktoken
+        
+        # Use GPT-4o model's encoding (cl100k_base)
+        encoding = tiktoken.get_encoding("cl100k_base")
+        
+        # Count tokens
+        tokens = encoding.encode(input_str)
+        return len(tokens)
+    except Exception:
+        raise Exception("Failed to count tokens")

ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py

@@ -1,23 +1,23 @@
-import os
 import json
-import tempfile
-from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 import logging
+import os
+import tempfile
 from dataclasses import asdict
-from ragaai_catalyst.tracers.utils.trace_json_converter import convert_json_format
+
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
 from ragaai_catalyst.tracers.agentic_tracing.tracers.base import TracerJSONEncoder
+from ragaai_catalyst.tracers.agentic_tracing.upload.trace_uploader import (
+    submit_upload_task,
+)
 from ragaai_catalyst.tracers.agentic_tracing.utils.system_monitor import SystemMonitor
-from ragaai_catalyst.tracers.agentic_tracing.upload.trace_uploader import submit_upload_task
-from ragaai_catalyst.tracers.agentic_tracing.utils.zip_list_of_unique_files import zip_list_of_unique_files
-from ragaai_catalyst.tracers.agentic_tracing.utils.trace_utils import format_interactions
-from ragaai_catalyst.tracers.utils.rag_trace_json_converter import rag_trace_json_converter
-from ragaai_catalyst.tracers.utils.convert_langchain_callbacks_output import convert_langchain_callbacks_output
-from ragaai_catalyst.tracers.upload_traces import UploadTraces
-import datetime
-import logging
-import asyncio
-import concurrent.futures
-from functools import partial
+from ragaai_catalyst.tracers.agentic_tracing.utils.trace_utils import (
+    format_interactions,
+)
+from ragaai_catalyst.tracers.agentic_tracing.utils.zip_list_of_unique_files import (
+    zip_list_of_unique_files,
+)
+from ragaai_catalyst.tracers.utils.trace_json_converter import convert_json_format
 
 logger = logging.getLogger("RagaAICatalyst")
 logging_level = (

ragaai_catalyst/tracers/tracer.py

@@ -404,66 +404,27 @@ class Tracer(AgenticTracing):
 
         def recursive_mask_values(obj, parent_key=None):
             """Apply masking to all values in nested structure."""
-            if isinstance(obj, dict):
-                if self.tracer_type == "langchain":
-                    # Special handling for LangChain data
-                    if isinstance(obj, dict):
-                        if obj.get("name", "") == "retrieve_documents.langchain.workflow":
-                            prompt_structured_data = {
-                                "traceloop.entity.input": json.dumps({
-                                    "kwargs": {
-                                        "input": masking_func(json.loads(obj.get("attributes", {}).get("traceloop.entity.input", "")).get("kwargs", {}).get("input", "")),
-                                    }
-                                })
-                            }     
-                            prompt_data = {
-                                "name": "retrieve_documents.langchain.workflow",
-                                "attributes": prompt_structured_data,
-                            }
-                            return prompt_data
-                        elif obj.get("name", "") == "PromptTemplate.langchain.task":
-                            context_structured_data = {
-                                "traceloop.entity.input": json.dumps({
-                                    "kwargs": {
-                                        "context": masking_func(json.loads(obj.get("attributes", {}).get("traceloop.entity.input", "")).get("kwargs", {}).get("context", "")),
-                                    }
-                                }),
-                                "traceloop.entity.output": json.dumps({
-                                    "kwargs": {
-                                        "text": masking_func(json.loads(obj.get("attributes", {}).get("traceloop.entity.output", "")).get("kwargs", {}).get("text", "")),
-                                    }
-                                })
-                            }
-                            context_data = {
-                                "name": "PromptTemplate.langchain.task",
-                                "attributes": context_structured_data,
-                            }
-                            return context_data
-                        elif obj.get("name", "") == "ChatOpenAI.langchain.task":
-                            response_structured_data = {"gen_ai.completion.0.content": masking_func(obj.get("attributes", {}).get("gen_ai.completion.0.content", "")),
-                                                        "gen_ai.prompt.0.content": masking_func(obj.get("attributes", {}).get("gen_ai.prompt.0.content", ""))}
-                            response_data = {
-                                "name": "ChatOpenAI.langchain.task",
-                                "attributes" : response_structured_data
-                            }
-                            return response_data
-                else:
+            try:
+                if isinstance(obj, dict):
                     return {k: recursive_mask_values(v, k) for k, v in obj.items()}
-            elif isinstance(obj, list):
-                return [recursive_mask_values(item, parent_key) for item in obj]
-            elif isinstance(obj, str):
-                # List of keys that should NOT be masked
-                excluded_keys = {
-                    'start_time', 'end_time', 'name', 'id', 
-                    'hash_id', 'parent_id', 'source_hash_id',
-                    'cost', 'type', 'feedback', 'error', 'ctx','telemetry.sdk.version',
-                    'telemetry.sdk.language','service.name'
-                }
-                # Apply masking only if the key is NOT in the excluded list
-                if parent_key and parent_key.lower() not in excluded_keys:
-                    return masking_func(obj)
-                return obj
-            else:
+                elif isinstance(obj, list):
+                    return [recursive_mask_values(item, parent_key) for item in obj]
+                elif isinstance(obj, str):
+                    # List of keys that should NOT be masked
+                    excluded_keys = {
+                        'start_time', 'end_time', 'name', 'id', 
+                        'hash_id', 'parent_id', 'source_hash_id',
+                        'cost', 'type', 'feedback', 'error', 'ctx','telemetry.sdk.version',
+                        'telemetry.sdk.language','service.name'
+                    }
+                    # Apply masking only if the key is NOT in the excluded list
+                    if parent_key and parent_key.lower() not in excluded_keys:
+                        return masking_func(obj)
+                    return obj
+                else:
+                    return obj
+            except Exception as e:
+                logger.error(f"Error masking value: {e}")
                 return obj
 
         def file_post_processor(original_trace_json_path: os.PathLike) -> os.PathLike:
@@ -953,4 +914,4 @@ class Tracer(AgenticTracing):
             self.dynamic_exporter.user_details = user_details
             self.metadata = user_metadata
         else:
-            logger.warning("metadata must be a dictionary")
+            logger.warning("metadata must be a dictionary")

ragaai_catalyst/tracers/utils/trace_json_converter.py

@@ -1,12 +1,18 @@
 import json
+import logging
 import sys
 import uuid
 from datetime import datetime
-from typing import final, List, Dict, Any, Optional
+from typing import Any, Dict, List
+
 import pytz
-import uuid
-from ragaai_catalyst.tracers.agentic_tracing.utils.llm_utils import calculate_llm_cost, get_model_cost
-import logging
+
+from ragaai_catalyst.tracers.agentic_tracing.utils.llm_utils import (
+    calculate_llm_cost,
+    count_tokens,
+    get_model_cost,
+)
+
 logger = logging.getLogger(__name__)
 
 def convert_time_format(original_time_str, target_timezone_str="Asia/Kolkata"):
@@ -52,7 +58,6 @@ def get_ordered_family(parent_children_mapping: Dict[str, Any]) -> List[str]:
 
 def get_spans(input_trace):
     data = input_trace.copy()
-    import uuid
     from collections import defaultdict
         
     name_counts = defaultdict(int)
@@ -123,15 +128,28 @@ def convert_json_format(input_trace, custom_model_cost, user_context, user_gt,ex
 
         final_trace["data"][0]["spans"] = spans
         
-
-        # TODO: each span has token value from prompt ,completion and total tokens. i want the sum of all these tokens for each span
         # Calculate token counts and costs from spans
         for span in spans:
             if "attributes" in span:
                 # Extract token counts
                 prompt_tokens = span["attributes"].get("llm.token_count.prompt", 0)
                 completion_tokens = span["attributes"].get("llm.token_count.completion", 0)
-                
+
+                # If prompt tokens or/and completion tokens are not present, will calculate it using tiktoken
+                try:
+                    if prompt_tokens == 0:
+                        prompt_value = span["attributes"].get("input.value")
+                        if prompt_value:
+                            prompt_tokens = count_tokens(prompt_value)
+                            logger.debug(f"Prompt tokens not present, calculated it: {prompt_tokens}")
+                    if completion_tokens == 0:
+                        completion_value = span["attributes"].get("output.value")
+                        if completion_value:
+                            completion_tokens = count_tokens(completion_value)
+                            logger.debug(f"Completion tokens not present, calculated it: {completion_tokens}")
+                except Exception as e:
+                    logger.warning(f"Failed to calculate token counts: {e}")
+
                 # Update token counts
                 final_trace["metadata"]["tokens"]["prompt_tokens"] += prompt_tokens
                 final_trace["metadata"]["tokens"]["completion_tokens"] += completion_tokens

ragaai_catalyst/tracers/utils/utils.py

@@ -60,3 +60,4 @@ def get_unique_key(input_data):
     unique_key = hash_object.hexdigest()
 
     return unique_key
+    

{ragaai_catalyst-2.2b0.dist-info → ragaai_catalyst-2.2.1.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: ragaai_catalyst
-Version: 2.2b0
+Version: 2.2.1
 Summary: RAGA AI CATALYST
-Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>, Tushar Kumar <tushar.kumar@raga.ai>
+Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>, Tushar Kumar <tushar.kumar@raga.ai>, Rishabh Pandey <rishabh.pandey@raga.ai>, Jyotsana C G <jyotsana@raga.ai>
 Requires-Python: <=3.13.2,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE

{ragaai_catalyst-2.2b0.dist-info → ragaai_catalyst-2.2.1.dist-info}/RECORD

@@ -31,7 +31,7 @@ ragaai_catalyst/tracers/distributed.py,sha256=MwlBwIxCAng-OI-7Ove_rkE1mTLeuW4Jw-
 ragaai_catalyst/tracers/langchain_callback.py,sha256=CB75zzG3-DkYTELj0vI1MOHQTY0MuQJfoHIXz9Cl8S8,34568
 ragaai_catalyst/tracers/llamaindex_callback.py,sha256=ZY0BJrrlz-P9Mg2dX-ZkVKG3gSvzwqBtk7JL_05MiYA,14028
 ragaai_catalyst/tracers/llamaindex_instrumentation.py,sha256=Ys_jLkvVqo12bKgXDmkp4TxJu9HkBATrFE8cIcTYxWw,14329
-ragaai_catalyst/tracers/tracer.py,sha256=4GRBnYtNCu-mXZdLuWyL3eDcgh1zdcZ9tokMUyX826w,44561
+ragaai_catalyst/tracers/tracer.py,sha256=hdQK3-zV14lBqz8B6gLrMLNtfV34BmJ4-7eiNeLABq8,41931
 ragaai_catalyst/tracers/upload_traces.py,sha256=w1clGGfdOMpStUJX40NAlxe6dcFdN4pwcezyII0bGYA,6994
 ragaai_catalyst/tracers/agentic_tracing/README.md,sha256=X4QwLb7-Jg7GQMIXj-SerZIgDETfw-7VgYlczOR8ZeQ,4508
 ragaai_catalyst/tracers/agentic_tracing/__init__.py,sha256=yf6SKvOPSpH-9LiKaoLKXwqj5sez8F_5wkOb91yp0oE,260
@@ -64,7 +64,7 @@ ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py,sha256=Kx
 ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py,sha256=YG601l1a29ov9VPu9Vl4RXxgL7l16k54_WWnoTNoG58,2064
 ragaai_catalyst/tracers/agentic_tracing/utils/generic.py,sha256=WwXT01xmp8MSr7KinuDCSK9a1ifpLcT7ajFkvYviG_A,1190
 ragaai_catalyst/tracers/agentic_tracing/utils/get_user_trace_metrics.py,sha256=vPZ4dn4EHFW0kqd1GyRpsYXbfrRrd0DXCmh-pzsDBNE,1109
-ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py,sha256=McKB7TQchmFcgg2h0zg-inuxxKaRjcwbqV_OnRzzYEw,22387
+ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py,sha256=Tdi4nnxsXywx-X-4FOvyzRpvB04rgyd3Mzf4o9k_BAI,22697
 ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json,sha256=2tzGw_cKCTPcfjEm7iGvFE6pTw7gMTPzeBov_MTaXNY,321336
 ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py,sha256=qmODERcFZhc8MX24boFCXkkh6sJ-vZngRHPvxhyWFeE,4347
 ragaai_catalyst/tracers/agentic_tracing/utils/supported_llm_provider.toml,sha256=LvFDivDIE96Zasp-fgDEqUJ5GEQZUawQucR3aOcSUTY,926
@@ -76,7 +76,7 @@ ragaai_catalyst/tracers/exporters/__init__.py,sha256=wQbaqyeIjVZxYprHCKZ9BeiqxeX
 ragaai_catalyst/tracers/exporters/dynamic_trace_exporter.py,sha256=JhYVSN9t6-9MFcg0IrYEJIyD6rg9O96Arg2poQSb5g8,6855
 ragaai_catalyst/tracers/exporters/file_span_exporter.py,sha256=RgGteu-NVGprXKkynvyIO5yOjpbtA41R3W_NzCjnkwE,6445
 ragaai_catalyst/tracers/exporters/raga_exporter.py,sha256=6xvjWXyh8XPkHKSLLmAZUQSvwuyY17ov8pv2VdfI0qA,17875
-ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py,sha256=exiAHrHlMPWwwlQnMFeAYwjeicgUjCL97OKPHl-g5GI,9275
+ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py,sha256=0ybb8GQMN-WERum-L9LSbSvWm6hPDj7ywV3_pfgBMBA,8935
 ragaai_catalyst/tracers/instrumentators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpacDA0U3wg6Ybw,64
 ragaai_catalyst/tracers/utils/convert_langchain_callbacks_output.py,sha256=SehrD7q8ytAiUYoWr406b4mWs3Lk0Rcy6Ekkihh22TI,1703
@@ -86,10 +86,10 @@ ragaai_catalyst/tracers/utils/langchain_tracer_extraction_logic.py,sha256=XS2_x2
 ragaai_catalyst/tracers/utils/model_prices_and_context_window_backup.json,sha256=C3uwkibJ08C9sOX-54kulZYmJlIpZ-SQpfE6HNGrjbM,343502
 ragaai_catalyst/tracers/utils/rag_extraction_logic_final.py,sha256=3ygkRT__lLDRflRttjzPu28tIA8cTCiGQVMQjqMItqQ,11309
 ragaai_catalyst/tracers/utils/rag_trace_json_converter.py,sha256=54IEZO-YRjUAahV5nw8KClXqTF1LhfDry_TsZ4KGow4,20467
-ragaai_catalyst/tracers/utils/trace_json_converter.py,sha256=K4X4yUmB01UtFX-_xmJsgFOAmzGe8qQ6SYQRHUyWlKs,9405
-ragaai_catalyst/tracers/utils/utils.py,sha256=ViygfJ7vZ7U0CTSA1lbxVloHp4NSlmfDzBRNCJuMhis,2374
-ragaai_catalyst-2.2b0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ragaai_catalyst-2.2b0.dist-info/METADATA,sha256=Uj1EV64kf4yuPEMNIf3VhgguogYZaFdqNE4bTALMrcA,17603
-ragaai_catalyst-2.2b0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-ragaai_catalyst-2.2b0.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
-ragaai_catalyst-2.2b0.dist-info/RECORD,,
+ragaai_catalyst/tracers/utils/trace_json_converter.py,sha256=YllXhkfyYOuQ0rsX9VqiUjMUeztDMs8TFMudIPkWvrY,10191
+ragaai_catalyst/tracers/utils/utils.py,sha256=L19LQGc8h08FFhyptBtixIHGG_e-VdSPsKs7JNaXnGE,2378
+ragaai_catalyst-2.2.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ragaai_catalyst-2.2.1.dist-info/METADATA,sha256=GjsoMgBgsiwIVyWzWkz_k8y08gnD2TsdvnGJbTM7its,17677
+ragaai_catalyst-2.2.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+ragaai_catalyst-2.2.1.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
+ragaai_catalyst-2.2.1.dist-info/RECORD,,