Update service

lion_service/imodel.py

@@ -9,6 +9,7 @@
 
 
 class iModel:
+
     def __init__(
         self,
         provider: str | Service,
@@ -25,7 +26,29 @@ def __init__(
         elif api_key_schema is not None:
             api_key = api_key_schema
 
+        if task == "chat":
+            match provider:
+                case "openai":
+                    task = "create_chat_completion"
+                case "anthropic":
+                    task = "create_message"
+                case "groq":
+                    task = "create_chat_completion"
+                case "perplexity":
+                    task = "create_chat_completion"
+
         if isinstance(provider, str):
+            if api_key is None:
+                match provider:
+                    case "openai":
+                        api_key = "OPENAI_API_KEY"
+                    case "anthropic":
+                        api_key = "ANTHROPIC_API_KEY"
+                    case "groq":
+                        api_key = "GROQ_API_KEY"
+                    case "perplexity":
+                        api_key = "PERPLEXIY_API_KEY"
+
             self.service = match_service(provider, api_key=api_key, **kwargs)
         elif isinstance(provider, Service):
             self.service = provider
@@ -67,6 +90,7 @@ def __init__(
         self.data_model = self.service.match_data_model(self.task)
 
     def parse_to_data_model(self, **kwargs):
+
         if kwargs.get("model") and self.model:
             if kwargs.get("model") != self.model:
                 raise ValueError(
@@ -90,8 +114,9 @@ def parse_to_data_model(self, **kwargs):
     async def invoke(self, **kwargs):
         return await self.request_model.invoke(**kwargs)
 
-    def list_tasks(self):
-        return self.service.list_tasks()
+    @property
+    def allowed_roles(self):
+        return self.service.allowed_roles
 
 
 __all__ = ["iModel"]

lion_service/service.py

@@ -3,6 +3,10 @@ class Service:
     def list_tasks(self):
         pass
 
+    @property
+    def allowed_roles(self):
+        return ["user", "assistant", "system"]
+
 
 def register_service(cls):
     original_init = cls.__init__

lion_service/version.py

@@ -1 +1 @@
-__version__ = "1.1.1"
+__version__ = "1.2.0"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "lion-service"
-version = "1.1.1"
+version = "1.2.0"
 description = "lion api service system"
 authors = [
     { name = "HaiyangLi", email = "[email protected]" },
@@ -56,3 +56,18 @@ exclude = [".git", "__pycache__", "build", "dist", ".venv"]
 asyncio_mode = "auto"
 testpaths = ["tests"]
 python_files = ["test_*.py"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["lion_service"]
+
+[tool.hatch.build]
+exclude = [
+    "dev/*",
+    "data/*",
+    "notebooks/*",
+    "tests/*",
+    "*.pyc",
+    "__pycache__",
+    "temp_logs/*",
+    "logs/*"
+]

tests/test_complete_request_info.py

@@ -0,0 +1,70 @@
+"""Tests for complete_request_info module."""
+
+import pytest
+
+from lion_service.complete_request_info import (
+    CompleteRequestInfo,
+    CompleteRequestTokenInfo,
+)
+
+
+def test_complete_request_info():
+    """Test CompleteRequestInfo model."""
+    timestamp = 1234567890.123
+    info = CompleteRequestInfo(timestamp=timestamp)
+    assert info.timestamp == timestamp
+
+
+def test_complete_request_token_info():
+    """Test CompleteRequestTokenInfo model."""
+    timestamp = 1234567890.123
+    token_usage = 100
+    info = CompleteRequestTokenInfo(
+        timestamp=timestamp, token_usage=token_usage
+    )
+    assert info.timestamp == timestamp
+    assert info.token_usage == token_usage
+
+
+def test_complete_request_token_info_inheritance():
+    """Test that CompleteRequestTokenInfo inherits from CompleteRequestInfo."""
+    info = CompleteRequestTokenInfo(timestamp=1234567890.123, token_usage=100)
+    assert isinstance(info, CompleteRequestInfo)
+
+
+def test_complete_request_info_validation():
+    """Test validation of CompleteRequestInfo."""
+    with pytest.raises(ValueError):
+        CompleteRequestInfo(timestamp="not a float")
+
+
+def test_complete_request_token_info_validation():
+    """Test validation of CompleteRequestTokenInfo."""
+    with pytest.raises(ValueError):
+        CompleteRequestTokenInfo(
+            timestamp=1234567890.123, token_usage="not an int"
+        )
+
+    with pytest.raises(ValueError):
+        CompleteRequestTokenInfo(timestamp="not a float", token_usage=100)
+
+
+def test_complete_request_info_model_fields():
+    """Test model fields are properly defined."""
+    info = CompleteRequestInfo(timestamp=1234567890.123)
+    assert "timestamp" in info.model_fields
+    assert (
+        info.model_fields["timestamp"].description
+        == "HTTP response generated time"
+    )
+
+
+def test_complete_request_token_info_model_fields():
+    """Test model fields are properly defined."""
+    info = CompleteRequestTokenInfo(timestamp=1234567890.123, token_usage=100)
+    assert "timestamp" in info.model_fields
+    assert "token_usage" in info.model_fields
+    assert (
+        info.model_fields["token_usage"].description
+        == "Number of tokens used in the request"
+    )

tests/test_imodel.py

@@ -0,0 +1,189 @@
+"""Tests for imodel module."""
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+from pydantic import BaseModel
+
+from lion_service.imodel import iModel
+from lion_service.service import Service
+
+
+class MockDataModel(BaseModel):
+    model: str = None
+
+
+class MockRequestModel:
+    def __init__(self):
+        self.invoke = AsyncMock(return_value="response")
+
+
+class MockService(Service):
+    def __init__(self, name=None):
+        super().__init__()
+        self.name = name or f"mock_{id(self)}"  # Unique name for each instance
+        self.match_data_model = Mock(return_value={"request": MockDataModel})
+        self._request_model = MockRequestModel()
+
+    def chat(self):
+        return self._request_model
+
+    def list_tasks(self):
+        return ["chat", "completion"]
+
+
+def test_imodel_initialization():
+    """Test iModel initialization with different parameters."""
+    service = MockService()
+
+    model = iModel(
+        provider=service, task="chat", model="gpt-4", api_key="test_key"
+    )
+
+    assert model.service == service
+    assert model.task == "chat"
+    assert model.model == "gpt-4"
+
+
+def test_imodel_initialization_with_string_provider():
+    """Test iModel initialization with provider name."""
+    with patch("lion_service.imodel.match_service") as mock_match:
+        mock_service = MockService()
+        mock_match.return_value = mock_service
+
+        model = iModel(
+            provider="mock_provider",
+            task="chat",
+            model="gpt-4",
+            api_key="test_key",
+        )
+
+        assert model.service == mock_service
+        assert model.task == "chat"
+        assert model.model == "gpt-4"
+        mock_match.assert_called_once_with("mock_provider", api_key="test_key")
+
+
+def test_imodel_with_invalid_provider():
+    """Test iModel initialization with invalid provider."""
+    with pytest.raises(ValueError, match="Invalid provider"):
+        iModel(provider=123)
+
+
+def test_imodel_with_invalid_task():
+    """Test iModel initialization with invalid task."""
+    service = MockService()
+
+    with pytest.raises(ValueError, match="No matching task found"):
+        iModel(provider=service, task="invalid_task")
+
+
+def test_imodel_with_ambiguous_task():
+    """Test iModel initialization with ambiguous task."""
+
+    class AmbiguousService(MockService):
+        def chat_completion(self):
+            return self._request_model
+
+        def chat_summary(self):
+            return self._request_model
+
+        def list_tasks(self):
+            return ["chat_completion", "chat_summary"]
+
+    service = AmbiguousService()
+
+    with pytest.raises(ValueError, match="Multiple possible tasks found"):
+        iModel(provider=service, task="chat")
+
+
+def test_parse_to_data_model():
+    """Test parsing parameters to data model."""
+    service = MockService()
+
+    model = iModel(provider=service, task="chat", model="gpt-4")
+
+    parsed = model.parse_to_data_model(temperature=0.7)
+    assert isinstance(parsed["request"], MockDataModel)
+    assert parsed["request"].model == "gpt-4"
+
+
+def test_parse_to_data_model_with_inconsistent_model():
+    """Test parsing with inconsistent model parameter."""
+    service = MockService()
+
+    model = iModel(provider=service, task="chat", model="gpt-4")
+
+    with pytest.raises(ValueError, match="Models are inconsistent"):
+        model.parse_to_data_model(model="gpt-3.5-turbo")
+
+
+@pytest.mark.asyncio
+async def test_invoke():
+    """Test invoke method."""
+    service = MockService()
+
+    model = iModel(provider=service, task="chat", model="gpt-4")
+
+    result = await model.invoke(
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+    assert result == "response"
+    model.request_model.invoke.assert_called_once_with(
+        messages=[{"role": "user", "content": "Hello"}]
+    )
+
+
+def test_imodel_with_api_key_schema():
+    """Test initialization with api_key_schema."""
+    with patch("lion_service.imodel.match_service") as mock_match:
+        mock_service = MockService()
+        mock_match.return_value = mock_service
+
+        model = iModel(
+            provider="mock_provider",
+            task="chat",
+            model="gpt-4",
+            api_key_schema="test_key",
+        )
+
+        mock_match.assert_called_once_with("mock_provider", api_key="test_key")
+
+
+def test_imodel_with_rate_limits():
+    """Test initialization with rate limits."""
+    service = MockService()
+
+    model = iModel(
+        provider=service,
+        task="chat",
+        model="gpt-4",
+        interval_tokens=1000,
+        interval_requests=10,
+    )
+
+    assert model.request_model is not None
+
+
+def test_imodel_with_additional_kwargs():
+    """Test initialization with additional kwargs."""
+    service = MockService()
+
+    model = iModel(
+        provider=service, task="chat", model="gpt-4", custom_param="value"
+    )
+
+    assert model.configs["custom_param"] == "value"
+
+
+def test_imodel_with_service_instance_and_api_key():
+    """Test initialization with both service instance and api_key."""
+    service = MockService()
+
+    with pytest.warns(
+        UserWarning,
+        match="A Service instance was provided along with api key info",
+    ):
+        model = iModel(
+            provider=service, task="chat", model="gpt-4", api_key="test_key"
+        )