https://ragflow.io/docs/dev/python_api_reference

RAGFlow.create_dataset(
    name: str,
    avatar: str = "",
    description: str = "",
    embedding_model: str = "BAAI/bge-large-zh-v1.5",
    language: str = "English",
    permission: str = "me", 
    chunk_method: str = "naive",
    parser_config: DataSet.ParserConfig = None
)

DataSet.update(update_message : dict)

DataSet.upload_documents(document_list : list[dict])

Document.update(update_message : dict)

DataSet.async_parse_documents(document_ids : list[str])

Document.add_chunk(content: str, important_keywords: list[str] = [])

Chunk.update(update_message : dict)

RAGFlow.retrieve(
    question: str = "",
    dataset_ids: list[str] = None,
    document_ids: list[str] = None,
    page: int = 1,
    page_size: int = 30,
    similarity_threshold: float = 0.2,
    vector_similarity_weight: float = 0.3,
    top_k: int = 1024,
    rerank_id: str = None,
    keyword: bool = False,
    highlight: bool = False
)

RAGFlow.create_chat(
    name: str, 
    avatar: str = "", 
    dataset_ids: list[str] = [], 
    llm: Chat.LLM = None, 
    prompt: Chat.Prompt = None
)

Chat.update(update_message : dict)

Chat.create_session(name: str = "New session")

Session.ask(question: str = "", stream: bool = False, **kwargs)

Agent.create_session(id, rag, **kwargs)

Session.ask(question: str = "", stream: bool = False)

# 需要修改源碼的示例邏輯
class CustomTerminologyProcessor:
    def __init__(self, terminology_mapping):
        self.terminology_mapping = terminology_mapping  # 同義詞映射表


    def process_query(self, query):
        # 專業(yè)術(shù)語標(biāo)準(zhǔn)化
        # 車間俚語轉(zhuǎn)換為標(biāo)準(zhǔn)術(shù)語
        processed_query = query
        for slang, standard_term in self.terminology_mapping.items():
            processed_query = processed_query.replace(slang, standard_term)
        return processed_query

# 這部分需要修改源碼，以下是概念性代碼
class CustomHybridIndexBuilder:
    def __init__(self, vector_db_client):
        self.client = vector_db_client


    def create_scalar_indices(self, collection_name, fields):
        # 創(chuàng)建設(shè)備編號(hào)、故障代碼等標(biāo)量索引
        for field in fields:
            self.client.create_index(collection_name, field, "scalar")


    def create_vector_indices(self, collection_name, fields):
        # 創(chuàng)建向量索引
        for field in fields:
            self.client.create_index(collection_name, field, {"index_type": "HNSW", "params": {"M": 16, "efConstruction": 200}})

# 查詢路由器 - 需要源碼修改實(shí)現(xiàn)
class QueryRouter:
    def route_query(self, query_text):
        if self._is_equipment_code(query_text):
            return "exact_match", {"field": "equipment_code"}
        elif self._is_fault_code(query_text):
            return "exact_match", {"field": "fault_code"}
        elif self._is_parameter_query(query_text):
            return "parameter_lookup", {"field": "parameter_name"}
        else:
            return "semantic_search", {"model": "embedding_model"}

# 上下文增強(qiáng)器 - 需要修改源碼實(shí)現(xiàn)
class ContextEnhancer:
    def enhance_query(self, query, session_history, equipment_metadata=None):
        # 添加設(shè)備上下文信息
        if equipment_metadata:
            query_context = f"設(shè)備型號(hào): {equipment_metadata['model']}, 生產(chǎn)年份: {equipment_metadata['year']}\n"
            query_context += query


        # 添加歷史查詢信息
        if session_history:
            relevant_history = self._extract_relevant_history(session_history, query)
            query_context = f"參考?xì)v史信息: {relevant_history}\n" + query_context


        return query_context

# 使用多模態(tài)模型生成圖片描述
processor = MultimodalDocumentProcessor(api_key="YOUR_API_KEY")
enhanced_docs = processor.process_pdf("設(shè)備手冊(cè).pdf")

