Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: use text instead of content for ChatMessage in Ollama #1239

Merged
merged 1 commit into from
Dec 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "OllamaChatGenerator":
return default_from_dict(cls, data)

def _message_to_dict(self, message: ChatMessage) -> Dict[str, str]:
return {"role": message.role.value, "content": message.content}
return {"role": message.role.value, "content": message.text}

def _build_message_from_ollama_response(self, ollama_response: ChatResponse) -> ChatMessage:
"""
Expand Down
8 changes: 4 additions & 4 deletions integrations/ollama/tests/test_chat_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ def test_build_message_from_ollama_response(self):
observed = OllamaChatGenerator(model=model)._build_message_from_ollama_response(ollama_response)

assert observed.role == "assistant"
assert observed.content == "Hello! How are you today?"
assert observed.text == "Hello! How are you today?"

@pytest.mark.integration
def test_run(self):
Expand All @@ -121,7 +121,7 @@ def test_run(self):

assert isinstance(response, dict)
assert isinstance(response["replies"], list)
assert answer in response["replies"][0].content
assert answer in response["replies"][0].text

@pytest.mark.integration
def test_run_with_chat_history(self):
Expand All @@ -137,7 +137,7 @@ def test_run_with_chat_history(self):

assert isinstance(response, dict)
assert isinstance(response["replies"], list)
assert "Manchester" in response["replies"][-1].content or "Glasgow" in response["replies"][-1].content
assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text

@pytest.mark.integration
def test_run_model_unavailable(self):
Expand Down Expand Up @@ -166,4 +166,4 @@ def test_run_with_streaming(self):

assert isinstance(response, dict)
assert isinstance(response["replies"], list)
assert "Manchester" in response["replies"][-1].content or "Glasgow" in response["replies"][-1].content
assert "Manchester" in response["replies"][-1].text or "Glasgow" in response["replies"][-1].text
Loading