feat: add basic tests

Author: diogoncalves

libs/core/poetry.lock

@@ -15,6 +15,13 @@ uvicorn = "^0.27"
 openai = "^1.26"
 tiktoken = "^0.7"
 
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4"
+pytest-asyncio = "^0"
+mypy = "^1.5"
+ruff = "^0.0.284"
+black = "^23.7"
+pre-commit = "^3.7"
 
 [build-system]
 requires = ["poetry-core"]

libs/core/pyproject.toml

@@ -0,0 +1,39 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+from llmstudio_core.providers.provider import BaseProvider, ChatRequest, ProviderError
+
+class MockProvider(BaseProvider):
+    async def aparse_response(self, response, **kwargs):
+        return response
+
+    def parse_response(self, response, **kwargs):
+        return response
+
+    def chat(self, chat_input, model, **kwargs):
+        # Mock the response to match expected structure
+        return MagicMock(choices=[MagicMock(finish_reason="stop")])
+
+    async def achat(self, chat_input, model, **kwargs):
+        # Mock the response to match expected structure
+        return MagicMock(choices=[MagicMock(finish_reason="stop")])
+    
+    def output_to_string(self, output):
+        # Handle string inputs
+        if isinstance(output, str):
+            return output
+        if output.choices[0].finish_reason == "stop":
+            return output.choices[0].message.content
+        return ""
+
+    @staticmethod
+    def _provider_config_name():
+        return "mock_provider"
+
+@pytest.fixture
+def mock_provider():
+    config = MagicMock()
+    config.models = {"test_model": MagicMock(input_token_cost=0.01, output_token_cost=0.02)}
+    config.id = "mock_provider"
+    tokenizer = MagicMock()
+    tokenizer.encode = lambda x: x.split()  # Simple tokenizer mock
+    return MockProvider(config=config, tokenizer=tokenizer)

libs/core/tests/unit_tests/__init__.py

@@ -0,0 +1,53 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+from llmstudio_core.providers.provider import BaseProvider, ChatRequest, ProviderError
+
+request = ChatRequest(chat_input="Hello", model="test_model")
+
+def test_chat(mock_provider):
+    mock_provider.generate_client = MagicMock(return_value=MagicMock())
+    mock_provider.handle_response = MagicMock(return_value=iter(["response"]))
+    
+    print(request.model_dump())
+    response = mock_provider.chat(request.chat_input, request.model)
+    
+    assert response is not None
+
+@pytest.mark.asyncio
+async def test_achat(mock_provider):
+    mock_provider.agenerate_client = AsyncMock(return_value=AsyncMock())
+    mock_provider.ahandle_response = AsyncMock(return_value=AsyncMock())
+    
+    print(request.model_dump())
+    response = await mock_provider.achat(request.chat_input, request.model)
+    
+    assert response is not None
+
+
+def test_validate_model(mock_provider):
+    request = ChatRequest(chat_input="Hello", model="test_model")
+    mock_provider.validate_model(request)  # Should not raise
+
+    request_invalid = ChatRequest(chat_input="Hello", model="invalid_model")
+    with pytest.raises(ProviderError):
+        mock_provider.validate_model(request_invalid)
+
+def test_calculate_metrics(mock_provider):
+    metrics = mock_provider.calculate_metrics(
+        input="Hello",
+        output="World",
+        model="test_model",
+        start_time=0,
+        end_time=1,
+        first_token_time=0.5,
+        token_times=(0.1, 0.2),
+        token_count=2
+    )
+    
+    assert metrics["input_tokens"] == pytest.approx(1)
+    assert metrics["output_tokens"] == pytest.approx(1)
+    assert metrics["cost_usd"] == pytest.approx(0.03)
+    assert metrics["latency_s"] == pytest.approx(1)
+    assert metrics["time_to_first_token_s"] == pytest.approx(0.5)
+    assert metrics["inter_token_latency_s"] == pytest.approx(0.15)
+    assert metrics["tokens_per_second"] == pytest.approx(2)

libs/core/tests/unit_tests/conftest.py

@@ -23,6 +23,17 @@ openai = "^1.26"
 tiktoken = "^0.7"
 python-dotenv = "^0"
 langchain-experimental = "^0.0"
+llmstudio-core = "^0"
+llmstudio-tracker = "^0"
+llmstudio-proxy = "^0"
+
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4"
+mypy = "^1.5"
+ruff = "^0.0.284"
+black = "^23.7"
+pre-commit = "^3.7"
 llmstudio-core = { path = "../core/", develop = true }
 llmstudio-tracker = { path = "../tracker/", develop = true }
 llmstudio-proxy = { path = "../proxy/", develop = true }

libs/core/tests/unit_tests/test_provider.py

@@ -15,10 +15,12 @@ keywords = ["ml", "ai", "llm", "llmops", "openai", "langchain", "chatgpt", "llms
 
 [tool.poetry.dependencies]
 python = "^3.9"
+
+[tool.poetry.group.dev.dependencies]
 llmstudio = { path = "libs/llmstudio/", develop = true }
 llmstudio-core = { path = "libs/core/", develop = true }
 llmstudio-tracker = { path = "libs/tracker/", develop = true }
-llmstudio-proxy = { path = "libs/proxy/", develop = true }
+langchain-proxy = { path = "libs/proxy/", develop = true }
 
 [tool.poetry.scripts]
 llmstudio = "llmstudio.cli:main"