Merge pull request #5 from hwchase17/harrison/chains

add initial chains

Author: hwchase17

langchain/chains/__init__.py

@@ -0,0 +1 @@
+"""Chains are easily reusable components which can be linked together."""

langchain/chains/base.py

@@ -0,0 +1,41 @@
+"""Base interface that all chains should implement."""
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+
+class Chain(ABC):
+    """Base interface that all chains should implement."""
+
+    @property
+    @abstractmethod
+    def input_keys(self) -> List[str]:
+        """Input keys this chain expects."""
+
+    @property
+    @abstractmethod
+    def output_keys(self) -> List[str]:
+        """Output keys this chain expects."""
+
+    def _validate_inputs(self, inputs: Dict[str, str]) -> None:
+        """Check that all inputs are present."""
+        missing_keys = set(self.input_keys).difference(inputs)
+        if missing_keys:
+            raise ValueError(f"Missing some input keys: {missing_keys}")
+
+    def _validate_outputs(self, outputs: Dict[str, str]) -> None:
+        if set(outputs) != set(self.output_keys):
+            raise ValueError(
+                f"Did not get output keys that were expected. "
+                f"Got: {set(outputs)}. Expected: {set(self.output_keys)}."
+            )
+
+    @abstractmethod
+    def _run(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        """Run the logic of this chain and return the output."""
+
+    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+        """Run the logic of this chain and add to output."""
+        self._validate_inputs(inputs)
+        outputs = self._run(inputs)
+        self._validate_outputs(outputs)
+        return {**inputs, **outputs}

langchain/chains/llm.py

@@ -0,0 +1,46 @@
+"""Chain that just formats a prompt and calls an LLM."""
+from typing import Any, Dict, List
+
+from pydantic import BaseModel, Extra
+
+from langchain.chains.base import Chain
+from langchain.llms.base import LLM
+from langchain.prompt import Prompt
+
+
+class LLMChain(Chain, BaseModel):
+    """Chain to run queries against LLMs."""
+
+    prompt: Prompt
+    llm: LLM
+    return_key: str = "text"
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Will be whatever keys the prompt expects."""
+        return self.prompt.input_variables
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Will always return text key."""
+        return [self.return_key]
+
+    def _run(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+        selected_inputs = {k: inputs[k] for k in self.prompt.input_variables}
+        prompt = self.prompt.template.format(**selected_inputs)
+
+        kwargs = {}
+        if "stop" in inputs:
+            kwargs["stop"] = inputs["stop"]
+        response = self.llm(prompt, **kwargs)
+        return {self.return_key: response}
+
+    def predict(self, **kwargs: Any) -> str:
+        """More user-friendly interface for interacting with LLMs."""
+        return self(kwargs)[self.return_key]

tests/unit_tests/chains/__init__.py

@@ -0,0 +1 @@
+"""Tests for correct functioning of chains."""

tests/unit_tests/chains/test_base.py

@@ -0,0 +1,50 @@
+"""Test logic on base chain class."""
+from typing import Dict, List
+
+import pytest
+from pydantic import BaseModel
+
+from langchain.chains.base import Chain
+
+
+class FakeChain(Chain, BaseModel):
+    """Fake chain class for testing purposes."""
+
+    be_correct: bool = True
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Input key of foo."""
+        return ["foo"]
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Output key of bar."""
+        return ["bar"]
+
+    def _run(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        if self.be_correct:
+            return {"bar": "baz"}
+        else:
+            return {"baz": "bar"}
+
+
+def test_bad_inputs() -> None:
+    """Test errors are raised if input keys are not found."""
+    chain = FakeChain()
+    with pytest.raises(ValueError):
+        chain({"foobar": "baz"})
+
+
+def test_bad_outputs() -> None:
+    """Test errors are raised if outputs keys are not found."""
+    chain = FakeChain(be_correct=False)
+    with pytest.raises(ValueError):
+        chain({"foo": "baz"})
+
+
+def test_correct_call() -> None:
+    """Test correct call of fake chain."""
+    chain = FakeChain()
+    output = chain({"foo": "bar"})
+    assert output == {"foo": "bar", "bar": "baz"}

tests/unit_tests/chains/test_llm.py

@@ -0,0 +1,36 @@
+"""Test LLM chain."""
+import pytest
+
+from langchain.chains.llm import LLMChain
+from langchain.prompt import Prompt
+from tests.unit_tests.llms.fake_llm import FakeLLM
+
+
+@pytest.fixture
+def fake_llm_chain() -> LLMChain:
+    """Fake LLM chain for testing purposes."""
+    prompt = Prompt(input_variables=["bar"], template="This is a {bar}:")
+    return LLMChain(prompt=prompt, llm=FakeLLM(), return_key="text1")
+
+
+def test_missing_inputs(fake_llm_chain: LLMChain) -> None:
+    """Test error is raised if inputs are missing."""
+    with pytest.raises(ValueError):
+        fake_llm_chain({"foo": "bar"})
+
+
+def test_valid_call(fake_llm_chain: LLMChain) -> None:
+    """Test valid call of LLM chain."""
+    output = fake_llm_chain({"bar": "baz"})
+    assert output == {"bar": "baz", "text1": "foo"}
+
+    # Test with stop words.
+    output = fake_llm_chain({"bar": "baz", "stop": ["foo"]})
+    # Response should be `bar` now.
+    assert output == {"bar": "baz", "stop": ["foo"], "text1": "bar"}
+
+
+def test_predict_method(fake_llm_chain: LLMChain) -> None:
+    """Test predict method works."""
+    output = fake_llm_chain.predict(bar="baz")
+    assert output == "foo"

tests/unit_tests/llms/fake_llm.py

@@ -0,0 +1,15 @@
+"""Fake LLM wrapper for testing purposes."""
+from typing import List, Optional
+
+from langchain.llms.base import LLM
+
+
+class FakeLLM(LLM):
+    """Fake LLM wrapper for testing purposes."""
+
+    def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        """Return `foo` if no stop words, otherwise `bar`."""
+        if stop is None:
+            return "foo"
+        else:
+            return "bar"