feat: init integration test

.github/workflows/ci.yml

@@ -21,6 +21,7 @@ jobs:
   ci:
     runs-on: ubuntu-latest
     strategy:
+      max-parallel: 1
       matrix:
         python-version: ["3.8", "3.9", "3.10", "3.11"]
 
@@ -41,10 +42,17 @@ jobs:
 
       - name: Install Dev dependencies
         run: |
-          pip install langchain openai pytest ruff mypy
+          pip install langchain openai pytest ruff mypy pymysql pytest-asyncio
 
       # TODO(yuanbohan): code coverage with pytest-cov
       - name: Test with pytest
+        env:
+          GREPTIMEAI_HOST: ${{ secrets.GREPTIMEAI_HOST }}
+          GREPTIMEAI_USERNAME: ${{ secrets.GREPTIMEAI_USERNAME }}
+          GREPTIMEAI_PASSWORD: ${{ secrets.GREPTIMEAI_PASSWORD }}
+          GREPTIMEAI_DATABASE: ${{ secrets.GREPTIMEAI_DATABASE }}
+          GREPTIMEAI_TOKEN: ${{ secrets.GREPTIMEAI_TOKEN }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         run: |
           pytest tests/ -v
 
@@ -56,3 +64,7 @@ jobs:
       - name: Static Type Checking with mypy
         run: |
           mypy $(git ls-files '*.py') --check-untyped-defs
+
+concurrency:
+  group: ${{ github.repository }}
+  cancel-in-progress: true

pyproject.toml

@@ -37,6 +37,8 @@ dev-dependencies = [
     "langchain>=0.0.27",
     "openai>=1.3.5",
     "mypy>=1.7.1",
+    "pymysql>=1.1.0",
+    "pytest-asyncio>=0.23.2",
 ]
 
 [tool.rye.scripts]

requirements-dev.lock

@@ -22,7 +22,6 @@ distro==1.8.0
 exceptiongroup==1.1.3
 frozenlist==1.4.0
 googleapis-common-protos==1.61.0
-greenlet==3.0.1
 h11==0.14.0
 httpcore==1.0.2
 httpx==0.25.1
@@ -47,16 +46,17 @@ opentelemetry-sdk==1.20.0
 opentelemetry-semantic-conventions==0.41b0
 packaging==23.2
 pluggy==1.3.0
-protobuf==4.24.4
+protobuf==4.25.1
 pydantic==2.4.2
 pydantic-core==2.10.1
+pymysql==1.1.0
 pytest==7.4.3
 pyyaml==6.0.1
 regex==2023.10.3
 requests==2.31.0
 ruff==0.1.3
 sniffio==1.3.0
-sqlalchemy==2.0.22
+sqlalchemy==2.0.23
 tenacity==8.2.3
 tiktoken==0.5.1
 tomli==2.0.1

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=15000
+            metrics_exporter, export_interval_millis=1000
         )
         metre_provider = MeterProvider(
             resource=resource, metric_readers=[self._metric_reader]

tests/integrations/database/db.py

@@ -0,0 +1,75 @@
+import logging
+import os
+from typing import Tuple
+
+import pymysql  # type: ignore
+
+from .model import Tables
+
+db = pymysql.connect(
+    host=os.getenv("GREPTIMEAI_HOST"),
+    user=os.getenv("GREPTIMEAI_USERNAME"),
+    passwd=os.getenv("GREPTIMEAI_PASSWORD"),
+    port=4002,
+    db=os.getenv("GREPTIMEAI_DATABASE"),
+)
+cursor = db.cursor()
+
+trace_sql = "SELECT model,prompt_tokens,completion_tokens FROM %s WHERE user_id = '%s'"
+trace_stream_sql = "SELECT model,completion_tokens FROM %s WHERE user_id = '%s'"
+metric_sql = "SELECT service_name,greptime_value FROM %s WHERE model = '%s'"
+truncate_sql = "TRUNCATE %s"
+
+
+def get_trace_data(user_id: str, is_stream: bool) -> Tuple:
+    """
+    get trace data for llm trace by user_id
+    :param user_id:
+    :return: model, prompt_tokens, completion_tokens
+    """
+    if is_stream:
+        cursor.execute(trace_stream_sql % (Tables.llm_trace, user_id))
+    else:
+        cursor.execute(trace_sql % (Tables.llm_trace, user_id))
+    trace = cursor.fetchone()
+    if trace is None:
+        raise Exception("trace data is None")
+    return trace
+
+
+def get_metric_data(table: str, model: str) -> Tuple:
+    """
+    get metric data by table and model
+    :param table:
+    :param model:
+    :return: service_name, greptime_value
+    """
+    cursor.execute(metric_sql % (table, model))
+    metric = cursor.fetchone()
+    if metric is None:
+        raise Exception("metric data is None")
+    return metric
+
+
+def truncate_tables():
+    """
+    truncate all tables
+    :return:
+    """
+    tables = [
+        "llm_completion_tokens",
+        "llm_completion_tokens_cost",
+        "llm_errors",
+        "llm_prompt_tokens",
+        "llm_prompt_tokens_cost",
+        "llm_request_duration_ms_bucket",
+        "llm_request_duration_ms_count",
+        "llm_request_duration_ms_sum",
+        "llm_traces_preview_v01",
+    ]
+    try:
+        cursor.executemany(truncate_sql, tables)
+        db.commit()
+    except Exception as e:
+        logging.error(e)
+        db.rollback()

tests/integrations/database/model.py

@@ -0,0 +1,88 @@
+class LlmTrace(object):
+    table_name = "llm_traces_preview_v01"
+
+    trace_id: str
+    span_id: str
+    parent_span_id: str
+    resource_attributes: str
+    scope_name: str
+    scope_version: str
+    scope_attributes: str
+    trace_state: str
+    span_name: str
+    span_kind: str
+    span_status_code: str
+    span_status_message: str
+    span_attributes: str
+    span_events: str
+    span_links: str
+    start: float
+    end: float
+    user_id: str
+    model: str
+    prompt_tokens: int
+    prompt_cost: float
+    completion_tokens: int
+    completion_cost: float
+    greptime_value: str
+    greptime_timestamp: float
+
+
+class LlmPromptToken(object):
+    table_name = "llm_prompt_tokens"
+
+    telemetry_sdk_language: str
+    telemetry_sdk_name: str
+    telemetry_sdk_version: str
+    service_name: str
+    span_name: str
+    model: str
+    greptime_value: str
+    greptime_timestamp: float
+
+
+class LlmPromptTokenCost(object):
+    table_name = "llm_prompt_tokens_cost"
+
+    telemetry_sdk_language: str
+    telemetry_sdk_name: str
+    telemetry_sdk_version: str
+    service_name: str
+    span_name: str
+    model: str
+    greptime_value: str
+    greptime_timestamp: float
+
+
+class LlmCompletionToken(object):
+    table_name = "llm_completion_tokens"
+
+    telemetry_sdk_language: str
+    telemetry_sdk_name: str
+    telemetry_sdk_version: str
+    service_name: str
+    span_name: str
+    model: str
+    greptime_value: str
+    greptime_timestamp: float
+
+
+class LlmCompletionTokenCost(object):
+    table_name = "llm_completion_tokens_cost"
+
+    telemetry_sdk_language: str
+    telemetry_sdk_name: str
+    telemetry_sdk_version: str
+    service_name: str
+    span_name: str
+    model: str
+    greptime_value: str
+    greptime_timestamp: float
+
+
+class Tables(object):
+    llm_trace = "llm_traces_preview_v01"
+    llm_prompt_tokens = "llm_prompt_tokens"
+    llm_prompt_tokens_cost = "llm_prompt_tokens_cost"
+    llm_completion_tokens = "llm_completion_tokens"
+    llm_completion_tokens_cost = "llm_completion_tokens_cost"

tests/integrations/openai_tracker/__init__.py

@@ -0,0 +1,10 @@
+from openai import AsyncOpenAI
+from openai import OpenAI
+
+from greptimeai import openai_patcher  # type: ignore
+
+async_client = AsyncOpenAI()
+openai_patcher.setup(client=async_client)
+
+client = OpenAI()
+openai_patcher.setup(client=client)

tests/integrations/openai_tracker/test_async.py

@@ -0,0 +1,43 @@
+import time
+import uuid
+
+import pytest
+
+from ..database.db import (
+    get_trace_data,
+    truncate_tables,
+)
+from ..openai_tracker import async_client
+
+
+@pytest.fixture
+def _truncate_tables():
+    truncate_tables()
+    yield
+
+
+@pytest.mark.asyncio
+async def test_chat_completion(_truncate_tables):
+    user_id = str(uuid.uuid4())
+    resp = await async_client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "1+1=",
+            }
+        ],
+        model="gpt-3.5-turbo",
+        user=user_id,
+        seed=1,
+    )
+    print("---------------------------------------", resp)
+    assert resp.choices[0].message.content == "2"
+
+    time.sleep(8)
+    trace = get_trace_data(user_id, False)
+
+    assert resp.model == trace[0]
+
+    if resp.usage:
+        assert resp.usage.prompt_tokens == trace[1]
+        assert resp.usage.completion_tokens == trace[2]

