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