from document_ai_agents.document_utils import extract_images_from_pdf
from document_ai_agents.image_utils import pil_image_to_base64_jpeg
from pathlib import Path

class DocumentParsingAgent:
    @classmethod
    def get_images(cls, state):
        """
        提取一個PDF的頁面為Base64編碼的JPEG圖像。
        """
        assert Path(state.document_path).is_file(), "File does not exist"
        # 從PDF中提取圖像
        images = extract_images_from_pdf(state.document_path)
        assert images, "No images extracted"
        # 轉換圖像到Base64編碼的JPEG
        pages_as_base64_jpeg_images = [pil_image_to_base64_jpeg(x) for x in images]
        return {"pages_as_base64_jpeg_images": pages_as_base64_jpeg_images}
extract_images_from_pdf：將PDF的每一頁提取為PIL圖像。

pil_image_to_base64_jpeg：將圖像轉換為Base64編碼的JPEG格式。

from pydantic import BaseModel, Field
from typing import Literal
import json
import google.generativeai as genai
from langchain_core.documents import Document

class DetectedLayoutItem(BaseModel):
    """
    針對頁面上檢測到的每個布局元素的架構。
    """
    element_type: Literal["Table", "Figure", "Image", "Text-block"] = Field(
        ..., 
        description="Type of detected item. Examples: Table, Figure, Image, Text-block."
    )
    summary: str = Field(..., description="A detailed description of the layout item.")

class LayoutElements(BaseModel):
    """
    針對頁面上的布局元素列表的架構。
    """
    layout_items: list[DetectedLayoutItem] = []

class FindLayoutItemsInput(BaseModel):
    """
    用于處理單個頁面的輸入模式。
    """
    document_path: str
    base64_jpeg: str
    page_number: int

class DocumentParsingAgent:
    def __init__(self, model_name="gemini-1.5-flash-002"):
        """
        使用適當的模式初始化LLM。
        """
        layout_elements_schema = prepare_schema_for_gemini(LayoutElements)
        self.model_name = model_name
        self.model = genai.GenerativeModel(
            self.model_name,
            generation_config={
                "response_mime_type": "application/json",
                "response_schema": layout_elements_schema,
            },
        )
    def find_layout_items(self, state: FindLayoutItemsInput):
        """
        Send a page image to the LLM for segmentation and summarization.
        """
        messages = [
            f"Find and summarize all the relevant layout elements in this PDF page in the following format: "
            f"{LayoutElements.schema_json()}. "
            f"Tables should have at least two columns and at least two rows. "
            f"The coordinates should overlap with each layout item.",
            {"mime_type": "image/jpeg", "data": state.base64_jpeg},
        ]
        # 向LLM發送提示信息
        result = self.model.generate_content(messages)
        data = json.loads(result.text)

        # 將JSON輸出轉換為文檔
        documents = [
            Document(
                page_content=item["summary"],
                metadata={
                    "page_number": state.page_number,
                    "element_type": item["element_type"],
                    "document_path": state.document_path,
                },
            )
            for item in data["layout_items"]
        ]
        return {"documents": documents}

from langgraph.types import Send

class DocumentParsingAgent:
    @classmethod
    def continue_to_find_layout_items(cls, state):
        """
        生成任務以并行處理每個頁面。
        """
        return [
            Send(
                "find_layout_items",
                FindLayoutItemsInput(
                    base64_jpeg=base64_jpeg,
                    page_number=i,
                    document_path=state.document_path,
                ),
            )
            for i, base64_jpeg in enumerate(state.pages_as_base64_jpeg_images)
        ]

from langgraph.graph import StateGraph, START, END

class DocumentParsingAgent:
    def build_agent(self):
        """
        使用狀態圖構建代理工作流。
        """
        builder = StateGraph(DocumentLayoutParsingState)

        # 添加節點，用于圖像提取和布局項檢測
        builder.add_node("get_images", self.get_images)
        builder.add_node("find_layout_items", self.find_layout_items)
        #定義圖形的流程
        builder.add_edge(START, "get_images")
        builder.add_conditional_edges("get_images", self.continue_to_find_layout_items)
        builder.add_edge("find_layout_items", END)

        self.graph = builder.compile()

