基于AI 代理的PDF文档总结
本文的目的是使用 AI 代理，以尽可能低的成本总结任何 PDF 文档。
AI 代理因其能够处理大量数据、做出明智的决策并自主执行任务而对推理和行动有用。以下是可能使用它们的原因：
- 自动化决策——AI 代理可以分析复杂场景，识别模式，并在没有人工干预的情况下做出最佳决策。
- 可扩展性——它们可以高效地处理大规模问题，同时处理多个请求。
- 适应性和学习能力——先进的 AI 代理，特别是那些具有强化学习或微调 LLM 的，可以适应新数据，随着时间的推移提高其推理能力。
- 多步骤任务执行——它们可以分解复杂的任务，计划行动并系统地执行它们。
- 与 API 和系统的交互——AI 代理可以集成到外部工具中，检索数据并自动执行工作流程。
- 增强生产力——它们减少了手动工作量，使人类有更多时间从事战略性工作。
1、目标
- 使用 AI 代理，以尽可能低的成本总结任何 PDF 文档。因此，如果 PDF 是纯文本，我们可以使用本地工具进行总结，或者如果包含图表等更复杂的内容，则使用 LLM 如 Gemini Gemini。
- 通常，我们会使用 AssistantAgent 或其他由 LLM 支持的东西来做出决定。然而，在这种情况下，我们将使用一个自定义代理，不使用 LLM 来根据 PDF 内容的一些自定义逻辑做出决定，从而限制初始 LLM 使用成本。
2、系统架构
关键组件：
- AutoGen (v0.4) —— AI 代理框架。为了这次演示我使用了 AutoGen，但也可以使用其他许多替代方案，因为这个演示不依赖于任何特定的原生 AutoGen 功能。
- Google Gemini 1.5 —— 用于图像转文本摘要的 Google LLM（付费选项）
- MarkItDown —— 用于纯文本摘要的 Python 库（免费选项）
代理架构：
3、代码实现
定义代理：
def get_filename_without_extension(file_path):
return os.path.splitext(os.path.basename(file_path))[0]
# 定义代理
class UserProxyAgent(BaseChatAgent):
def __init__(self, name: str) -> None:
super().__init__(name, "A human user participating in the chat.")
@property
def produced_message_types(self) -> List[type[ChatMessage]]:
return [TextMessage]
async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
user_input = await asyncio.get_event_loop().run_in_executor(None, input, "Enter source PDF document Path:")
print("input:",user_input)
return Response(chat_message=TextMessage(content=user_input, source=self.name))
async def on_reset(self, cancellation_token: CancellationToken) -> None:
print("User Proxy Reset")
# 自定义 UserProxyAgent 以与 Gemini 交互
class GeminiAgent(BaseChatAgent):
def __init__(self, name: str, output_dir: str, mime_type : str = None):
super().__init__(name=name, description="An agent that converts images to text.")
self._output_dir = output_dir
self._mime_type = mime_type
@property
def produced_message_types(self) -> Sequence[type[ChatMessage]]:
return (TextMessage,)
async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
# 使用 Google Cloud SDK 调用 Gemini
if not messages:
return Response(chat_message=StopMessage(content="No File Name received!", source=self.name))
last_message = messages[-1].content.strip()
filePath = last_message
print("file:",filePath)
response = await self.process(filePath=filePath, mime_type=self._mime_type)
# 返回从 Gemini 生成的响应
final_msg = TextMessage(content=f"created:{filePath}", source=self.name)
return Response(chat_message=final_msg)
async def on_reset(self, cancellation_token: CancellationToken) -> None:
pass
async def process(self, filePath: str, mime_type: str = None) -> str:
try:
model = GenerativeModel("gemini-1.5-flash-002")
if(mime_type == "application/pdf"):
print("summarise pdf...")
encoded=self.encode_pdf(filePath)
#print("encoded", base64.b64decode(encoded))
attachment_file = Part.from_data( mime_type="application/pdf",
data=encoded)
prompt = "Can you summarise content including images from the attached pdf in markdown?"
else:
print("summarise image...")
attachment_file = Part.from_image(Image.load_from_file(filePath))
prompt = "Describe this image?"
# 查询模型
print("total_tokens: ", model.count_tokens([attachment_file, prompt]))
response = model.generate_content([attachment_file, prompt])
print("usage_metadata: ", response.usage_metadata)
filename = os.path.join(self._output_dir, get_filename_without_extension(filePath))+"_gemini.md"
with open(filename, "w", encoding="utf-8") as file:
file.write(response.text)
return f"created:{filename}"
except Exception as e:
print("Exception : ", traceback.format_exc())
return f"Error calling Vertex AI: {str(e)}"
def encode_pdf(self, pdf_path):
return Path(pdf_path).read_bytes()
class MarkItDownAgent(BaseChatAgent):
def __init__(self, name, output_dir: str):
super().__init__(name, "An agent to convert to Mark Down output.")
self._output_dir = output_dir
@property
def produced_message_types(self) -> Sequence[type[ChatMessage]]:
return (TextMessage,)
async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
if not messages:
return Response(chat_message=StopMessage(content="No File Name received!", source=self.name))
last_message = messages[-1].content.strip()
filePath = last_message
return await self.process(filePath)
async def process(self, filePath: str) -> Response:
inner_messages: List[ChatMessage] = []
md = MarkItDown()
inner_messages.append(TextMessage(content=f'Will Summarise Free option using MarkDown...', source=self.name))
result = md.convert(filePath)
filename = os.path.join(self._output_dir, get_filename_without_extension(filePath))+"_markitdown.md"
print("content:",result.text_content)
with open(filename, "w", encoding="utf-8") as file:
file.write(result.text_content)
#print(result.text_content)
final_msg = TextMessage(content=f"created:{filename}", source=self.name)
inner_messages.append(final_msg)
return Response(chat_message=final_msg)
async def on_reset(self, cancellation_token: CancellationToken) -> None:
pass
class DocumentTriageJobAgent(BaseChatAgent):
def __init__(self, name: str):
super().__init__(name, "An agent to decide on how which agent to use for processing document.")
@property
def produced_message_types(self) -> Sequence[type[ChatMessage]]:
return (TextMessage,)
async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
if not messages:
return Response(chat_message=StopMessage(content="No File Name received!", source=self.name))
last_message = messages[-1].content.strip()
filePath = last_message
return await self.process(filePath)
async def process(self, filePath: str) -> Response:
inner_messages: List[ChatMessage] = []
if not(os.path.isfile(filePath)):
return Response(chat_message=StopMessage(content="File path provided isn't a valid file!", source=self.name), inner_messages=inner_messages)
classifier_result = self.classifier(filePath)
inner_messages.append(TextMessage(content=f'pdf classifier_result: {classifier_result}', source=self.name))
print("document type:",classifier_result)
if(classifier_result=="TEXT"):
agent = MarkItDownAgent(name="markdown_agent", output_dir=output_dir)
elif (classifier_result=="IMAGE"):
agent = GeminiAgent(name="llm_agent", output_dir=output_dir, mime_type="application/pdf")
await Console(
agent.on_messages_stream(
[TextMessage(content=filePath, source="user")], CancellationToken()
)
)
#print("final:",inner_messages)
final_msg = TextMessage(content="Done!", source=self.name)
return Response(chat_message=final_msg, inner_messages=inner_messages)
async def on_reset(self, cancellation_token: CancellationToken) -> None:
pass
def classifier(self, pdf_file):
with open(pdf_file,"rb") as f:
pdf = fitz.open(f)
res = []
for page in pdf:
img_refs = page.get_image_info(xrefs=True)
if img_refs != []:
#print("Page", page.number, "images:", [i["xref"] for i in img_refs])
res.append([i["xref"] for i in img_refs])
if len(res) == 0:
return("TEXT")
else:
return("IMAGE")
运行代理：
# 初始化代理
async def main():
# 初始化代理
user_proxy_agent = UserProxyAgent(name="user_proxy_agent")
document_triage_agent = DocumentTriageJobAgent(name="document_triage_agent")
markitdown_agent = MarkItDownAgent(name="markdown_agent", output_dir=output_dir)
llm_agent = GeminiAgent(name="llm_agent", output_dir=output_dir, mime_type="application/pdf")
# 定义终止条件
termination = TextMentionTermination("Done!") | HandoffTermination(user_proxy_agent) | StopMessageTermination()
# 创建聊天组
group_chat = RoundRobinGroupChat(
participants=[user_proxy_agent, document_triage_agent],
termination_condition=termination
)
# 程序化提供初始消息
initial_message = TextMessage(content="Enter Source file Name:", source="user")
# 开始聊天
stream = group_chat.run_stream(cancellation_token=CancellationToken())
# 等待流处理消息
await Console(stream)
nest_asyncio.apply()
asyncio.run(main())
4、测试 PDF
通过 AI 代理运行了两个 PDF 文件，一个是纯文本，另一个包含图表作为图像。
纯文本摘要示例
带图像的文本摘要示例
5、结束语
- 使用 MarkItDown 时，PDF 的摘要免费，而 LLM 则根据文档内容的大小收费。
- Google Gemini 1.5 Flash 在文档摘要方面表现良好，尤其是在包含图像内容时。最初，我打算使用 Google 的 DocAI，但 Gemini 更便宜且更容易使用。预计 Google Gemini 2.0 Pro 将会更好。
汇智网翻译整理，转载请标明出处