tests/integrations/openai_tracker/test_raw_response.py

@@ -0,0 +1,44 @@
+import json
+import time
+import uuid
+
+import pytest
+
+from ..database.db import (
+    get_trace_data,
+    truncate_tables,
+)
+from ..openai_tracker import client
+
+
+@pytest.fixture
+def _truncate_tables():
+    truncate_tables()
+    yield
+
+
+def test_chat_completion(_truncate_tables):
+    user_id = str(uuid.uuid4())
+    resp = client.with_raw_response.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "1+1=",
+            }
+        ],
+        model="gpt-3.5-turbo",
+        user=user_id,
+        seed=1,
+    )
+
+    data = json.loads(resp.content)
+    assert data["choices"][0]["message"]["content"] == "2"
+
+    time.sleep(6)
+    trace = get_trace_data(user_id, False)
+
+    assert data["model"] == trace[0]
+
+    if "usage" in data:
+        assert data["usage"]["prompt_tokens"] == trace[1]
+        assert data["usage"]["completion_tokens"] == trace[2]

tests/integrations/openai_tracker/test_retry.py

@@ -0,0 +1 @@
+# test_retry

tests/integrations/openai_tracker/test_stream.py

@@ -0,0 +1,48 @@
+import time
+import uuid
+
+import pytest
+
+from ..database.db import (
+    get_trace_data,
+    truncate_tables,
+)
+from ..openai_tracker import client
+
+
+@pytest.fixture
+def _truncate_tables():
+    truncate_tables()
+    yield
+
+
+def test_chat_completion(_truncate_tables):
+    user_id = str(uuid.uuid4())
+    resp = client.chat.completions.create(
+        messages=[
+            {
+                "role": "user",
+                "content": "1+1=",
+            }
+        ],
+        model="gpt-3.5-turbo",
+        user=user_id,
+        seed=1,
+        stream=True,
+    )
+
+    ans = ""
+    model = ""
+    for item in resp:
+        model = item.model
+        for choice in item.choices:
+            if choice.delta.content:
+                ans += choice.delta.content
+
+    assert ans == "2"
+    time.sleep(6)
+
+    trace = get_trace_data(user_id, True)
+
+    assert trace[0] in model  # model
+    assert trace[1] == 2  # completion tokens