From 714658a0c80e733f82a5a1cb7f33c84574251eb1 Mon Sep 17 00:00:00 2001
From: Nicholas Albion <nalbion@yahoo.com>
Date: Thu, 21 Sep 2023 00:25:25 +1000
Subject: [PATCH] tidy up tests

---
 pilot/helpers/AgentConvo.py        | 3 ++-
 pilot/utils/test_llm_connection.py | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/pilot/helpers/AgentConvo.py b/pilot/helpers/AgentConvo.py
index d6163d8..59ce3ae 100644
--- a/pilot/helpers/AgentConvo.py
+++ b/pilot/helpers/AgentConvo.py
@@ -23,7 +23,8 @@ class AgentConvo:
         agent: An instance of the agent participating in the conversation.
     """
     def __init__(self, agent):
-        self.messages = []
+        # [{'role': 'system'|'user'|'assistant', 'content': ''}, ...]
+        self.messages: list[dict] = []
         self.branches = {}
         self.log_to_user = True
         self.agent = agent
diff --git a/pilot/utils/test_llm_connection.py b/pilot/utils/test_llm_connection.py
index 93cdc5d..35c86e7 100644
--- a/pilot/utils/test_llm_connection.py
+++ b/pilot/utils/test_llm_connection.py
@@ -49,14 +49,14 @@ class TestLlmConnection:
                                                     })
 
         messages = convo.messages
-        # messages = [{"role": "user", "content": "I want to create a website"}]
 
         # When
         response = create_gpt_chat_completion(messages, '', function_calls=ARCHITECTURE)
+
         # Then
-        # You are and experienced software architect...
-        # You are working in a software development agency...
         assert len(convo.messages) == 2
+        assert convo.messages[0]['content'].startswith('You are an experienced software architect')
+        assert convo.messages[1]['content'].startswith('You are working in a software development agency')
         assert response is not None
         assert len(response) > 0
         technologies: list[str] = response['function_calls']['arguments']['technologies']