如何在LangChain中为LLM传递枚举值的附加信息

我在Python里用LangChain调用GPT-4这类大模型提取结构化信息，目标是用标签给公司分类。我定义了基于BaseModel的输出类：

from langchain_core.pydantic_v1 import BaseModel

class Company(BaseModel):
    industry: list[Industry]
    customer: list[Customer]

目前功能正常，但有些标签比较特殊，我想给大模型传递更多信息辅助标签选择。于是用aenum库的Enum给枚举值加了文档字符串：

from aenum import Enum

class Industry(Enum):
    _init_ = 'value __doc__'
    it = "Information Technology", "All kinds of computer stuff"
    agriculture = "Agriculture", "Farming, irrigation, fertilizers etc."

class Customer(Enum):
    _init_ = 'value __doc__'
    B2C = "B2C", "Companies selling directly to consumers"
    B2B = "B2B", "Companies selling to other businesses"

现在枚举值带了说明，但没法直接把这些文档字符串传给大模型。用.with_structured_output()或PydanticOutputParser时，枚举成员的文档字符串不会被包含进去：

from langchain_core.output_parsers import PydanticOutputParser

parser = PydanticOutputParser(pydantic_object=Company)

parser.get_format_instructions()
# 输出的格式说明里，枚举的描述只有"An enumeration."，看不到每个成员的文档字符串
# 'The output should be formatted as a JSON instance that conforms to the JSON schema below.
# As an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}
# the object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.
# Here is the output schema:
# ```
# {"properties": {"industry": {"type": "array", "items": {"$ref": "#/definitions/Industry"}}, "customer": {"type": "array", "items": {"$ref": "#/definitions/Customer"}}}, "required": ["industry", "customer"], "definitions": {"Industry": {"title": "Industry", "description": "An enumeration.", "enum": ["Information Technology", "Agriculture"]}, "Customer": {"title": "Customer", "description": "An enumeration.", "enum": ["B2C", "B2B"]}}}
#```'

我可以写自定义提示词来传递这些文档字符串，但想知道有没有更直接的实现方式。

解决方案 #

方法1：自定义枚举类的JSON Schema生成 #

LangChain的PydanticOutputParser依赖Pydantic默认的JSON Schema生成逻辑，默认不会包含枚举成员的文档字符串。可以通过重写枚举类的__json_schema__方法，把成员说明注入到schema中：

from aenum import Enum
from langchain_core.pydantic_v1 import BaseModel

class DocumentedEnum(Enum):
    _init_ = 'value __doc__'

    @classmethod
    def __json_schema__(cls):
        schema = super().__json_schema__()
        # 拼接所有枚举成员的说明文本
        enum_options = "\n".join([f"- {member.value}: {member.__doc__}" for member in cls])
        # 更新schema的描述字段
        schema["description"] = f"枚举选项说明:\n{enum_options}"
        return schema

# 继承自定义的DocumentedEnum
class Industry(DocumentedEnum):
    it = "Information Technology", "All kinds of computer stuff"
    agriculture = "Agriculture", "Farming, irrigation, fertilizers etc."

class Customer(DocumentedEnum):
    B2C = "B2C", "Companies selling directly to consumers"
    B2B = "B2B", "Companies selling to other businesses"

class Company(BaseModel):
    industry: list[Industry]
    customer: list[Customer]

现在再生成格式说明，枚举的描述会包含所有成员的文档字符串：

parser = PydanticOutputParser(pydantic_object=Company)
print(parser.get_format_instructions())

方法2：自定义格式说明模板 #

如果不想修改枚举类的继承关系，可以手动修改PydanticOutputParser生成的格式说明，注入枚举成员信息：

from langchain_core.output_parsers import PydanticOutputParser

# 提取枚举成员说明的工具函数
def extract_enum_docs(enum_cls):
    return "\n".join([f"- {member.value}: {member.__doc__}" for member in enum_cls])

parser = PydanticOutputParser(pydantic_object=Company)
original_instructions = parser.get_format_instructions()

# 替换原始说明中的枚举描述
updated_instructions = original_instructions.replace(
    '"description": "An enumeration."',
    f'"description": "枚举选项说明:\n{extract_enum_docs(Industry)}"'
).replace(
    '"description": "An enumeration."',
    f'"description": "枚举选项说明:\n{extract_enum_docs(Customer)}"'
)

print(updated_instructions)

方法3：使用LangChain官方枚举类型（推荐） #

LangChain提供了原生的Enum类型支持，会自动把枚举成员的文档字符串包含到格式说明中，无需依赖第三方库：

from langchain_core.pydantic_v1 import BaseModel
from langchain_core.enums import Enum as LangChainEnum

class Industry(LangChainEnum):
    it = "Information Technology"
    agriculture = "Agriculture"

    # 直接给成员添加文档字符串
    it.__doc__ = "All kinds of computer stuff"
    agriculture.__doc__ = "Farming, irrigation, fertilizers etc."

class Customer(LangChainEnum):
    B2C = "B2C"
    B2B = "B2B"

    B2C.__doc__ = "Companies selling directly to consumers"
    B2B.__doc__ = "Companies selling to other businesses"

class Company(BaseModel):
    industry: list[Industry]
    customer: list[Customer]

使用with_structured_output()调用模型时，大模型会自动收到包含枚举成员说明的格式要求：

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o")
structured_llm = llm.with_structured_output(Company)

# 示例调用
result = structured_llm.invoke("一家开发云服务器并向科技创业公司销售的企业")
print(result)

内容的提问来源于stack exchange，提问作者antti