Merge pull request #140 from aristotle-ai/model_cost_latency_metadata

Model-cost-latency as metadata

Author: kiranscaria

ragaai_catalyst/tracers/tracer.py

@@ -525,6 +525,14 @@ def stop(self):
                 combined_metadata.update(user_detail['trace_user_detail']['metadata'])
             if additional_metadata:
                 combined_metadata.update(additional_metadata)
+                
+            model_cost_latency_metadata = {}
+            if additional_metadata:
+                model_cost_latency_metadata["model"] = additional_metadata["model_name"]
+                model_cost_latency_metadata["total_cost"] = additional_metadata["cost"]
+                model_cost_latency_metadata["total_latency"] = additional_metadata["latency"]
+                model_cost_latency_metadata["recorded_on"] = datetime.datetime.now().astimezone().isoformat()
+                combined_metadata.update(model_cost_latency_metadata)
 
             langchain_traces = langchain_tracer_extraction(data, self.user_context)
             final_result = convert_langchain_callbacks_output(langchain_traces)

ragaai_catalyst/tracers/upload_traces.py

@@ -28,11 +28,14 @@ def _create_dataset_schema_with_trace(self, additional_metadata_keys=None, addit
             "response":{"columnType": "response"},
             "context": {"columnType": "context"},
             "llm_model": {"columnType":"pipeline"},
-            "recorded_on": {"columnType": "metadata"},
+            "recorded_on": {"columnType": "timestamp"},
             "embed_model": {"columnType":"pipeline"},
             "log_source": {"columnType": "metadata"},
             "vector_store":{"columnType":"pipeline"},
-            "feedback": {"columnType":"feedBack"}
+            "feedback": {"columnType":"feedBack"},
+            "model": {"columnType": "metadata"},
+            "total_cost": {"columnType": "metadata", "dataType": "numerical"},
+            "total_latency": {"columnType": "metadata", "dataType": "numerical"},
         }
 
         if additional_metadata_keys: