mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
|
"""Fake Chat Model wrapper for testing purposes."""
|
||
|
|
from typing import List, Optional
|
||
|
|
|
||
|
|
from langchain.callbacks.manager import (
|
||
|
|
AsyncCallbackManagerForLLMRun,
|
||
|
|
CallbackManagerForLLMRun,
|
||
|
|
)
|
||
|
|
from langchain.chat_models.base import SimpleChatModel
|
||
|
|
from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatResult
|
||
|
|
|
||
|
|
|
||
|
|
class FakeChatModel(SimpleChatModel):
|
||
|
|
"""Fake Chat Model wrapper for testing purposes."""
|
||
|
|
|
||
|
|
def _call(
|
||
|
|
self,
|
||
|
|
messages: List[BaseMessage],
|
||
|
|
stop: Optional[List[str]] = None,
|
||
|
|
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
||
|
|
) -> str:
|
||
|
|
return "fake response"
|
||
|
|
|
||
|
|
async def _agenerate(
|
||
|
|
self,
|
||
|
|
messages: List[BaseMessage],
|
||
|
|
stop: Optional[List[str]] = None,
|
||
|
|
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
|
||
|
|
) -> ChatResult:
|
||
|
|
output_str = "fake response"
|
||
|
|
message = AIMessage(content=output_str)
|
||
|
|
generation = ChatGeneration(message=message)
|
||
|
|
return ChatResult(generations=[generation])
|