mirror of
https://github.com/kennethreitz/langchain.git
synced 2026-06-05 23:00:18 +00:00
Zander Chase
b0859c9b18
Handle the new retriever events in a way that (I think) is entirely
backwards compatible? Needs more testing for some of the chain changes
and all.
This creates an entire new run type, however. We could also just treat
this as an event within a chain run presumably (same with memory)
Adds a subclass initializer that upgrades old retriever implementations
to the new schema, along with tests to ensure they work.
First commit doesn't upgrade any of our retriever implementations (to
show that we can pass the tests along with additional ones testing the
upgrade logic).
Second commit upgrades the known universe of retrievers in langchain.
- [X] Add callback handling methods for retriever start/end/error (open
to renaming to 'retrieval' if you want that)
- [X] Update BaseRetriever schema to support callbacks
- [X] Tests for upgrading old "v1" retrievers for backwards
compatibility
- [X] Update existing retriever implementations to implement the new
interface
- [X] Update calls within chains to .{a]get_relevant_documents to pass
the child callback manager
- [X] Update the notebooks/docs to reflect the new interface
- [X] Test notebooks thoroughly
Not handled:
- Memory pass throughs: retrieval memory doesn't have a parent callback
manager passed through the method
---------
Co-authored-by: Nuno Campos <nuno@boringbits.io>
Co-authored-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
375 lines
8.3 KiB
Python
375 lines
8.3 KiB
Python
"""A fake callback handler for testing purposes."""
|
|
from itertools import chain
|
|
from typing import Any, Dict, List, Optional
|
|
from uuid import UUID
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
|
|
from langchain.schema import BaseMessage
|
|
|
|
|
|
class BaseFakeCallbackHandler(BaseModel):
|
|
"""Base fake callback handler for testing."""
|
|
|
|
starts: int = 0
|
|
ends: int = 0
|
|
errors: int = 0
|
|
text: int = 0
|
|
ignore_llm_: bool = False
|
|
ignore_chain_: bool = False
|
|
ignore_agent_: bool = False
|
|
ignore_retriever_: bool = False
|
|
ignore_chat_model_: bool = False
|
|
|
|
# add finer-grained counters for easier debugging of failing tests
|
|
chain_starts: int = 0
|
|
chain_ends: int = 0
|
|
llm_starts: int = 0
|
|
llm_ends: int = 0
|
|
llm_streams: int = 0
|
|
tool_starts: int = 0
|
|
tool_ends: int = 0
|
|
agent_actions: int = 0
|
|
agent_ends: int = 0
|
|
chat_model_starts: int = 0
|
|
retriever_starts: int = 0
|
|
retriever_ends: int = 0
|
|
retriever_errors: int = 0
|
|
|
|
|
|
class BaseFakeCallbackHandlerMixin(BaseFakeCallbackHandler):
|
|
"""Base fake callback handler mixin for testing."""
|
|
|
|
def on_llm_start_common(self) -> None:
|
|
self.llm_starts += 1
|
|
self.starts += 1
|
|
|
|
def on_llm_end_common(self) -> None:
|
|
self.llm_ends += 1
|
|
self.ends += 1
|
|
|
|
def on_llm_error_common(self) -> None:
|
|
self.errors += 1
|
|
|
|
def on_llm_new_token_common(self) -> None:
|
|
self.llm_streams += 1
|
|
|
|
def on_chain_start_common(self) -> None:
|
|
("CHAIN START")
|
|
self.chain_starts += 1
|
|
self.starts += 1
|
|
|
|
def on_chain_end_common(self) -> None:
|
|
self.chain_ends += 1
|
|
self.ends += 1
|
|
|
|
def on_chain_error_common(self) -> None:
|
|
self.errors += 1
|
|
|
|
def on_tool_start_common(self) -> None:
|
|
self.tool_starts += 1
|
|
self.starts += 1
|
|
|
|
def on_tool_end_common(self) -> None:
|
|
self.tool_ends += 1
|
|
self.ends += 1
|
|
|
|
def on_tool_error_common(self) -> None:
|
|
self.errors += 1
|
|
|
|
def on_agent_action_common(self) -> None:
|
|
print("AGENT ACTION")
|
|
self.agent_actions += 1
|
|
self.starts += 1
|
|
|
|
def on_agent_finish_common(self) -> None:
|
|
self.agent_ends += 1
|
|
self.ends += 1
|
|
|
|
def on_chat_model_start_common(self) -> None:
|
|
print("STARTING CHAT MODEL")
|
|
self.chat_model_starts += 1
|
|
self.starts += 1
|
|
|
|
def on_text_common(self) -> None:
|
|
self.text += 1
|
|
|
|
def on_retriever_start_common(self) -> None:
|
|
self.starts += 1
|
|
self.retriever_starts += 1
|
|
|
|
def on_retriever_end_common(self) -> None:
|
|
self.ends += 1
|
|
self.retriever_ends += 1
|
|
|
|
def on_retriever_error_common(self) -> None:
|
|
self.errors += 1
|
|
self.retriever_errors += 1
|
|
|
|
|
|
class FakeCallbackHandler(BaseCallbackHandler, BaseFakeCallbackHandlerMixin):
|
|
"""Fake callback handler for testing."""
|
|
|
|
@property
|
|
def ignore_llm(self) -> bool:
|
|
"""Whether to ignore LLM callbacks."""
|
|
return self.ignore_llm_
|
|
|
|
@property
|
|
def ignore_chain(self) -> bool:
|
|
"""Whether to ignore chain callbacks."""
|
|
return self.ignore_chain_
|
|
|
|
@property
|
|
def ignore_agent(self) -> bool:
|
|
"""Whether to ignore agent callbacks."""
|
|
return self.ignore_agent_
|
|
|
|
@property
|
|
def ignore_retriever(self) -> bool:
|
|
"""Whether to ignore retriever callbacks."""
|
|
return self.ignore_retriever_
|
|
|
|
def on_llm_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_llm_start_common()
|
|
|
|
def on_llm_new_token(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_llm_new_token_common()
|
|
|
|
def on_llm_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_llm_end_common()
|
|
|
|
def on_llm_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_llm_error_common()
|
|
|
|
def on_chain_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_chain_start_common()
|
|
|
|
def on_chain_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_chain_end_common()
|
|
|
|
def on_chain_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_chain_error_common()
|
|
|
|
def on_tool_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_tool_start_common()
|
|
|
|
def on_tool_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_tool_end_common()
|
|
|
|
def on_tool_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_tool_error_common()
|
|
|
|
def on_agent_action(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_agent_action_common()
|
|
|
|
def on_agent_finish(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_agent_finish_common()
|
|
|
|
def on_text(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_text_common()
|
|
|
|
def on_retriever_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_retriever_start_common()
|
|
|
|
def on_retriever_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_retriever_end_common()
|
|
|
|
def on_retriever_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
self.on_retriever_error_common()
|
|
|
|
def __deepcopy__(self, memo: dict) -> "FakeCallbackHandler":
|
|
return self
|
|
|
|
|
|
class FakeCallbackHandlerWithChatStart(FakeCallbackHandler):
|
|
def on_chat_model_start(
|
|
self,
|
|
serialized: Dict[str, Any],
|
|
messages: List[List[BaseMessage]],
|
|
*,
|
|
run_id: UUID,
|
|
parent_run_id: Optional[UUID] = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
assert all(isinstance(m, BaseMessage) for m in chain(*messages))
|
|
self.on_chat_model_start_common()
|
|
|
|
|
|
class FakeAsyncCallbackHandler(AsyncCallbackHandler, BaseFakeCallbackHandlerMixin):
|
|
"""Fake async callback handler for testing."""
|
|
|
|
@property
|
|
def ignore_llm(self) -> bool:
|
|
"""Whether to ignore LLM callbacks."""
|
|
return self.ignore_llm_
|
|
|
|
@property
|
|
def ignore_chain(self) -> bool:
|
|
"""Whether to ignore chain callbacks."""
|
|
return self.ignore_chain_
|
|
|
|
@property
|
|
def ignore_agent(self) -> bool:
|
|
"""Whether to ignore agent callbacks."""
|
|
return self.ignore_agent_
|
|
|
|
async def on_llm_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_llm_start_common()
|
|
|
|
async def on_llm_new_token(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_llm_new_token_common()
|
|
|
|
async def on_llm_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_llm_end_common()
|
|
|
|
async def on_llm_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_llm_error_common()
|
|
|
|
async def on_chain_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_chain_start_common()
|
|
|
|
async def on_chain_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_chain_end_common()
|
|
|
|
async def on_chain_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_chain_error_common()
|
|
|
|
async def on_tool_start(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_tool_start_common()
|
|
|
|
async def on_tool_end(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_tool_end_common()
|
|
|
|
async def on_tool_error(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_tool_error_common()
|
|
|
|
async def on_agent_action(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_agent_action_common()
|
|
|
|
async def on_agent_finish(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_agent_finish_common()
|
|
|
|
async def on_text(
|
|
self,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
self.on_text_common()
|
|
|
|
def __deepcopy__(self, memo: dict) -> "FakeAsyncCallbackHandler":
|
|
return self
|