使用 openGauss 和 Docling 构建 RAG ​

from docling.document_converter import DocumentConverter

converter = DocumentConverter()

source = "https://docs.opengauss.org/zh/docs/latest/datavec/datavec_overview.html"

doc = converter.convert(source).document
print(doc)

from docling_core.transforms.chunker import HierarchicalChunker

chunker = HierarchicalChunker()

chunks = chunker.chunk(doc)

for index, chunk in enumerate(chunks, start=1):
    print(f"\n===== chunk {index} =====")
    print(chunk)

texts: list[Any] = []
for chunk in chunks:
    text = (chunk.text or "").strip()
    if text:
        texts.append(text)

for index, text in enumerate(texts, start=1):
    print(f"\n===== text {index} =====")
    print(text)

import os
from typing import List, Any

from openai import OpenAI


client = OpenAI(
    api_key=os.getenv(
        "DASHSCOPE_API_KEY"
    ),
    base_url="https://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/compatible-mode/v1",
)

def batch_embedder(texts: List[str]) -> List[List[float]]:
    vectors: List[List[float]] = []
    max_batch_size: int = 10

    for i in range(0, len(texts), max_batch_size):
        batch = texts[i: i + max_batch_size]
        response = client.embeddings.create(
            input=batch,
            model="text-embedding-v4",
        )
        vectors.extend(item.embedding for item in response.data)
    return vectors

vector = batch_embedder(["Is it okay?"])
print(len(vector[0]))
print(vector[0])

import uuid
import json

vectors = batch_embedder(texts)
dim = len(vectors[0])

rows: list[Any] = []
for i, (text, vector) in enumerate(zip(texts, vectors)):
    rows.append(
        (
            str(uuid.uuid4()),
            text,
            vector,
            json.dumps({"source": source, "chunk_index": i}, ensure_ascii=False),
        )
    )

print(dim, rows[0])

import psycopg2

conn = psycopg2.connect(
    dbname = os.getenv("OG_DBNAME", "postgres"),
    user = os.getenv("OG_USER", "gaussdb"),
    password = os.getenv("OG_PASSWORD", "openGauss@123"),
    host = os.getenv("OG_HOST", "localhost"),
    port = int(os.getenv("OG_PORT", "5432"))
)
cursor = conn.cursor()

print(conn)
print(cursor)

from psycopg2 import sql

table_name = "docling_test"

cursor.execute(
    sql.SQL(
        """
        DROP TABLE IF EXISTS public.{table_name};
        """
    ).format(table_name=sql.Identifier(table_name))
)

cursor.execute(
    sql.SQL(
        """
        CREATE TABLE IF NOT EXISTS public.{table_name} (
            id TEXT PRIMARY KEY,
            content TEXT NOT NULL,
            embedding vector({dim}) NOT NULL,
            metadata JSONB,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """
    ).format(table_name=sql.Identifier(table_name), dim=sql.Literal(dim)),
)
conn.commit()

cursor.execute(
    sql.SQL(
        """
        SELECT column_name, data_type
        FROM information_schema.columns
        WHERE table_schema='public' AND table_name={table_name}
        ORDER BY ordinal_position;
    	"""
    ).format(table_name=sql.Literal(table_name)),
)

print(cursor.fetchall())

cursor.executemany(
    sql.SQL(
        """
        INSERT INTO public.{table_name} (id, content, embedding, metadata)
        VALUES (%s, %s, %s, %s::jsonb)
        """
    )
    .format(table_name=sql.Identifier(table_name))
    .as_string(cursor),
    rows,
)
conn.commit()

cursor.execute(
    sql.SQL(
        """
        SELECT id, content FROM public.{table_name}
        LIMIT 3;
        """
    )
    .format(table_name=sql.Identifier(table_name))
    .as_string(cursor), 
)

print(cursor.fetchall())

question = "openGauss支持哪些向量数据类型"

question_vec = batch_embedder([question])[0]

print(question_vec)

topk: int = 5

cursor.execute(
    sql.SQL(
        """
        SELECT content FROM public.{table_name}
        ORDER BY embedding <-> %s::vector
        LIMIT %s::int;
        """
    )
    .format(table_name=sql.Identifier(table_name)),
    (question_vec, topk),
)

contexts = [row[0] for row in cursor.fetchall()]

print(contexts)

context_text: str = ""
for i, c in enumerate(contexts, start=1):
    context_text += f"[{i}] {c}\n\n"
    
print(context_text)

from openai.types.chat import ChatCompletionMessageParam

SYSTEM_PROMPT = """你是一个智能问答助手。
你的任务是根据提供的参考资料（Context）回答用户问题。
请遵循以下规则：
1. 仅根据参考资料回答，不要编造信息。
2. 如果资料中没有答案，请直接说“根据提供的资料，无法回答该问题”。
3. 回答要简洁、准确。"""