if __name__ == "__main__":
    _state = DocumentLayoutParsingState(
        document_path="path/to/document.pdf"
    )
    agent = DocumentParsingAgent()

    # 步驟1：從PDF中提取圖像
    result_images = agent.get_images(_state)
    _state.pages_as_base64_jpeg_images = result_images["pages_as_base64_jpeg_images"]

    #步驟2：處理第一頁（作為一個示例）
    result_layout = agent.find_layout_items(
        FindLayoutItemsInput(
            base64_jpeg=_state.pages_as_base64_jpeg_images[0],
            page_number=0,
            document_path=_state.document_path,
        )
    )
    # 顯示處理結果
    for item in result_layout["documents"]:
        print(item.page_content)
        print(item.metadata["element_type"])

class DocumentRAGAgent:
    def index_documents(self, state: DocumentRAGState):
        """
        將解析后的文檔索引到向量存儲區中。
        """
        assert state.documents, "Documents should have at least one element"
        # 檢查該文檔是否已被編入索引
        if self.vector_store.get(where={"document_path": state.document_path})["ids"]:
            logger.info(
                "Documents for this file are already indexed, exiting this node"
            )
            return  #如果已經完成，跳過索引
        # 將解析后的文檔添加到向量存儲區中
        self.vector_store.add_documents(state.documents)
        logger.info(f"Indexed {len(state.documents)} documents for {state.document_path}")

class DocumentRAGAgent:
    def answer_question(self, state: DocumentRAGState):
        """
        檢索相關的數據塊，并生成針對用戶問題的響應。
        """
        # 根據查詢檢索前k個相關文檔
        relevant_documents: list[Document] = self.retriever.invoke(state.question)

        # 檢索相應的頁面圖像（避免重復）
        images = list(
            set(
                [
                    state.pages_as_base64_jpeg_images[doc.metadata["page_number"]]
                    for doc in relevant_documents
                ]
            )
        )
        logger.info(f"Responding to question: {state.question}")
        #構建提示：結合圖像、相關總結和問題
        messages = (
            [{"mime_type": "image/jpeg", "data": base64_jpeg} for base64_jpeg in images]
            + [doc.page_content for doc in relevant_documents]
            + [
                f"Answer this question using the context images and text elements only: {state.question}",
            ]
        )
        #使用LLM生成響應
        response = self.model.generate_content(messages)
        return {"response": response.text, "relevant_documents": relevant_documents}

class DocumentRAGAgent:
    def build_agent(self):
        """
        構建RAG代理的工作流。
        """
        builder = StateGraph(DocumentRAGState)
        # 添加用于編制索引和回答問題的節點
        builder.add_node("index_documents", self.index_documents)
        builder.add_node("answer_question", self.answer_question)
        # 定義工作流
        builder.add_edge(START, "index_documents")
        builder.add_edge("index_documents", "answer_question")
        builder.add_edge("answer_question", END)
        self.graph = builder.compile()

if __name__ == "__main__":
    from pathlib import Path

  # 導入要解析文檔的第一個代理
    from document_ai_agents.document_parsing_agent import (
        DocumentLayoutParsingState,
        DocumentParsingAgent,
    )
    # 步驟1：使用第一個代理來解析文檔
    state1 = DocumentLayoutParsingState(
        document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf")
    )
    agent1 = DocumentParsingAgent()
    result1 = agent1.graph.invoke(state1)
    #步驟2：設置第二個代理進行檢索和應答
    state2 = DocumentRAGState(
        question="Who was acknowledged in this paper?",
        document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf"),
        pages_as_base64_jpeg_images=result1["pages_as_base64_jpeg_images"],
        documents=result1["documents"],
    )
    agent2 = DocumentRAGAgent()
    # 索引文檔
    agent2.graph.invoke(state2)
    # 回答第一個問題
    result2 = agent2.graph.invoke(state2)
    print(result2["response"])
    # 回答第二個問題
    state3 = DocumentRAGState(
        question="What is the macro average when fine-tuning on PubLayNet using M-RCNN?",
        document_path=str(Path(__file__).parents[1] / "data" / "docs.pdf"),
        pages_as_base64_jpeg_images=result1["pages_as_base64_jpeg_images"],
        documents=result1["documents"],
    )
    result3 = agent2.graph.invoke(state3)
    print(result3["response"])