Skip to content

Commit

Permalink
(ChatOpenAI) Add model_name to LLMResult.llm_output
Browse files Browse the repository at this point in the history
This makes sure OpenAI and ChatOpenAI have the same llm_output, and allow tracking usage per model.
Same work for OpenAI was done in langchain-ai#1713.
  • Loading branch information
mario-at-intercom committed Mar 24, 2023
1 parent 4f364db commit 43e08fa
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 13 deletions.
25 changes: 12 additions & 13 deletions langchain/chat_models/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,16 +91,6 @@ def _convert_message_to_dict(message: BaseMessage) -> dict:
return message_dict


def _create_chat_result(response: Mapping[str, Any]) -> ChatResult:
generations = []
for res in response["choices"]:
message = _convert_dict_to_message(res["message"])
gen = ChatGeneration(message=message)
generations.append(gen)
llm_output = {"token_usage": response["usage"]}
return ChatResult(generations=generations, llm_output=llm_output)


class ChatOpenAI(BaseChatModel, BaseModel):
"""Wrapper around OpenAI Chat large language models.
Expand Down Expand Up @@ -237,7 +227,7 @@ def _combine_llm_outputs(self, llm_outputs: List[Optional[dict]]) -> dict:
overall_token_usage[k] += v
else:
overall_token_usage[k] = v
return {"token_usage": overall_token_usage}
return {"token_usage": overall_token_usage, "model_name": self.model_name}

def _generate(
self, messages: List[BaseMessage], stop: Optional[List[str]] = None
Expand All @@ -262,7 +252,7 @@ def _generate(
)
return ChatResult(generations=[ChatGeneration(message=message)])
response = self.completion_with_retry(messages=message_dicts, **params)
return _create_chat_result(response)
return self._create_chat_result(response)

def _create_message_dicts(
self, messages: List[BaseMessage], stop: Optional[List[str]]
Expand All @@ -275,6 +265,15 @@ def _create_message_dicts(
message_dicts = [_convert_message_to_dict(m) for m in messages]
return message_dicts, params

def _create_chat_result(self, response: Mapping[str, Any]) -> ChatResult:
generations = []
for res in response["choices"]:
message = _convert_dict_to_message(res["message"])
gen = ChatGeneration(message=message)
generations.append(gen)
llm_output = {"token_usage": response["usage"], "model_name": self.model_name}
return ChatResult(generations=generations, llm_output=llm_output)

async def _agenerate(
self, messages: List[BaseMessage], stop: Optional[List[str]] = None
) -> ChatResult:
Expand Down Expand Up @@ -307,7 +306,7 @@ async def _agenerate(
response = await acompletion_with_retry(
self, messages=message_dicts, **params
)
return _create_chat_result(response)
return self._create_chat_result(response)

@property
def _identifying_params(self) -> Mapping[str, Any]:
Expand Down
19 changes: 19 additions & 0 deletions tests/integration_tests/chat_models/test_openai.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
"""Test ChatOpenAI wrapper."""


import pytest

from langchain.callbacks.base import CallbackManager
Expand Down Expand Up @@ -78,6 +79,24 @@ def test_chat_openai_streaming() -> None:
assert isinstance(response, BaseMessage)


def test_chat_openai_llm_output_contains_model_name() -> None:
"""Test llm_output contains model_name."""
chat = ChatOpenAI(max_tokens=10)
message = HumanMessage(content="Hello")
llm_result = chat.generate([[message]])
assert llm_result.llm_output is not None
assert llm_result.llm_output["model_name"] == chat.model_name


def test_chat_openai_streaming_llm_output_contains_model_name() -> None:
"""Test llm_output contains model_name."""
chat = ChatOpenAI(max_tokens=10, streaming=True)
message = HumanMessage(content="Hello")
llm_result = chat.generate([[message]])
assert llm_result.llm_output is not None
assert llm_result.llm_output["model_name"] == chat.model_name


def test_chat_openai_invalid_streaming_params() -> None:
"""Test that streaming correctly invokes on_llm_new_token callback."""
with pytest.raises(ValueError):
Expand Down

0 comments on commit 43e08fa

Please sign in to comment.