Add new beta StructuredPrompt

libs/core/langchain_core/prompts/structured.py

@@ -0,0 +1,130 @@
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    Type,
+    Union,
+)
+
+from langchain_core._api.beta_decorator import beta
+from langchain_core.language_models.base import BaseLanguageModel
+from langchain_core.messages import MessageLikeRepresentation
+from langchain_core.prompts.chat import (
+    BaseChatPromptTemplate,
+    BaseMessagePromptTemplate,
+    ChatPromptTemplate,
+    MessagesPlaceholder,
+    _convert_to_message,
+)
+from langchain_core.pydantic_v1 import BaseModel
+from langchain_core.runnables.base import (
+    Other,
+    Runnable,
+    RunnableSequence,
+    RunnableSerializable,
+)
+
+
+@beta()
+class StructuredPrompt(ChatPromptTemplate):
+    schema_: Union[Dict, Type[BaseModel]]
+
+    @classmethod
+    def from_messages_and_schema(
+        cls,
+        messages: Sequence[MessageLikeRepresentation],
+        schema: Union[Dict, Type[BaseModel]],
+    ) -> ChatPromptTemplate:
+        """Create a chat prompt template from a variety of message formats.
+
+        Examples:
+
+            Instantiation from a list of message templates:
+
+            .. code-block:: python
+
+                template = ChatPromptTemplate.from_messages([
+                    ("human", "Hello, how are you?"),
+                    ("ai", "I'm doing well, thanks!"),
+                    ("human", "That's good to hear."),
+                ])
+
+            Instantiation from mixed message formats:
+
+            .. code-block:: python
+
+                template = ChatPromptTemplate.from_messages([
+                    SystemMessage(content="hello"),
+                    ("human", "Hello, how are you?"),
+                ])
+
+        Args:
+            messages: sequence of message representations.
+                  A message can be represented using the following formats:
+                  (1) BaseMessagePromptTemplate, (2) BaseMessage, (3) 2-tuple of
+                  (message type, template); e.g., ("human", "{user_input}"),
+                  (4) 2-tuple of (message class, template), (4) a string which is
+                  shorthand for ("human", template); e.g., "{user_input}"
+
+        Returns:
+            a chat prompt template
+        """
+        _messages = [_convert_to_message(message) for message in messages]
+
+        # Automatically infer input variables from messages
+        input_vars: Set[str] = set()
+        partial_vars: Dict[str, Any] = {}
+        for _message in _messages:
+            if isinstance(_message, MessagesPlaceholder) and _message.optional:
+                partial_vars[_message.variable_name] = []
+            elif isinstance(
+                _message, (BaseChatPromptTemplate, BaseMessagePromptTemplate)
+            ):
+                input_vars.update(_message.input_variables)
+
+        return cls(
+            input_variables=sorted(input_vars),
+            messages=_messages,
+            partial_variables=partial_vars,
+            schema_=schema,
+        )
+
+    def __or__(
+        self,
+        other: Union[
+            Runnable[Any, Other],
+            Callable[[Any], Other],
+            Callable[[Iterator[Any]], Iterator[Other]],
+            Mapping[str, Union[Runnable[Any, Other], Callable[[Any], Other], Any]],
+        ],
+    ) -> RunnableSerializable[Dict, Other]:
+        if isinstance(other, BaseLanguageModel) or hasattr(
+            other, "with_structured_output"
+        ):
+            return RunnableSequence(self, other.with_structured_output(self.schema_))
+        else:
+            return super().__or__(other)
+
+    def pipe(
+        self,
+        *others: Union[Runnable[Any, Other], Callable[[Any], Other]],
+        name: Optional[str] = None,
+    ) -> RunnableSerializable[Dict, Other]:
+        if (
+            others
+            and isinstance(others[0], BaseLanguageModel)
+            or hasattr(others[0], "with_structured_output")
+        ):
+            return RunnableSequence(
+                self,
+                others[0].with_structured_output(self.schema_),
+                *others[1:],
+                name=name,
+            )
+        else:
+            return super().pipe(*others, name=name)

libs/core/tests/unit_tests/prompts/test_structured.py

@@ -0,0 +1,42 @@
+from ast import Dict
+from functools import partial
+from typing import Type, Union
+
+from langchain_core.prompts.structured import StructuredPrompt
+from langchain_core.pydantic_v1 import BaseModel
+from langchain_core.runnables.base import Runnable, RunnableLambda
+from tests.unit_tests.fake.chat_model import FakeListChatModel
+
+
+def _fake_runnable(schema: Type[BaseModel], _) -> BaseModel:
+    return schema(name="yo", value=42)
+
+
+class FakeStructuredChatModel(FakeListChatModel):
+    """Fake ChatModel for testing purposes."""
+
+    def with_structured_output(self, schema: Union[Dict, Type[BaseModel]]) -> Runnable:
+        return RunnableLambda(partial(_fake_runnable, schema))
+
+    @property
+    def _llm_type(self) -> str:
+        return "fake-messages-list-chat-model"
+
+
+def test_structured_prompt() -> None:
+    class OutputSchema(BaseModel):
+        name: str
+        value: int
+
+    prompt = StructuredPrompt.from_messages_and_schema(
+        [
+            ("human", "I'm very structured, how about you?"),
+        ],
+        OutputSchema,
+    )
+
+    model = FakeStructuredChatModel(responses=[])
+
+    chain = prompt | model
+
+    assert chain.invoke({"hello": "there"}) == OutputSchema(name="yo", value=42)