USER_PROMPT_TEMPLATE = """请根据以下参考资料回答问题。

<参考资料>
{contexts_text}
</参考资料>

<问题>
{question}
</问题>

请开始回答："""

messages: List[ChatCompletionMessageParam] = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {
        "role": "user",
        "content": USER_PROMPT_TEMPLATE.format(
            contexts_text=context_text, question=question
        ),
    },
]

completion = client.chat.completions.create(
    model="qwen3.5-plus",
    messages=messages,
    extra_body={"enable_thinking": True},
    stream=True,
)
is_answering = False
print("\n" + "=" * 20 + "思考过程" + "=" * 20)
for chunk in completion:
    delta = chunk.choices[0].delta
    if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None:
        if not is_answering:
            print(delta.reasoning_content, end="", flush=True)
    if hasattr(delta, "content") and delta.content:
        if not is_answering:
            print("\n" + "=" * 20 + "完整回复" + "=" * 20)
            is_answering = True
        print(delta.content, end="", flush=True)

#!/usr/bin/env python3
"""RAG (Retrieval-Augmented Generation) module based on Docling and openGauss."""

import os
import json
import hashlib
from typing import List, Any
from tqdm import tqdm

from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam
from docling.document_converter import DocumentConverter
from docling_core.transforms.chunker import HierarchicalChunker

import psycopg2
from psycopg2 import sql


converter = DocumentConverter()
chunker = HierarchicalChunker()

client = OpenAI(
    api_key=os.getenv(
        "DASHSCOPE_API_KEY",
        "sk-03cac48e1e4f4d9badc4f02a73756dc6"
    ),
    base_url="https://{WorkspaceId}.cn-beijing.maas.aliyuncs.com/compatible-mode/v1",
)

DB_CONFIG = {
    "dbname": os.getenv("OG_DBNAME", "postgres"),
    "user": os.getenv("OG_USER", "gaussdb"),
    "password": os.getenv("OG_PASSWORD", "openGauss@123"),
    "host": os.getenv("OG_HOST", "localhost"),
    "port": int(os.getenv("OG_PORT", "5432")),
}

TABLE_NAME = os.getenv("OG_TABLE", "docling")


def reasoning(messages: list[ChatCompletionMessageParam]) -> str:
    """Send messages to the LLM and return the response with streaming output."""
    completion = client.chat.completions.create(
        model="qwen3.5-plus",
        messages=messages,
        extra_body={"enable_thinking": True},
        stream=True,
    )
    is_answering = False
    answer_chunks = []
    print("\n" + "=" * 20 + "思考过程" + "=" * 20)
    for chunk in completion:
        delta = chunk.choices[0].delta
        if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None:
            if not is_answering:
                print(delta.reasoning_content, end="", flush=True)
        if hasattr(delta, "content") and delta.content:
            if not is_answering:
                print("\n" + "=" * 20 + "完整回复" + "=" * 20)
                is_answering = True
            print(delta.content, end="", flush=True)
            answer_chunks.append(delta.content)
    print()
    return "".join(answer_chunks)


def batch_embedder(texts: List[str]) -> List[List[float]]:
    """Generate embedding vectors for a batch of texts."""
    vectors: List[List[float]] = []
    max_batch_size: int = 10

    for i in range(0, len(texts), max_batch_size):
        batch = texts[i: i + max_batch_size]
        response = client.embeddings.create(
            input=batch,
            model="text-embedding-v4",
        )
        vectors.extend(item.embedding for item in response.data)
    return vectors


def embedder(text: str) -> List[float]:
    """Generate an embedding vector for a single text."""
    return batch_embedder([text])[0]


def text_chunks(source: str) -> List[str]:
    """Convert a document source into a list of text chunks."""
    doc = converter.convert(source).document
    chunks = chunker.chunk(doc)
    texts: list[Any] = []
    for chunk in chunks:
        text = (chunk.text or "").strip()
        if text:
            texts.append(text)
    return texts


def batch_text_chunks(sources: List[str]) -> List[tuple]:
    """Convert multiple document sources into (source, chunk) tuples."""
    items: List[Any] = []
    for s in sources:
        for t in text_chunks(s):
            items.append((s, t))
    return items


def compute_doc_id(source: str, text: str) -> str:
    """Compute a SHA-256 document ID from a source and its text chunk."""
    return hashlib.sha256(f"{source}\n{text}".encode("utf-8")).hexdigest()


def filter_new_items(cursor, table_name: str, items: List[tuple]) -> List[tuple]:
    """Filter out items whose document IDs already exist in the database table."""
    ids = [compute_doc_id(s, t) for s, t in items]
    try:
        cursor.execute(
            sql.SQL(
                """
                SELECT id FROM public.{table_name} 
                WHERE id = ANY(%s);
                """
            ).format(table_name=sql.Identifier(table_name)),
            (ids,),
        )
        existing = {row[0] for row in cursor.fetchall()}
    except psycopg2.errors.UndefinedTable:
        cursor.connection.rollback()
        existing = set()

    new_items: List[Any] = []
    seen = set()
    for item, doc_id in zip(items, ids):
        if doc_id in existing or doc_id in seen:
            continue
        seen.add(doc_id)
        new_items.append(item)
    return new_items


