From 66dad7e21f6639d4556dbb1532b71e2cd9fce53d Mon Sep 17 00:00:00 2001
From: Shankar <90070882+shankar-landing-ai@users.noreply.github.com>
Date: Tue, 14 May 2024 11:10:08 -0700
Subject: [PATCH] Keep GPT4-Turbo for old vision agent (#82)

* adding gpt4t as default model for old vision agent

* adding test case

* fix test case
---
 tests/test_llm.py                  | 13 +++++++++++++
 vision_agent/agent/vision_agent.py |  8 +++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/tests/test_llm.py b/tests/test_llm.py
index b10ba7fc..40c8f8bc 100644
--- a/tests/test_llm.py
+++ b/tests/test_llm.py
@@ -36,6 +36,19 @@ def test_chat_with_mock(openai_llm_mock):  # noqa: F811
     )
 
 
+@pytest.mark.parametrize(
+    "openai_llm_mock_turbo", ["mocked response"], indirect=["openai_llm_mock_turbo"]
+)
+def openai_llm_mock_turbo(openai_llm_mock_2):  # noqa: F811
+    llm = OpenAILLM()
+    response = llm.chat([{"role": "user", "content": "test prompt"}])
+    assert response == "mocked response"
+    openai_llm_mock_turbo.chat.completions.create.assert_called_once_with(
+        model="gpt-4-turbo",
+        messages=[{"role": "user", "content": "test prompt"}],
+    )
+
+
 @pytest.mark.parametrize(
     "openai_llm_mock", ["mocked response"], indirect=["openai_llm_mock"]
 )
diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index 0e63a70b..42fcbee6 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -464,15 +464,17 @@ def __init__(
             report_progress_callback: a callback to report the progress of the agent. This is useful for streaming logs in a web application where multiple VisionAgent instances are running in parallel. This callback ensures that the progress are not mixed up.
         """
         self.task_model = (
-            OpenAILLM(json_mode=True, temperature=0.1)
+            OpenAILLM(model_name="gpt-4-turbo", json_mode=True, temperature=0.0)
             if task_model is None
             else task_model
         )
         self.answer_model = (
-            OpenAILLM(temperature=0.1) if answer_model is None else answer_model
+            OpenAILLM(model_name="gpt-4-turbo", temperature=0.0)
+            if answer_model is None
+            else answer_model
         )
         self.reflect_model = (
-            OpenAILMM(json_mode=True, temperature=0.1)
+            OpenAILMM(model_name="gpt-4-turbo", json_mode=True, temperature=0.0)
             if reflect_model is None
             else reflect_model
         )