Skip to content

Commit

Permalink
top_k and top_p transposed in vertexai (langchain-ai#5673)
Browse files Browse the repository at this point in the history
Fix transposed properties in vertexai model


Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
  • Loading branch information
2 people authored and Undertone0809 committed Jun 19, 2023
1 parent 7d503c3 commit 41bbd0c
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 2 deletions.
4 changes: 2 additions & 2 deletions langchain/llms/vertexai.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ def _default_params(self) -> Dict[str, Any]:
base_params = {
"temperature": self.temperature,
"max_output_tokens": self.max_output_tokens,
"top_k": self.top_p,
"top_p": self.top_k,
"top_k": self.top_k,
"top_p": self.top_p,
}
return {**base_params}

Expand Down
30 changes: 30 additions & 0 deletions tests/integration_tests/chat_models/test_vertexai.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
Your end-user credentials would be used to make the calls (make sure you've run
`gcloud auth login` first).
"""
from unittest.mock import Mock, patch

import pytest

from langchain.chat_models import ChatVertexAI
Expand Down Expand Up @@ -86,3 +88,31 @@ def test_vertexai_single_call_failes_no_message() -> None:
str(exc_info.value)
== "You should provide at least one message to start the chat!"
)


def test_vertexai_args_passed() -> None:
response_text = "Goodbye"
user_prompt = "Hello"
prompt_params = {
"max_output_tokens": 1,
"temperature": 10000.0,
"top_k": 10,
"top_p": 0.5,
}

# Mock the library to ensure the args are passed correctly
with patch(
"vertexai.language_models._language_models.ChatSession.send_message"
) as send_message:
mock_response = Mock(text=response_text)
send_message.return_value = mock_response

model = ChatVertexAI(**prompt_params)
message = HumanMessage(content=user_prompt)
response = model([message])

assert response.content == response_text
send_message.assert_called_once_with(
user_prompt,
**prompt_params,
)

0 comments on commit 41bbd0c

Please sign in to comment.