Much more testing

Author: pamelafox

.github/workflows/python.yaml

@@ -30,4 +30,4 @@ jobs:
           run: black . --check --verbose
         - name: Run unit tests
           run: |
-            python3 -m pytest
+            python3 -m pytest -s -vv --cov --cov-fail-under=99

tests/messages.py

@@ -1,19 +1,3 @@
-user_message = {
-    "message": {
-        "role": "user",
-        "content": "Hello, how are you?",
-    },
-    "count": 13,
-}
-
-user_message_unicode = {
-    "message": {
-        "role": "user",
-        "content": "á",
-    },
-    "count": 8,
-}
-
 system_message_short = {
     "message": {
         "role": "system",
@@ -30,6 +14,14 @@
     "count": 25,
 }
 
+system_message_long = {
+    "message": {
+        "role": "system",
+        "content": "Assistant helps the company employees with their healthcare plan questions, and questions about the employee handbook. Be brief in your answers.",
+    },
+    "count": 31,
+}
+
 system_message_unicode = {
     "message": {
         "role": "system",
@@ -47,6 +39,68 @@
     "count": 20,  # Less tokens in older vision preview models
 }
 
+user_message = {
+    "message": {
+        "role": "user",
+        "content": "Hello, how are you?",
+    },
+    "count": 13,
+}
+
+user_message_unicode = {
+    "message": {
+        "role": "user",
+        "content": "á",
+    },
+    "count": 8,
+}
+
+user_message_perf = {
+    "message": {
+        "role": "user",
+        "content": "What happens in a performance review?",
+    },
+    "count": 14,
+}
+
+assistant_message_perf = {
+    "message": {
+        "role": "assistant",
+        "content": "During the performance review at Contoso Electronics, the supervisor will discuss the employee's performance over the past year and provide feedback on areas for improvement. They will also provide an opportunity for the employee to discuss their goals and objectives for the upcoming year. The review is a two-way dialogue between managers and employees, and employees will receive a written summary of their performance review which will include a rating of their performance, feedback, and goals and objectives for the upcoming year [employee_handbook-3.pdf].",
+    },
+    "count": 106,
+}
+
+assistant_message_perf_short = {
+    "message": {
+        "role": "assistant",
+        "content": "The supervisor will discuss the employee's performance and provide feedback on areas for improvement. They will also provide an opportunity for the employee to discuss their goals and objectives for the upcoming year. The review is a two-way dialogue between managers and employees, and employees will receive a written summary of their performance review which will include a rating of their performance, feedback, and goals for the upcoming year [employee_handbook-3.pdf].",
+    },
+    "count": 91,
+}
+
+user_message_dresscode = {
+    "message": {
+        "role": "user",
+        "content": "Is there a dress code?",
+    },
+    "count": 13,
+}
+
+assistant_message_dresscode = {
+    "message": {
+        "role": "assistant",
+        "content": "Yes, there is a dress code at Contoso Electronics. Look sharp! [employee_handbook-1.pdf]",
+    },
+    "count": 30,
+}
+user_message_pm = {
+    "message": {
+        "role": "user",
+        "content": "What does a Product Manager do?",
+    },
+    "count": 14,
+}
 text_and_image_message = {
     "message": {
         "role": "user",
@@ -63,3 +117,20 @@
     },
     "count": 266,
 }
+
+MESSAGE_COUNTS = [
+    system_message,
+    system_message_short,
+    system_message_long,
+    system_message_unicode,
+    system_message_with_name,
+    user_message,
+    user_message_unicode,
+    user_message_perf,
+    user_message_dresscode,
+    user_message_pm,
+    assistant_message_perf,
+    assistant_message_perf_short,
+    assistant_message_dresscode,
+    text_and_image_message,
+]

tests/test_messagebuilder.py

@@ -1,7 +1,20 @@
 import pytest
 from openai_messages_token_helper import build_messages, count_tokens_for_message
 
-from .messages import system_message_short, system_message_unicode, user_message, user_message_unicode
+from .messages import (
+    assistant_message_dresscode,
+    assistant_message_perf,
+    assistant_message_perf_short,
+    system_message_long,
+    system_message_short,
+    system_message_unicode,
+    text_and_image_message,
+    user_message,
+    user_message_dresscode,
+    user_message_perf,
+    user_message_pm,
+    user_message_unicode,
+)
 
 
 def test_messagebuilder():
@@ -10,6 +23,15 @@ def test_messagebuilder():
     assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_short["count"]
 
 
+def test_messagebuilder_imagemessage():
+    messages = build_messages(
+        "gpt-35-turbo",
+        system_message_short["message"]["content"],
+        new_user_message=text_and_image_message["message"]["content"],
+    )
+    assert messages == [system_message_short["message"], text_and_image_message["message"]]
+
+
 def test_messagebuilder_append():
     messages = build_messages(
         "gpt-35-turbo", system_message_short["message"]["content"], new_user_message=user_message["message"]["content"]
@@ -55,3 +77,120 @@ def test_messagebuilder_model_fallback():
     assert messages == [system_message_short["message"], user_message["message"]]
     assert count_tokens_for_message(model, messages[0], default_to_cl100k=True) == system_message_short["count"]
     assert count_tokens_for_message(model, messages[1], default_to_cl100k=True) == user_message["count"]
+
+
+def test_messagebuilder_pastmessages():
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],  # 12 tokens
+        past_messages=[
+            user_message_perf["message"],  # 14 tokens
+            assistant_message_perf["message"],  # 106 tokens
+        ],
+        new_user_message=user_message_pm["message"]["content"],  # 14 tokens
+        max_tokens=3000,
+    )
+    assert messages == [
+        system_message_short["message"],
+        user_message_perf["message"],
+        assistant_message_perf["message"],
+        user_message_pm["message"],
+    ]
+
+
+def test_messagebuilder_pastmessages_truncated():
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],  # 12 tokens
+        past_messages=[
+            user_message_perf["message"],  # 14 tokens
+            assistant_message_perf["message"],  # 106 tokens
+        ],
+        new_user_message=user_message_pm["message"]["content"],  # 14 tokens
+        max_tokens=10,
+    )
+    assert messages == [system_message_short["message"], user_message_pm["message"]]
+
+
+def test_messagebuilder_pastmessages_truncated_longer():
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],  # 12 tokens
+        past_messages=[
+            user_message_perf["message"],  # 14 tokens
+            assistant_message_perf["message"],  # 106 tokens
+            user_message_dresscode["message"],  # 13 tokens
+            assistant_message_dresscode["message"],  # 30 tokens
+        ],
+        new_user_message=user_message_pm["message"]["content"],  # 14 tokens
+        max_tokens=69,
+    )
+    assert messages == [
+        system_message_short["message"],
+        user_message_dresscode["message"],
+        assistant_message_dresscode["message"],
+        user_message_pm["message"],
+    ]
+
+
+def test_messagebuilder_pastmessages_truncated_break_pair():
+    """Tests that the truncation breaks the pair of messages."""
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],  # 12 tokens
+        past_messages=[
+            user_message_perf["message"],  # 14 tokens
+            assistant_message_perf_short["message"],  # 91 tokens
+            user_message_dresscode["message"],  # 13 tokens
+            assistant_message_dresscode["message"],  # 30 tokens
+        ],
+        new_user_message=user_message_pm["message"]["content"],  # 14 tokens
+        max_tokens=160,
+    )
+    assert messages == [
+        system_message_short["message"],
+        assistant_message_perf_short["message"],
+        user_message_dresscode["message"],
+        assistant_message_dresscode["message"],
+        user_message_pm["message"],
+    ]
+
+
+def test_messagebuilder_system():
+    """Tests that the system message token count is considered."""
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_long["message"]["content"],  # 31 tokens
+        past_messages=[
+            user_message_perf["message"],  # 14 tokens
+            assistant_message_perf["message"],  # 106 tokens
+            user_message_dresscode["message"],  # 13 tokens
+            assistant_message_dresscode["message"],  # 30 tokens
+        ],
+        new_user_message=user_message_pm["message"]["content"],  # 14 tokens
+        max_tokens=36,
+    )
+    assert messages == [system_message_long["message"], user_message_pm["message"]]
+
+
+def test_messagebuilder_system_fewshots():
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],
+        new_user_message=user_message_pm["message"]["content"],
+        past_messages=[],
+        few_shots=[
+            {"role": "user", "content": "How did crypto do last year?"},
+            {"role": "assistant", "content": "Summarize Cryptocurrency Market Dynamics from last year"},
+            {"role": "user", "content": "What are my health plans?"},
+            {"role": "assistant", "content": "Show available health plans"},
+        ],
+    )
+    # Make sure messages are in the right order
+    assert messages[0]["role"] == "system"
+    assert messages[1]["role"] == "user"
+    assert messages[2]["role"] == "assistant"
+    assert messages[3]["role"] == "user"
+    assert messages[4]["role"] == "assistant"
+    assert messages[5]["role"] == "user"
+    assert messages[5]["content"] == user_message_pm["message"]["content"]

tests/verify_openai.py

@@ -3,15 +3,7 @@
 import azure.identity
 import openai
 from dotenv import load_dotenv
-from messages import (
-    system_message,
-    system_message_short,
-    system_message_unicode,
-    system_message_with_name,
-    text_and_image_message,
-    user_message,
-    user_message_unicode,
-)
+from messages import MESSAGE_COUNTS
 
 # Setup the OpenAI client to use either Azure OpenAI or OpenAI API
 load_dotenv()
@@ -32,15 +24,7 @@
     MODEL_NAME = os.getenv("OPENAI_MODEL")
 
 # Test the token count for each message
-for message_count_pair in [
-    user_message,
-    user_message_unicode,
-    system_message,
-    system_message_short,
-    system_message_unicode,
-    system_message_with_name,
-    text_and_image_message,
-]:
+for message_count_pair in MESSAGE_COUNTS:
     response = client.chat.completions.create(
         model=MODEL_NAME,
         temperature=0.7,