From ea247e70ed9083bd8fe2ee055869a7a14333b441 Mon Sep 17 00:00:00 2001
From: Torsten Raudssus <torsten@raudss.us>
Date: Sun, 12 Jan 2025 21:21:01 +0100
Subject: [PATCH] Fixing response content handling on groq and openai

---
 pyproject.toml                   |  2 +-
 tackleberry/runtime/groq.py      |  2 +-
 tackleberry/runtime/openai.py    |  2 +-
 tests/test_110_standard_query.py | 46 ++++++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_110_standard_query.py

diff --git a/pyproject.toml b/pyproject.toml
index 8875559..a1eef47 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "tackleberry"
-version = "0.1.0"
+version = "0.1.1"
 description = "Tackleberry (or TB) is helping you tackle the access to AI"
 authors = [
     { name = "Torsten Raudßus", email = "torsten@raudssus.de" },
diff --git a/tackleberry/runtime/groq.py b/tackleberry/runtime/groq.py
index 9dfb810..3e9013a 100644
--- a/tackleberry/runtime/groq.py
+++ b/tackleberry/runtime/groq.py
@@ -43,7 +43,7 @@ def chat_context(self, chat: TBChat, context: TBContext, struct: BaseModel = Non
                 model=chat.model.name,
                 messages=self.get_messages_from_context(context),
             )
-            return response.content
+            return response.choices[0].message.content
 
     def __str__(self):
         return f"TB Runtime Groq {hex(id(self))}"
diff --git a/tackleberry/runtime/openai.py b/tackleberry/runtime/openai.py
index 5374fb7..ceab119 100644
--- a/tackleberry/runtime/openai.py
+++ b/tackleberry/runtime/openai.py
@@ -42,7 +42,7 @@ def chat_context(self, chat: TBChat, context: TBContext, struct: BaseModel = Non
                 model=chat.model.name,
                 messages=self.get_messages_from_context(context),
             )
-            return response.content
+            return response.choices[0].message.content
 
     def __str__(self):
         return f"TB Runtime OpenAI {hex(id(self))}"
diff --git a/tests/test_110_standard_query.py b/tests/test_110_standard_query.py
new file mode 100644
index 0000000..e635c57
--- /dev/null
+++ b/tests/test_110_standard_query.py
@@ -0,0 +1,46 @@
+import unittest
+import warnings
+import os
+from unittest.mock import patch
+import requests
+from tackleberry import TB
+
+import sys
+
+class TestTB(unittest.TestCase):
+
+    def test_010_openai(self):
+        """Test OpenAI"""
+        if os.environ.get("OPENAI_API_KEY"):
+            chat = TB.chat('openai/gpt-4o-mini')
+            test_resp = chat.query("Say test")
+        else:
+            warnings.warn("Can't test OpenAI runtime without OPENAI_API_KEY", UserWarning)
+    
+    def test_020_anthropic(self):
+        """Test Anthropic"""
+        if os.environ.get("ANTHROPIC_API_KEY"):
+            chat = TB.chat('anthropic/claude-3-5-haiku-20241022')
+            test_resp = chat.query("Say test")
+        else:
+            warnings.warn("Can't test Anthropic runtime without ANTHROPIC_API_KEY", UserWarning)
+    
+    def test_030_groq(self):
+        """Test Groq"""
+        if os.environ.get("GROQ_API_KEY"):
+            chat = TB.chat('groq/llama3-8b-8192')
+            test_resp = chat.query("Say test")
+        else:
+            warnings.warn("Can't test Groq runtime without GROQ_API_KEY", UserWarning)
+
+    def test_040_ollama(self):
+        """Test Ollama"""
+        ollama_model = os.environ.get("TACKLEBERRY_OLLAMA_TEST_MODEL") or 'gemma2:2b'
+        if (os.environ.get("OLLAMA_HOST") or os.environ.get("OLLAMA_PROXY_URL")) and ollama_model:
+            chat = TB.chat('ollama/'+ollama_model)
+            test_resp = chat.query("Say test")
+        else:
+            warnings.warn("Can't test Ollama runtime without explicit setting OLLAMA_HOST or OLLAMA_PROXY_URL", UserWarning)
+
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file