added tests (#14)

* added tests

* updated pyproject

* fixed flake8

tests/fixtures.py

@@ -0,0 +1,29 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture
+def openai_llm_mock(request):
+ content = request.param
+ # Note the path here is adjusted to where OpenAI is used, not where it's defined
+ with patch("vision_agent.llm.llm.OpenAI") as mock:
+ # Setup a mock response structure that matches what your code expects
+ mock_instance = mock.return_value
+ mock_instance.chat.completions.create.return_value = MagicMock(
+ choices=[MagicMock(message=MagicMock(content=content))]
+ )
+ yield mock_instance
+
+
+@pytest.fixture
+def openai_lmm_mock(request):
+ content = request.param
+ # Note the path here is adjusted to where OpenAI is used, not where it's defined
+ with patch("vision_agent.lmm.lmm.OpenAI") as mock:
+ # Setup a mock response structure that matches what your code expects
+ mock_instance = mock.return_value
+ mock_instance.chat.completions.create.return_value = MagicMock(
+ choices=[MagicMock(message=MagicMock(content=content))]
+ )
+ yield mock_instance

tests/test_llm.py

@@ -0,0 +1,59 @@
+import pytest
+
+from vision_agent.llm.llm import OpenAILLM
+from vision_agent.tools import CLIP
+from vision_agent.tools.tools import GroundingDINO
+
+from .fixtures import openai_llm_mock # noqa: F401
+
+
+@pytest.mark.parametrize(
+ "openai_llm_mock", ["mocked response"], indirect=["openai_llm_mock"]
+)
+def test_generate_with_mock(openai_llm_mock): # noqa: F811
+ llm = OpenAILLM()
+ response = llm.generate("test prompt")
+ assert response == "mocked response"
+ openai_llm_mock.chat.completions.create.assert_called_once_with(
+ model="gpt-4-turbo-preview",
+ messages=[{"role": "user", "content": "test prompt"}],
+ )
+
+
+@pytest.mark.parametrize(
+ "openai_llm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_llm_mock"],
+)
+def test_generate_classifier(openai_llm_mock): # noqa: F811
+ llm = OpenAILLM()
+ prompt = "Can you generate a cat classifier?"
+ classifier = llm.generate_classifier(prompt)
+ assert isinstance(classifier, CLIP)
+ assert classifier.prompt == "cat"
+
+
+@pytest.mark.parametrize(
+ "openai_llm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_llm_mock"],
+)
+def test_generate_detector(openai_llm_mock): # noqa: F811
+ llm = OpenAILLM()
+ prompt = "Can you generate a cat detector?"
+ detector = llm.generate_detector(prompt)
+ assert isinstance(detector, GroundingDINO)
+ assert detector.prompt == "cat"
+
+
+@pytest.mark.parametrize(
+ "openai_llm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_llm_mock"],
+)
+def test_generate_segmentor(openai_llm_mock): # noqa: F811
+ llm = OpenAILLM()
+ prompt = "Can you generate a cat segmentor?"
+ segmentor = llm.generate_detector(prompt)
+ assert isinstance(segmentor, GroundingDINO)
+ assert segmentor.prompt == "cat"

tests/test_lmm.py

@@ -0,0 +1,72 @@
+import tempfile
+
+import pytest
+from PIL import Image
+
+from vision_agent.lmm.lmm import OpenAILMM
+from vision_agent.tools import CLIP, GroundingDINO, GroundingSAM
+
+from .fixtures import openai_lmm_mock # noqa: F401
+
+
+def create_temp_image(image_format="jpeg"):
+ temp_file = tempfile.NamedTemporaryFile(suffix=f".{image_format}", delete=False)
+ image = Image.new("RGB", (100, 100), color=(255, 0, 0))
+ image.save(temp_file, format=image_format)
+ temp_file.seek(0)
+ return temp_file.name
+
+
+@pytest.mark.parametrize(
+ "openai_lmm_mock", ["mocked response"], indirect=["openai_lmm_mock"]
+)
+def test_generate_with_mock(openai_lmm_mock): # noqa: F811
+ temp_image = create_temp_image()
+ lmm = OpenAILMM()
+ response = lmm.generate("test prompt", image=temp_image)
+ assert response == "mocked response"
+ assert (
+ "image_url"
+ in openai_lmm_mock.chat.completions.create.call_args.kwargs["messages"][0][
+ "content"
+ ][1]
+ )
+
+
+@pytest.mark.parametrize(
+ "openai_lmm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_lmm_mock"],
+)
+def test_generate_classifier(openai_lmm_mock): # noqa: F811
+ lmm = OpenAILMM()
+ prompt = "Can you generate a cat classifier?"
+ classifier = lmm.generate_classifier(prompt)
+ assert isinstance(classifier, CLIP)
+ assert classifier.prompt == "cat"
+
+
+@pytest.mark.parametrize(
+ "openai_lmm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_lmm_mock"],
+)
+def test_generate_classifier(openai_lmm_mock): # noqa: F811
+ lmm = OpenAILMM()
+ prompt = "Can you generate a cat classifier?"
+ detector = lmm.generate_detector(prompt)
+ assert isinstance(detector, GroundingDINO)
+ assert detector.prompt == "cat"
+
+
+@pytest.mark.parametrize(
+ "openai_lmm_mock",
+ ['{"Parameters": {"prompt": "cat"}}'],
+ indirect=["openai_lmm_mock"],
+)
+def test_generate_classifier(openai_lmm_mock): # noqa: F811
+ lmm = OpenAILMM()
+ prompt = "Can you generate a cat classifier?"
+ segmentor = lmm.generate_segmentor(prompt)
+ assert isinstance(segmentor, GroundingSAM)
+ assert segmentor.prompt == "cat"

vision_agent/llm/llm.py

@@ -1,6 +1,7 @@
 import json
 from abc import ABC, abstractmethod
 from typing import Mapping, cast
+from openai import OpenAI
 
 from vision_agent.tools import (
  CHOOSE_PARAMS,
@@ -22,8 +23,6 @@ class OpenAILLM(LLM):
  r"""An LLM class for any OpenAI LLM model."""
 
  def __init__(self, model_name: str = "gpt-4-turbo-preview"):
- from openai import OpenAI
-
  self.model_name = model_name
  self.client = OpenAI()
 

vision_agent/lmm/lmm.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, List, Mapping, Optional, Union, cast
 
 import requests
+from openai import OpenAI
 
 from vision_agent.tools import (
  CHOOSE_PARAMS,
@@ -59,17 +60,18 @@ def generate(
  json=data,
  )
  resp_json: Dict[str, Any] = res.json()
- if resp_json["statusCode"] != 200:
- _LOGGER.error(f"Request failed: {resp_json['data']}")
+ if (
+ "statusCode" in resp_json and resp_json["statusCode"] != 200
+ ) or "statusCode" not in resp_json:
+ _LOGGER.error(f"Request failed: {resp_json}")
+ raise ValueError(f"Request failed: {resp_json}")
  return cast(str, resp_json["data"])
 
 
 class OpenAILMM(LMM):
  r"""An LMM class for the OpenAI GPT-4 Vision model."""
 
  def __init__(self, model_name: str = "gpt-4-vision-preview"):
- from openai import OpenAI
-
  self.model_name = model_name
  self.client = OpenAI()