mirror of https://github.com/hwchase17/langchain
Nuno Campos
12 months ago
committed by
GitHub
parent
8c150ad7f6
commit
34ffb94770
@ -1,168 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
from typing import (
|
|
||||||
TYPE_CHECKING,
|
|
||||||
Any,
|
|
||||||
AsyncIterator,
|
|
||||||
Dict,
|
|
||||||
Iterator,
|
|
||||||
Mapping,
|
|
||||||
Optional,
|
|
||||||
Union,
|
|
||||||
)
|
|
||||||
|
|
||||||
from langchain.schema.runnable.base import Input, Other, Output, RunnableSerializable
|
|
||||||
from langchain.schema.runnable.config import RunnableConfig
|
|
||||||
from langchain.schema.runnable.passthrough import RunnablePassthrough
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from langchain.callbacks.manager import (
|
|
||||||
AsyncCallbackManagerForChainRun,
|
|
||||||
CallbackManagerForChainRun,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class PutLocalVar(RunnablePassthrough):
|
|
||||||
key: Union[str, Mapping[str, str]]
|
|
||||||
"""The key(s) to use for storing the input variable(s) in local state.
|
|
||||||
|
|
||||||
If a string is provided then the entire input is stored under that key. If a
|
|
||||||
Mapping is provided, then the map values are gotten from the input and
|
|
||||||
stored in local state under the map keys.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, key: Union[str, Mapping[str, str]], **kwargs: Any) -> None:
|
|
||||||
super().__init__(key=key, **kwargs)
|
|
||||||
|
|
||||||
def _concat_put(
|
|
||||||
self,
|
|
||||||
input: Other,
|
|
||||||
*,
|
|
||||||
config: Optional[RunnableConfig] = None,
|
|
||||||
replace: bool = False,
|
|
||||||
) -> None:
|
|
||||||
if config is None:
|
|
||||||
raise ValueError(
|
|
||||||
"PutLocalVar should only be used in a RunnableSequence, and should "
|
|
||||||
"therefore always receive a non-null config."
|
|
||||||
)
|
|
||||||
if isinstance(self.key, str):
|
|
||||||
if self.key not in config["locals"] or replace:
|
|
||||||
config["locals"][self.key] = input
|
|
||||||
else:
|
|
||||||
config["locals"][self.key] += input
|
|
||||||
elif isinstance(self.key, Mapping):
|
|
||||||
if not isinstance(input, Mapping):
|
|
||||||
raise TypeError(
|
|
||||||
f"Received key of type Mapping but input of type {type(input)}. "
|
|
||||||
f"input is expected to be of type Mapping when key is Mapping."
|
|
||||||
)
|
|
||||||
for input_key, put_key in self.key.items():
|
|
||||||
if put_key not in config["locals"] or replace:
|
|
||||||
config["locals"][put_key] = input[input_key]
|
|
||||||
else:
|
|
||||||
config["locals"][put_key] += input[input_key]
|
|
||||||
else:
|
|
||||||
raise TypeError(
|
|
||||||
f"`key` should be a string or Mapping[str, str], received type "
|
|
||||||
f"{(type(self.key))}."
|
|
||||||
)
|
|
||||||
|
|
||||||
def invoke(
|
|
||||||
self, input: Other, config: Optional[RunnableConfig] = None, **kwargs: Any
|
|
||||||
) -> Other:
|
|
||||||
self._concat_put(input, config=config, replace=True)
|
|
||||||
return super().invoke(input, config=config, **kwargs)
|
|
||||||
|
|
||||||
async def ainvoke(
|
|
||||||
self,
|
|
||||||
input: Other,
|
|
||||||
config: Optional[RunnableConfig] = None,
|
|
||||||
**kwargs: Optional[Any],
|
|
||||||
) -> Other:
|
|
||||||
self._concat_put(input, config=config, replace=True)
|
|
||||||
return await super().ainvoke(input, config=config, **kwargs)
|
|
||||||
|
|
||||||
def transform(
|
|
||||||
self,
|
|
||||||
input: Iterator[Other],
|
|
||||||
config: Optional[RunnableConfig] = None,
|
|
||||||
**kwargs: Optional[Any],
|
|
||||||
) -> Iterator[Other]:
|
|
||||||
for chunk in super().transform(input, config=config, **kwargs):
|
|
||||||
self._concat_put(chunk, config=config)
|
|
||||||
yield chunk
|
|
||||||
|
|
||||||
async def atransform(
|
|
||||||
self,
|
|
||||||
input: AsyncIterator[Other],
|
|
||||||
config: Optional[RunnableConfig] = None,
|
|
||||||
**kwargs: Optional[Any],
|
|
||||||
) -> AsyncIterator[Other]:
|
|
||||||
async for chunk in super().atransform(input, config=config, **kwargs):
|
|
||||||
self._concat_put(chunk, config=config)
|
|
||||||
yield chunk
|
|
||||||
|
|
||||||
|
|
||||||
class GetLocalVar(
|
|
||||||
RunnableSerializable[Input, Union[Output, Dict[str, Union[Input, Output]]]]
|
|
||||||
):
|
|
||||||
key: str
|
|
||||||
"""The key to extract from the local state."""
|
|
||||||
passthrough_key: Optional[str] = None
|
|
||||||
"""The key to use for passing through the invocation input.
|
|
||||||
|
|
||||||
If None, then only the value retrieved from local state is returned. Otherwise a
|
|
||||||
dictionary ``{self.key: <<retrieved_value>>, self.passthrough_key: <<input>>}``
|
|
||||||
is returned.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, key: str, **kwargs: Any) -> None:
|
|
||||||
super().__init__(key=key, **kwargs)
|
|
||||||
|
|
||||||
def _get(
|
|
||||||
self,
|
|
||||||
input: Input,
|
|
||||||
run_manager: Union[CallbackManagerForChainRun, Any],
|
|
||||||
config: RunnableConfig,
|
|
||||||
) -> Union[Output, Dict[str, Union[Input, Output]]]:
|
|
||||||
if self.passthrough_key:
|
|
||||||
return {
|
|
||||||
self.key: config["locals"][self.key],
|
|
||||||
self.passthrough_key: input,
|
|
||||||
}
|
|
||||||
else:
|
|
||||||
return config["locals"][self.key]
|
|
||||||
|
|
||||||
async def _aget(
|
|
||||||
self,
|
|
||||||
input: Input,
|
|
||||||
run_manager: AsyncCallbackManagerForChainRun,
|
|
||||||
config: RunnableConfig,
|
|
||||||
) -> Union[Output, Dict[str, Union[Input, Output]]]:
|
|
||||||
return self._get(input, run_manager, config)
|
|
||||||
|
|
||||||
def invoke(
|
|
||||||
self, input: Input, config: Optional[RunnableConfig] = None
|
|
||||||
) -> Union[Output, Dict[str, Union[Input, Output]]]:
|
|
||||||
if config is None:
|
|
||||||
raise ValueError(
|
|
||||||
"GetLocalVar should only be used in a RunnableSequence, and should "
|
|
||||||
"therefore always receive a non-null config."
|
|
||||||
)
|
|
||||||
|
|
||||||
return self._call_with_config(self._get, input, config)
|
|
||||||
|
|
||||||
async def ainvoke(
|
|
||||||
self,
|
|
||||||
input: Input,
|
|
||||||
config: Optional[RunnableConfig] = None,
|
|
||||||
**kwargs: Optional[Any],
|
|
||||||
) -> Union[Output, Dict[str, Union[Input, Output]]]:
|
|
||||||
if config is None:
|
|
||||||
raise ValueError(
|
|
||||||
"GetLocalVar should only be used in a RunnableSequence, and should "
|
|
||||||
"therefore always receive a non-null config."
|
|
||||||
)
|
|
||||||
|
|
||||||
return await self._acall_with_config(self._aget, input, config)
|
|
||||||
@ -1,94 +0,0 @@
|
|||||||
from typing import Any, Callable, Type
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from langchain.llms import FakeListLLM
|
|
||||||
from langchain.prompts import PromptTemplate
|
|
||||||
from langchain.schema.runnable import (
|
|
||||||
GetLocalVar,
|
|
||||||
PutLocalVar,
|
|
||||||
Runnable,
|
|
||||||
RunnablePassthrough,
|
|
||||||
RunnableSequence,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("method", "input", "output"),
|
|
||||||
[
|
|
||||||
(lambda r, x: r.invoke(x), "foo", "foo"),
|
|
||||||
(lambda r, x: r.batch(x), ["foo", "bar"], ["foo", "bar"]),
|
|
||||||
(lambda r, x: list(r.stream(x))[0], "foo", "foo"),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_put_get(method: Callable, input: Any, output: Any) -> None:
|
|
||||||
runnable = PutLocalVar("input") | GetLocalVar("input")
|
|
||||||
assert method(runnable, input) == output
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("method", "input", "output"),
|
|
||||||
[
|
|
||||||
(lambda r, x: r.ainvoke(x), "foo", "foo"),
|
|
||||||
(lambda r, x: r.abatch(x), ["foo", "bar"], ["foo", "bar"]),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
async def test_put_get_async(method: Callable, input: Any, output: Any) -> None:
|
|
||||||
runnable = PutLocalVar("input") | GetLocalVar("input")
|
|
||||||
assert await method(runnable, input) == output
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("runnable", "error"),
|
|
||||||
[
|
|
||||||
(PutLocalVar("input"), ValueError),
|
|
||||||
(GetLocalVar("input"), ValueError),
|
|
||||||
(PutLocalVar("input") | GetLocalVar("missing"), KeyError),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_incorrect_usage(runnable: RunnableSequence, error: Type[Exception]) -> None:
|
|
||||||
with pytest.raises(error):
|
|
||||||
runnable.invoke("foo")
|
|
||||||
|
|
||||||
|
|
||||||
def test_get_in_map() -> None:
|
|
||||||
runnable: Runnable = PutLocalVar("input") | {"bar": GetLocalVar("input")}
|
|
||||||
assert runnable.invoke("foo") == {"bar": "foo"}
|
|
||||||
|
|
||||||
|
|
||||||
def test_put_in_map() -> None:
|
|
||||||
runnable: Runnable = {"bar": PutLocalVar("input")} | GetLocalVar("input")
|
|
||||||
with pytest.raises(KeyError):
|
|
||||||
runnable.invoke("foo")
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"runnable",
|
|
||||||
[
|
|
||||||
PutLocalVar("input") | GetLocalVar("input", passthrough_key="output"),
|
|
||||||
(
|
|
||||||
PutLocalVar("input")
|
|
||||||
| {"input": RunnablePassthrough()}
|
|
||||||
| PromptTemplate.from_template("say {input}")
|
|
||||||
| FakeListLLM(responses=["hello"])
|
|
||||||
| GetLocalVar("input", passthrough_key="output")
|
|
||||||
),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("method", "input", "output"),
|
|
||||||
[
|
|
||||||
(lambda r, x: r.invoke(x), "hello", {"input": "hello", "output": "hello"}),
|
|
||||||
(lambda r, x: r.batch(x), ["hello"], [{"input": "hello", "output": "hello"}]),
|
|
||||||
(
|
|
||||||
lambda r, x: list(r.stream(x))[0],
|
|
||||||
"hello",
|
|
||||||
{"input": "hello", "output": "hello"},
|
|
||||||
),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_put_get_sequence(
|
|
||||||
runnable: RunnableSequence, method: Callable, input: Any, output: Any
|
|
||||||
) -> None:
|
|
||||||
assert method(runnable, input) == output
|
|
||||||
Loading…
Reference in New Issue