wehub-resource-sync
73 lines
3.0 KiB
Python
73 lines
3.0 KiB
Python
import os
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
from gui_agents.s3.core.mllm import LMMAgent
|
|
from gui_agents.s3.core.engine import LMMEngineOpenAI
|
|
|
|
|
|
class TestProviders(unittest.TestCase):
|
|
def setUp(self):
|
|
# Clear env vars before each test
|
|
if "OLLAMA_HOST" in os.environ:
|
|
del os.environ["OLLAMA_HOST"]
|
|
if "DEEPSEEK_API_KEY" in os.environ:
|
|
del os.environ["DEEPSEEK_API_KEY"]
|
|
if "QWEN_API_KEY" in os.environ:
|
|
del os.environ["QWEN_API_KEY"]
|
|
if "DEEPSEEK_ENDPOINT_URL" in os.environ:
|
|
del os.environ["DEEPSEEK_ENDPOINT_URL"]
|
|
if "QWEN_ENDPOINT_URL" in os.environ:
|
|
del os.environ["QWEN_ENDPOINT_URL"]
|
|
|
|
def test_ollama_missing_config(self):
|
|
"""Test that Ollama raises ValueError if no endpoint is provided"""
|
|
with self.assertRaises(ValueError) as cm:
|
|
LMMAgent(engine_params={"engine_type": "ollama", "model": "llama3"})
|
|
self.assertIn("Ollama endpoint must be provided", str(cm.exception))
|
|
|
|
def test_ollama_valid_config_param(self):
|
|
"""Test Ollama init with base_url param"""
|
|
agent = LMMAgent(
|
|
engine_params={
|
|
"engine_type": "ollama",
|
|
"model": "llama3",
|
|
"base_url": "http://example.com/v1",
|
|
}
|
|
)
|
|
self.assertIsInstance(agent.engine, LMMEngineOpenAI)
|
|
self.assertEqual(agent.engine.base_url, "http://example.com/v1")
|
|
|
|
def test_ollama_valid_config_env(self):
|
|
"""Test Ollama init with OLLAMA_HOST env var"""
|
|
with patch.dict(os.environ, {"OLLAMA_HOST": "http://env-host:11434"}):
|
|
agent = LMMAgent(engine_params={"engine_type": "ollama", "model": "llama3"})
|
|
self.assertIsInstance(agent.engine, LMMEngineOpenAI)
|
|
# Check for /v1 addition
|
|
self.assertEqual(agent.engine.base_url, "http://env-host:11434/v1")
|
|
|
|
def test_deepseek_init(self):
|
|
"""Test DeepSeek initialization"""
|
|
with patch.dict(os.environ, {"DEEPSEEK_API_KEY": "sk-test"}):
|
|
agent = LMMAgent(
|
|
engine_params={"engine_type": "deepseek", "model": "deepseek-coder"}
|
|
)
|
|
self.assertIsInstance(agent.engine, LMMEngineOpenAI)
|
|
# Default URL
|
|
self.assertEqual(agent.engine.base_url, "https://api.deepseek.com/v1")
|
|
# (Note: engine.py logic resolves default at generate() time or if client created,
|
|
# but init just stores what's passed. Let's verify prompt generation to ensure it doesn't crash on init)
|
|
|
|
def test_qwen_init(self):
|
|
"""Test Qwen initialization"""
|
|
with patch.dict(os.environ, {"QWEN_API_KEY": "sk-qwen"}):
|
|
agent = LMMAgent(engine_params={"engine_type": "qwen", "model": "qwen-max"})
|
|
self.assertIsInstance(agent.engine, LMMEngineOpenAI)
|
|
self.assertEqual(
|
|
agent.engine.base_url,
|
|
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|