From 748792c2b0f8824b5a1ee8979c6a8cfa0d8ef540 Mon Sep 17 00:00:00 2001
From: Anyc66666666 <1662048992@qq.com>
Date: Thu, 7 Dec 2023 15:48:36 +0800
Subject: [PATCH] feat: support to test trace and metric

---
 src/greptimeai/collector/collection.py        |  2 +-
 .../integrations/openai_tracker/test_sync.py  | 48 ++++++++++++-------
 2 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/src/greptimeai/collector/collection.py b/src/greptimeai/collector/collection.py
index 0b8cb22..5f17a80 100644
--- a/src/greptimeai/collector/collection.py
+++ b/src/greptimeai/collector/collection.py
@@ -367,7 +367,7 @@ def _setup_otel_exporter(self):
             timeout=5,
         )
         self._metric_reader = PeriodicExportingMetricReader(
-            metrics_exporter, export_interval_millis=3000
+            metrics_exporter, export_interval_millis=1000
         )
         metre_provider = MeterProvider(
             resource=resource, metric_readers=[self._metric_reader]
diff --git a/tests/integrations/openai_tracker/test_sync.py b/tests/integrations/openai_tracker/test_sync.py
index d4f6a41..ddb97c2 100644
--- a/tests/integrations/openai_tracker/test_sync.py
+++ b/tests/integrations/openai_tracker/test_sync.py
@@ -2,8 +2,8 @@
 import time
 import uuid
 
-from openai import OpenAI
 import pymysql
+from openai import OpenAI
 
 from greptimeai import openai_patcher
 
@@ -108,6 +108,9 @@ class Number(object):
 client = OpenAI()
 openai_patcher.setup(client=client)
 
+trace_sql = "SELECT model,prompt_tokens,completion_tokens FROM %s WHERE user_id = '%s'"
+metric_sql = "SELECT model,service_name,greptime_value FROM %s WHERE model = '%s'"
+
 
 def test_chat_completion():
     user_id = str(uuid.uuid4())
@@ -125,16 +128,23 @@ def test_chat_completion():
     assert resp.choices[0].message.content == "2"
 
     time.sleep(5)
-    sql = (
-        "select model,prompt_tokens,completion_tokens from %s where user_id = '%s'"
-        % (LlmTrace.table_name, user_id)
-    )
-    cursor.execute(sql)
-    result = cursor.fetchone()
+    cursor.execute(trace_sql % (LlmTrace.table_name, user_id))
+    trace = cursor.fetchone()
+    cursor.execute(metric_sql % (LlmPromptToken.table_name, resp.model))
+    prompt_token = cursor.fetchone()
+    cursor.execute(metric_sql % (LlmCompletionToken.table_name, resp.model))
+    completion_token = cursor.fetchone()
+
+    assert resp.model == trace[0]
 
-    assert resp.model == result[0]
-    assert resp.usage.prompt_tokens == result[1]
-    assert resp.usage.completion_tokens == result[2]
+    assert resp.usage.prompt_tokens == trace[1]
+    assert resp.usage.prompt_tokens == prompt_token[2]
+
+    assert resp.usage.completion_tokens == trace[2]
+    assert resp.usage.completion_tokens == completion_token[2]
+
+    assert "openai" == prompt_token[1]
+    assert "openai" == completion_token[1]
 
 
 def test_embedding():
@@ -146,12 +156,14 @@ def test_embedding():
     )
 
     time.sleep(5)
-    sql = "select model,prompt_tokens from %s where user_id = '%s'" % (
-        LlmTrace.table_name,
-        user_id,
-    )
-    cursor.execute(sql)
-    result = cursor.fetchone()
+    cursor.execute(trace_sql % (LlmTrace.table_name, user_id))
+    trace = cursor.fetchone()
+    cursor.execute(metric_sql % (LlmPromptToken.table_name, resp.model))
+    prompt_token = cursor.fetchone()
+
+    assert resp.model == trace[0]
+
+    assert resp.usage.prompt_tokens == trace[1]
+    assert resp.usage.prompt_tokens == prompt_token[2]
 
-    assert resp.model == result[0]
-    assert resp.usage.prompt_tokens == result[1]
+    assert "openai" == prompt_token[1]