def build_vector_rows(items: List[tuple]):
    """Build database rows with embeddings from (source, text) tuples."""
    texts = [t for _, t in items]
    vectors = batch_embedder(texts)
    dim = len(vectors[0])

    rows: list[Any] = []
    for i, ((source, text), vector) in enumerate(zip(items, vectors)):
        rows.append(
            (
                compute_doc_id(source, text),
                text,
                vector,
                json.dumps({"source": source, "chunk_index": i}, ensure_ascii=False),
            )
        )
    return rows, dim


def create_connection():
    """Create and return a database connection and cursor."""
    conn = psycopg2.connect(
        dbname=DB_CONFIG["dbname"],
        user=DB_CONFIG["user"],
        password=DB_CONFIG["password"],
        host=DB_CONFIG["host"],
        port=DB_CONFIG["port"],
    )
    return conn, conn.cursor()


def close_connection(conn, cursor):
    """Close the database cursor and connection."""
    cursor.close()
    conn.close()


def create_table(conn, cursor, table_name: str, dim: int):
    """Create the vector store table if it does not already exist."""
    cursor.execute(
        sql.SQL(
            """
            CREATE TABLE IF NOT EXISTS public.{table_name} (
                id TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                embedding vector({dim}) NOT NULL,
                metadata JSONB,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
            """
        ).format(table_name=sql.Identifier(table_name), dim=sql.Literal(dim)),
    )
    conn.commit()


def insert_rows(conn, cursor, table_name: str, rows):
    """Insert vector rows into the database table."""
    if not rows:
        return
    stmt = (
        sql.SQL(
            """
            INSERT INTO public.{table_name} (id, content, embedding, metadata)
            VALUES (%s, %s, %s, %s::jsonb)
            """
        )
        .format(table_name=sql.Identifier(table_name))
        .as_string(cursor)
    )

    for row in tqdm(rows, desc="Inserting", unit="row"):
        cursor.execute(stmt, row)
    conn.commit()


def retrieve(cursor, table_name: str, question: str, top_k: int = 5) -> List[str]:
    """Retrieve the top-k most relevant document chunks for a question."""
    question_vec = embedder(question)
    cursor.execute(
        sql.SQL(
            """
            SELECT content FROM public.{table_name}
            ORDER BY embedding <-> %s::vector
            LIMIT %s::int;
            """
        ).format(table_name=sql.Identifier(table_name)),
        (question_vec, top_k),
    )
    return [row[0] for row in cursor.fetchall()]


def rag(cursor, table_name: str, question: str, top_k: int = 5) -> str:
    """Run the full RAG pipeline: retrieve relevant chunks and generate an answer."""
    contexts = retrieve(cursor, table_name, question, top_k)

    contexts_text: str = ""
    for i, c in enumerate(contexts, start=1):
        contexts_text += f"[{i}] {c}\n\n"

    SYSTEM_PROMPT = """你是一个智能问答助手。
你的任务是根据提供的参考资料（Context）回答用户问题。
请遵循以下规则：
1. 仅根据参考资料回答，不要编造信息。
2. 如果资料中没有答案，请直接说“根据提供的资料，无法回答该问题”。
3. 回答要简洁、准确。"""

    USER_PROMPT_TEMPLATE = """请根据以下参考资料回答问题。

<参考资料>
{contexts_text}
</参考资料>

<问题>
{question}
</问题>

请开始回答："""

    messages: List[ChatCompletionMessageParam] = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": USER_PROMPT_TEMPLATE.format(
                contexts_text=contexts_text, question=question
            ),
        },
    ]
    return reasoning(messages)


if __name__ == "__main__":
    sources = [
        "https://docs.opengauss.org/zh/docs/latest/datavec/datavec_overview.html",
        "https://docs.opengauss.org/zh/docs/latest/datavec/vector_index.html",
        "https://docs.opengauss.org/zh/docs/latest/datavec/datavec_architecture.html",
    ]
    table_name = TABLE_NAME

    items = batch_text_chunks(sources)

    conn, cursor = create_connection()
    try:
        new_items = filter_new_items(cursor, table_name, items)

        if not new_items:
            print(f"[SKIP] No new chunks. table=public.{table_name}")
        else:
            rows, dim = build_vector_rows(new_items)
            create_table(conn, cursor, table_name, dim)
            insert_rows(conn, cursor, table_name, rows)
            print(f"[OK]   Inserted {len(rows)} rows. table=public.{table_name}, dim={dim}")

        question = "openGauss支持哪些向量数据类型"    # pylint: disable=invalid-name
        print("问题:", question)
        rag(cursor, table_name, question, top_k=5)
    finally:
        close_connection(conn, cursor)