forked from Archives/langchain
4ee47926ca
### Add on_chat_message_start to callback manager and base tracer Goal: trace messages directly to permit reloading as chat messages (store in an integration-agnostic way) Add an `on_chat_message_start` method. Fall back to `on_llm_start()` for handlers that don't have it implemented. Does so in a non-backwards-compat breaking way (for now)
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])
|