使用 openGauss 和 DeepSeek 构建 RAG ​

import random
import string
import urllib.request
from tqdm import tqdm

def chunk_source(urls: List[str]) -> List[str]:
    """Chunk text into chunks."""
    
    all_chunks: List[str] = []

    for url in tqdm(urls, desc="Chunking text"):
        random_name = "".join(random.choices(string.ascii_letters + string.digits, k=8))
        file_path = f"./{random_name}.md"
        try:
            urllib.request.urlretrieve(url, file_path)
            with open(file_path, "r", encoding="utf-8") as file:
                file_text = file.read()
            text_lines = file_text.split("# ")
            text_lines = [line for line in text_lines if line.strip()]
            all_chunks.extend(text_lines)
        finally:
            if os.path.exists(file_path):
                os.remove(file_path)

    return all_chunks

urls = [
    "https://raw.gitcode.com/opengauss/docs/raw/master/docs/zh/datavec/datavec_overview.md"
]

chunk_texts = chunk_source(urls)
print(chunk_texts)

from typing import List
import numpy as np
from fastembed import TextEmbedding

embedding_model = TextEmbedding(model_name="BAAI/bge-small-zh-v1.5")

def embedder(texts: List[str]) -> List[List[float]]:
    """ Embed text into embedding vector """

    response = list(embedding_model.embed(texts))
    result = [vec.tolist() if isinstance(vec, np.ndarray) else vec for vec in response]
    return result

vector = embedder(["Is it okay?", "Please help me!"])
print(len(vector[0]))
print(vector)

import uuid
from typing import Any

vectors = embedder(chunk_texts)
dim = len(vectors[0])

rows: list[Any] = []
for chunk_text, vector in zip(chunk_texts, vectors):
    rows.append(
        (
            str(uuid.uuid4()),
            chunk_text,
            vector
        )
    )

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 = "deepseek_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,
            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)
        VALUES (%s, %s, %s)
        """
    )
    .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 = 3

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
        ),
    },
]

import os
from openai import OpenAI

deepseek_client = OpenAI(
    api_key=os.getenv(
        "DEEPSEEK_API_KEY"
    ),
    base_url="https://api.deepseek.com",
)

completion = deepseek_client.chat.completions.create(
    model="deepseek-v4-pro",
    messages=messages,
    stream=True,
    reasoning_effort="high",
    extra_body={"thinking": {"type": "enabled"}}
)

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)

result = "".join(answer_chunks)
print(result)

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

import os
import json
import hashlib
import random
import string
import urllib.request
from typing import List, Any
from tqdm import tqdm
import numpy as np

from fastembed import TextEmbedding
from openai import OpenAI
from openai.types.chat import ChatCompletionMessageParam

import psycopg2
from psycopg2 import sql


embedding_model = TextEmbedding(model_name="BAAI/bge-small-zh-v1.5")
deepseek_client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY", "sk-776af5c01af54c77ae6aa9d46fe7f14e"),
    base_url="https://api.deepseek.com",
)

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", "deepseek")


def reasoning(messages: list[ChatCompletionMessageParam]) -> str:
    """Send messages to the LLM and return the response with streaming output."""

    completion = deepseek_client.chat.completions.create(
        model="deepseek-v4-pro",
        messages=messages,
        stream=True,
        reasoning_effort="high",
        extra_body={"thinking": {"type": "enabled"}},
    )

    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)

    return "".join(answer_chunks)


def batch_embedder(texts: List[str]) -> List[List[float]]:
    """Embed text into embedding vector"""

    response = list(embedding_model.embed(texts))
    result = [vec.tolist() if isinstance(vec, np.ndarray) else vec for vec in response]
    return result


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

    return batch_embedder([text])[0]


def text_chunks(url: str) -> List[str]:
    """Chunk text into chunks."""

    all_chunks: List[str] = []

    random_name = "".join(random.choices(string.ascii_letters + string.digits, k=8))
    file_path = f"./{random_name}.md"
    try:
        urllib.request.urlretrieve(url, file_path)
        with open(file_path, "r", encoding="utf-8") as file:
            file_text = file.read()
        text_lines = file_text.split("# ")
        text_lines = [line for line in text_lines if line.strip()]
        all_chunks.extend(text_lines)
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)

    return all_chunks


def batch_text_chunks(urls: List[str]) -> List[tuple]:
    """Convert multiple document sources into (source, chunk) tuples."""
    items: List[Any] = []
    for url in tqdm(urls, desc="Batching text chunks"):
        for text in text_chunks(url):
            items.append((url, text))
    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(u, t) for u, 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, ((url, text), vector) in enumerate(zip(items, vectors)):
        rows.append(
            (
                compute_doc_id(url, text),
                text,
                vector,
                json.dumps({"url": url, "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__":
    urls = [
        "https://raw.gitcode.com/opengauss/docs/raw/master/docs/zh/datavec/datavec_overview.md"
    ]
    table_name = TABLE_NAME

    items = batch_text_chunks(urls)

    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)