Revert "types: enable passing messages with arbitrary role (#462)"

This reverts commit 07eec6d517.
This commit is contained in:
ParthSareen 2025-03-20 17:37:44 -04:00
parent 07eec6d517
commit f8a2e5034d
2 changed files with 2 additions and 33 deletions

View File

@ -256,7 +256,7 @@ class Message(SubscriptableBaseModel):
Chat message.
"""
role: str
role: Literal['user', 'assistant', 'system', 'tool']
"Assumed role of the message. Response messages has role 'assistant' or 'tool'."
content: Optional[str] = None

View File

@ -4,16 +4,14 @@ import os
import re
import tempfile
from pathlib import Path
from typing import Any
import pytest
from httpx import Response as httpxResponse
from pydantic import BaseModel, ValidationError
from pytest_httpserver import HTTPServer, URIPattern
from werkzeug.wrappers import Request, Response
from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools
from ollama._types import Image, Message
from ollama._types import Image
PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC'
PNG_BYTES = base64.b64decode(PNG_BASE64)
@ -1183,32 +1181,3 @@ async def test_async_client_connection_error():
with pytest.raises(ConnectionError) as exc_info:
await client.show('model')
assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
def test_arbitrary_roles_accepted_in_message():
_ = Message(role='somerandomrole', content="I'm ok with you adding any role message now!")
def _mock_request(*args: Any, **kwargs: Any) -> Response:
return httpxResponse(status_code=200, content="{'response': 'Hello world!'}")
def test_arbitrary_roles_accepted_in_message_request(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(Client, '_request', _mock_request)
client = Client()
client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}])
async def _mock_request_async(*args: Any, **kwargs: Any) -> Response:
return httpxResponse(status_code=200, content="{'response': 'Hello world!'}")
@pytest.mark.asyncio
async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(AsyncClient, '_request', _mock_request_async)
client = AsyncClient()
await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}])