mirror of
https://github.com/hwchase17/langchain
synced 2024-11-06 03:20:49 +00:00
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
"""Fake Chat Model wrapper for testing purposes."""
|
|
from typing import Any, List, Mapping, 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,
|
|
**kwargs: Any,
|
|
) -> str:
|
|
return "fake response"
|
|
|
|
async def _agenerate(
|
|
self,
|
|
messages: List[BaseMessage],
|
|
stop: Optional[List[str]] = None,
|
|
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
|
|
**kwargs: Any,
|
|
) -> ChatResult:
|
|
output_str = "fake response"
|
|
message = AIMessage(content=output_str)
|
|
generation = ChatGeneration(message=message)
|
|
return ChatResult(generations=[generation])
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
return "fake-chat-model"
|
|
|
|
@property
|
|
def _identifying_params(self) -> Mapping[str, Any]:
|
|
return {"key": "fake"}
|