# 配置保留圖片的數(shù)據(jù)集
dataset = rag_object.create_dataset(
    name="圖文設(shè)備手冊(cè)",
    chunk_method="paper",
    parser_cnotallow={"layout_recognize": True}
)


# 上傳原始PDF文檔
with open("設(shè)備手冊(cè).pdf", "rb") as f:
    dataset.upload_documents([{"display_name": "設(shè)備手冊(cè).pdf", "blob": f.read()}])

assistant = rag_object.create_chat(
    name="圖文設(shè)備專家",
    dataset_ids=[dataset.id],
    prompt=Chat.Prompt(
        prompt="""你是設(shè)備維修專家。回答時(shí)，請(qǐng)同時(shí)提供:
        1. 文字描述解釋故障和解決方案
        2. 引用相關(guān)圖片，包括圖片描述
        3. 告訴用戶可以參考哪些圖片獲取更多信息
        {knowledge}"""
    )
)

from ragflow_sdk import RAGFlow
import fitz  # PyMuPDF
import re
import pandas as pd
import io
import logging
from typing import List, Dict, Any, Tuple


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class DynamicChunker:
    """針對(duì)混合格式文檔的自定義分塊器"""


    def __init__(self, api_key: str, base_url: str):
        self.rag_object = RAGFlow(api_key=api_key, base_url=base_url)


    def process_document(self, pdf_path: str, dataset_name: str = None) -> str:
        """處理PDF文檔并實(shí)現(xiàn)自定義分塊"""
        # 1. 創(chuàng)建或獲取數(shù)據(jù)集
        if dataset_name:
            datasets = self.rag_object.list_datasets(name=dataset_name)
            if datasets:
                dataset = datasets[0]
                logger.info(f"使用現(xiàn)有數(shù)據(jù)集: {dataset.name}, ID: {dataset.id}")
            else:
                dataset = self.rag_object.create_dataset(
                    name=dataset_name,
                    embedding_model="BAAI/bge-large-zh-v1.5",
                    language="Chinese",
                    chunk_method="one"  # 使用"one"方式，保持文檔完整性
                )
                logger.info(f"創(chuàng)建新數(shù)據(jù)集: {dataset.name}, ID: {dataset.id}")
        else:
            dataset = self.rag_object.create_dataset(
                name="混合文檔庫",
                embedding_model="BAAI/bge-large-zh-v1.5",
                language="Chinese",
                chunk_method="one"
            )
            logger.info(f"創(chuàng)建新數(shù)據(jù)集: {dataset.name}, ID: {dataset.id}")


        # 2. 上傳文檔
        with open(pdf_path, "rb") as f:
            document_blob = f.read()


        docs = dataset.upload_documents([{
            "display_name": pdf_path.split("/")[-1],
            "blob": document_blob
        }])


        if not docs:
            logger.error("文檔上傳失敗")
            return None


        doc = docs[0]
        logger.info(f"文檔上傳成功, ID: {doc.id}")


        # 3. 分析文檔結(jié)構(gòu)并創(chuàng)建自定義分塊
        chunks = self._create_custom_chunks(pdf_path)
        logger.info(f"創(chuàng)建了 {len(chunks)} 個(gè)自定義分塊")


        # 4. 手動(dòng)添加分塊到文檔
        chunk_ids = []
        for chunk in chunks:
            content = chunk["content"]
            keywords = chunk["keywords"]
            chunk_type = chunk["type"]


            # 添加分塊類型標(biāo)記以便未來檢索
            content_with_marker = f"[{chunk_type}]\n{content}"


            # 添加分塊
            new_chunk = doc.add_chunk(
                cnotallow=content_with_marker,
                important_keywords=keywords
            )


            chunk_ids.append(new_chunk.id)
            logger.info(f"添加分塊: 類型={chunk_type}, ID={new_chunk.id}, 關(guān)鍵詞數(shù)量={len(keywords)}")


        return doc.id


    def _create_custom_chunks(self, pdf_path: str) -> List[Dict[str, Any]]:
        """分析PDF并創(chuàng)建自定義分塊"""
        pdf_doc = fitz.open(pdf_path)
        chunks = []


        # 用于存儲(chǔ)當(dāng)前上下文
        current_section = ""
        current_text_chunk = ""
        current_table_data = []
        current_keywords = []


        for page_idx, page in enumerate(pdf_doc):
            # 提取頁面文本
            text = page.get_text()


            # 檢測(cè)章節(jié)標(biāo)題
            section_headers = self._detect_section_headers(text)
            if section_headers:
                # 如果有積累的文本塊，先保存
                if current_text_chunk:
                    chunks.append({
                        "content": current_text_chunk,
                        "keywords": list(set(current_keywords)),
                        "type": "text",
                        "section": current_section
                    })
                    current_text_chunk = ""
                    current_keywords = []


                # 更新當(dāng)前章節(jié)
                current_section = section_headers[0]
                current_keywords.append(current_section)


            # 檢測(cè)表格
            tables = self._detect_tables(page)
            if tables:
                # 如果有積累的文本塊，先保存
                if current_text_chunk:
                    chunks.append({
                        "content": current_text_chunk,
                        "keywords": list(set(current_keywords)),
                        "type": "text",
                        "section": current_section
                    })
                    current_text_chunk = ""
                    current_keywords = []


                # 處理表格
                for table in tables:
                    table_content = self._format_table(table)
                    table_keywords = self._extract_keywords_from_table(table)


                    chunks.append({
                        "content": table_content,
                        "keywords": list(set(table_keywords + current_keywords)),
                        "type": "table",
                        "section": current_section
                    })


            # 檢測(cè)圖片
            images = self._detect_images(page)
            for img_idx, img in enumerate(images):
                # 提取圖片周圍的文本作為上下文
                img_context = self._extract_image_context(text, img["bbox"])


                # 如果找到圖片相關(guān)文本
                if img_context:
                    # 如果有積累的文本塊，先保存
                    if current_text_chunk:
                        chunks.append({
                            "content": current_text_chunk,
                            "keywords": list(set(current_keywords)),
                            "type": "text",
                            "section": current_section
                        })
                        current_text_chunk = ""
                        current_keywords = []


                    # 創(chuàng)建圖文塊
                    image_text_content = f"圖片描述: {img_context}\n圖片位置: 第{page_idx+1}頁"
                    image_keywords = self._extract_keywords(img_context)


                    chunks.append({
                        "content": image_text_content,
                        "keywords": list(set(image_keywords + current_keywords)),
                        "type": "image",
                        "section": current_section
                    })


                    # 防止重復(fù)處理同一段文本
                    text = text.replace(img_context, "", 1)


            # 處理剩余文本
            if text.strip():
                # 按段落拆分
                paragraphs = self._split_into_paragraphs(text)


                for para in paragraphs:
                    if len(para.strip()) < 10:  # 忽略太短的段落
                        continue


                    # 累積文本直到達(dá)到合適的大小
                    current_text_chunk += para + "\n\n"
                    current_keywords.extend(self._extract_keywords(para))


                    # 檢查是否應(yīng)該創(chuàng)建新的文本塊
                    if len(current_text_chunk) > 1500:  # 大約300-500個(gè)詞
                        chunks.append({
                            "content": current_text_chunk,
                            "keywords": list(set(current_keywords)),
                            "type": "text",
                            "section": current_section
                        })
                        current_text_chunk = ""
                        current_keywords = []


        # 處理最后一個(gè)文本塊
        if current_text_chunk:
            chunks.append({
                "content": current_text_chunk,
                "keywords": list(set(current_keywords)),
                "type": "text",
                "section": current_section
            })


        return chunks


    def _detect_section_headers(self, text: str) -> List[str]:
        """檢測(cè)章節(jié)標(biāo)題"""
        # 這是一個(gè)簡(jiǎn)化的實(shí)現(xiàn)，可以根據(jù)實(shí)際文檔格式調(diào)整
        header_patterns = [
            r"^第[一二三四五六七八九十\d]+章\s*(.+)$",
            r"^\d+\.\d*\s+(.+)$",
            r"^[一二三四五六七八九十]+[、.．]\s*(.+)$"
        ]


        headers = []
        for line in text.split("\n"):
            line = line.strip()
            for pattern in header_patterns:
                match = re.match(pattern, line)
                if match:
                    headers.append(line)
                    break


        return headers


    def _detect_tables(self, page) -> List[Any]:
        """檢測(cè)頁面中的表格"""
        # 這里使用簡(jiǎn)化的檢測(cè)邏輯，實(shí)際應(yīng)用中可能需要更復(fù)雜的表格檢測(cè)算法
        # 例如，可以尋找包含多個(gè)垂直和水平線的區(qū)域
        tables = []


        # 簡(jiǎn)單表格檢測(cè)：查找水平線和垂直線的集中區(qū)域
        # 這里僅作為占位示例，實(shí)際實(shí)現(xiàn)會(huì)更復(fù)雜
        horizontal_lines = []
        vertical_lines = []


        for drawing in page.get_drawings():
            for item in drawing["items"]:
                if item["type"] == "l":  # 線段
                    x0, y0, x1, y1 = item["rect"]
                    if abs(y1 - y0) < 2:  # 水平線
                        horizontal_lines.append((x0, y0, x1, y1))
                    elif abs(x1 - x0) < 2:  # 垂直線
                        vertical_lines.append((x0, y0, x1, y1))


        # 簡(jiǎn)單的表格判定：有足夠多的水平線和垂直線
        if len(horizontal_lines) >= 3 and len(vertical_lines) >= 2:
            # 找出所有線的邊界，作為表格邊界
            min_x = min([min(l[0], l[2]) for l in horizontal_lines + vertical_lines])
            max_x = max([max(l[0], l[2]) for l in horizontal_lines + vertical_lines])
            min_y = min([min(l[1], l[3]) for l in horizontal_lines + vertical_lines])
            max_y = max([max(l[1], l[3]) for l in horizontal_lines + vertical_lines])


            # 提取表格區(qū)域的文本
            table_rect = fitz.Rect(min_x, min_y, max_x, max_y)
            table_text = page.get_text("text", clip=table_rect)


            # 簡(jiǎn)化的表格結(jié)構(gòu)化（實(shí)際應(yīng)用中需要更復(fù)雜的邏輯）
            table_rows = table_text.split("\n")
            table = []
            for row in table_rows:
                if row.strip():
                    cells = row.split()
                    if len(cells) >= 2:  # 至少有2個(gè)單元格才視為有效行
                        table.append(cells)


            if table:
                tables.append(table)


        return tables


    def _detect_images(self, page) -> List[Dict[str, Any]]:
        """檢測(cè)頁面中的圖片"""
        images = []


        # 獲取頁面上的圖片對(duì)象
        img_list = page.get_images(full=True)


        for img_idx, img_info in enumerate(img_list):
            xref = img_info[0]
            base_image = page.parent.extract_image(xref)


            if base_image:
                # 尋找圖片在頁面上的位置
                img_rects = []
                for rect in page.get_image_rects(xref):
                    img_rects.append(rect)


                if img_rects:
                    # 使用第一個(gè)找到的位置
                    bbox = img_rects[0]


                    images.append({
                        "bbox": [bbox.x0, bbox.y0, bbox.x1, bbox.y1]
                    })


        return images


    def _extract_image_context(self, text: str, bbox: List[float], context_size: int = 200) -> str:
        """提取圖片周圍的文本作為上下文"""
        # 在文本中查找可能的圖片標(biāo)題，如"圖1"，"Figure 1"等
        caption_patterns = [
            r"圖\s*\d+[\.．:]?\s*(.+?)(?:\n|$)",
            r"Figure\s*\d+[\.:]?\s*(.+?)(?:\n|$)",
            r"圖表\s*\d+[\.．:]?\s*(.+?)(?:\n|$)"
        ]


        for pattern in caption_patterns:
            matches = re.finditer(pattern, text, re.IGNORECASE)
            for match in matches:
                return match.group(0)


        # 如果找不到明確的圖片標(biāo)題，則嘗試提取圖片周圍的文本
        # 這是一個(gè)簡(jiǎn)化實(shí)現(xiàn)，實(shí)際應(yīng)用中可能需要更復(fù)雜的邏輯
        lines = text.split("\n")
        total_length = 0
        start_line = 0


        # 估計(jì)圖片在文本中的位置
        # 這是一個(gè)非常粗略的估計(jì)，實(shí)際應(yīng)用中需要更精確的方法
        relative_position = bbox[1] / 1000  # 假設(shè)頁面高度為1000
        target_position = int(len(text) * relative_position)


        # 找到大致對(duì)應(yīng)的行
        for i, line in enumerate(lines):
            total_length += len(line) + 1  # +1 for newline
            if total_length > target_position:
                start_line = max(0, i - 2)  # 從前兩行開始
                break


        # 提取上下文（當(dāng)前行及其前后幾行）
        context_lines = lines[max(0, start_line):min(len(lines), start_line + 5)]
        return "\n".join(line for line in context_lines if len(line.strip()) > 5)


    def _format_table(self, table: List[List[str]]) -> str:
        """格式化表格為文本格式"""
        if not table:
            return ""


        # 創(chuàng)建pandas DataFrame
        df = pd.DataFrame(table[1:], columns=table[0] if len(table) > 1 else None)


        # 轉(zhuǎn)換為字符串形式
        result = io.StringIO()
        df.to_csv(result, sep="|", index=False)
        return result.getvalue()


    def _extract_keywords_from_table(self, table: List[List[str]]) -> List[str]:
        """從表格中提取關(guān)鍵詞"""
        keywords = []


        # 表頭作為關(guān)鍵詞
        if table and len(table) > 0:
            keywords.extend(table[0])


        # 第一列可能是行標(biāo)題，也加入關(guān)鍵詞
        for row in table[1:] if len(table) > 1 else []:
            if row and len(row) > 0:
                keywords.append(row[0])


        return keywords


    def _split_into_paragraphs(self, text: str) -> List[str]:
        """將文本拆分為段落"""
        # 按照多個(gè)換行符拆分
        paragraphs = re.split(r"\n\s*\n", text)
        return [p.strip() for p in paragraphs if p.strip()]


    def _extract_keywords(self, text: str) -> List[str]:
        """從文本中提取關(guān)鍵詞"""
        # 這是一個(gè)簡(jiǎn)化的關(guān)鍵詞提取方法
        # 實(shí)際應(yīng)用中可以使用更復(fù)雜的NLP技術(shù)，如TF-IDF、TextRank等


        # 1. 移除停用詞和標(biāo)點(diǎn)
        stop_words = {"的", "了", "和", "與", "或", "在", "是", "有", "被", "將", "把"}
        cleaned_text = re.sub(r'[^\w\s]', ' ', text)
        words = cleaned_text.split()
        filtered_words = [w for w in words if w not in stop_words and len(w) > 1]


        # 2. 簡(jiǎn)單詞頻統(tǒng)計(jì)
        word_count = {}
        for word in filtered_words:
            if word in word_count:
                word_count[word] += 1
            else:
                word_count[word] = 1


        # 3. 選擇頻率最高的幾個(gè)詞作為關(guān)鍵詞
        sorted_words = sorted(word_count.items(), key=lambda x: x[1], reverse=True)
        return [word for word, count in sorted_words[:10] if count > 1]


# 使用示例
if __name__ == "__main__":
    chunker = DynamicChunker(
        api_key="YOUR_API_KEY",
        base_url="http://YOUR_BASE_URL:9380"
    )


    doc_id = chunker.process_document("設(shè)備維修手冊(cè).pdf", "動(dòng)態(tài)分塊測(cè)試")
    print(f"處理完成，文檔ID: {doc_id}")