Add Text2Text Assistant as LLM for LangChain

text2text/langchain/test_text2text_assistant.py

@@ -0,0 +1,10 @@
+import pytest
+
+from text2text.langchain.text2text_assistant import Text2TextAssistant
+
+@pytest.mark.requires("langchain")
+def test_llm_inference() -> None:
+    input_text = 'Say "hello, world" back to me'
+    llm = Text2TextAssistant()
+    result = llm(input_text)
+    assert "hello" in result.lower() 

text2text/langchain/text2text_assistant.py

@@ -0,0 +1,28 @@
+from typing import Any, List, Mapping, Optional
+
+import text2text as t2t
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.base import LLM
+
+class Text2TextAssistant(LLM):
+    model: t2t.Assistant = t2t.Assistant()
+
+    @property
+    def _llm_type(self) -> str:
+        return "Text2Text"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs
+    ) -> str:
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+        return self.model.transform([prompt], **kwargs)[0]
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {"type": self._llm_type}