From 8cdb82bd9dbb7b5684c66658ab31f03f6fe0f065 Mon Sep 17 00:00:00 2001 From: Kye Date: Fri, 22 Dec 2023 21:59:57 -0500 Subject: [PATCH] [CLEANUP][ADA] --- swarms/models/simple_ada.py | 22 --------- tests/models/test_ada.py | 91 ------------------------------------- 2 files changed, 113 deletions(-) delete mode 100644 swarms/models/simple_ada.py delete mode 100644 tests/models/test_ada.py diff --git a/swarms/models/simple_ada.py b/swarms/models/simple_ada.py deleted file mode 100644 index e96995c4..00000000 --- a/swarms/models/simple_ada.py +++ /dev/null @@ -1,22 +0,0 @@ -from openai import OpenAI - -client = OpenAI() - - -def get_ada_embeddings( - text: str, model: str = "text-embedding-ada-002" -): - """ - Simple function to get embeddings from ada - - Usage: - >>> get_ada_embeddings("Hello World") - >>> get_ada_embeddings("Hello World", model="text-embedding-ada-001") - - """ - - text = text.replace("\n", " ") - - return client.embeddings.create(input=[text], model=model)[ - "data" - ][0]["embedding"] diff --git a/tests/models/test_ada.py b/tests/models/test_ada.py deleted file mode 100644 index 43895e79..00000000 --- a/tests/models/test_ada.py +++ /dev/null @@ -1,91 +0,0 @@ -# test_embeddings.py - -import pytest -import openai -from unittest.mock import patch -from swarms.models.simple_ada import ( - get_ada_embeddings, -) # Adjust this import path to your project structure -from os import getenv -from dotenv import load_dotenv - -load_dotenv() - - -# Fixture for test texts -@pytest.fixture -def test_texts(): - return [ - "Hello World", - "This is a test string with newline\ncharacters", - "A quick brown fox jumps over the lazy dog", - ] - - -# Basic Test -def test_get_ada_embeddings_basic(test_texts): - with patch("openai.resources.Embeddings.create") as mock_create: - # Mocking the OpenAI API call - mock_create.return_value = { - "data": [{"embedding": [0.1, 0.2, 0.3]}] - } - - for text in test_texts: - embedding = get_ada_embeddings(text) - assert embedding == [ - 0.1, - 0.2, - 0.3, - ], "Embedding does not match expected output" - mock_create.assert_called_with( - input=[text.replace("\n", " ")], - model="text-embedding-ada-002", - ) - - -# Parameterized Test -@pytest.mark.parametrize( - "text, model, expected_call_model", - [ - ( - "Hello World", - "text-embedding-ada-002", - "text-embedding-ada-002", - ), - ( - "Hello World", - "text-embedding-ada-001", - "text-embedding-ada-001", - ), - ], -) -def test_get_ada_embeddings_models(text, model, expected_call_model): - with patch("openai.resources.Embeddings.create") as mock_create: - mock_create.return_value = { - "data": [{"embedding": [0.1, 0.2, 0.3]}] - } - - _ = get_ada_embeddings(text, model=model) - mock_create.assert_called_with( - input=[text], model=expected_call_model - ) - - -# Exception Test -def test_get_ada_embeddings_exception(): - with patch("openai.resources.Embeddings.create") as mock_create: - mock_create.side_effect = openai.OpenAIError("Test error") - with pytest.raises(openai.OpenAIError): - get_ada_embeddings("Some text") - - -# Tests for environment variable loading -def test_env_var_loading(monkeypatch): - monkeypatch.setenv("OPENAI_API_KEY", "testkey123") - with patch("openai.resources.Embeddings.create"): - assert ( - getenv("OPENAI_API_KEY") == "testkey123" - ), "Environment variable for API key is not set correctly" - - -# ... more tests to cover other aspects such as different input types, large inputs, invalid inputs, etc.