使用 BentoML 和 openGauss 构建 RAG ​

import os
from typing import List, Any

import psycopg2
from psycopg2 import sql


class OpenGaussClient:
    """A simple openGauss client using psycopg2."""

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

    def drop_table(self, table_name: str):
        """Drop table."""
        self.cursor.execute(
            sql.SQL(
                """
                DROP TABLE IF EXISTS public.{table_name};
                """
            ).format(table_name=sql.Identifier(table_name))
        )
        self.conn.commit()

    def create_table(self, table_name: str, dim: int):
        """Create table."""
        self.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)),
        )
        self.conn.commit()

    def insert_into_table(self, table_name: str, rows: List[Any]) -> None:
        """Insert into the table."""
        self.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(self.cursor),
            rows,
        )
        self.conn.commit()

    def query(self, table_name: str, question_vec: List[float], top_k: int = 5) -> List[str]:
        """Execute a vector similarity query and return the results."""
        self.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 self.cursor.fetchall()]

    def close(self):
        """Close the connection."""
        self.cursor.close()
        self.conn.close()

import random
import string
import urllib.request

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

    for url in urls:
        if url.startswith("http://") or url.startswith("https://"):
            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)
        else:
            with open(url, "r", encoding="utf-8") as f:
                file_text = f.read()
            text_lines = file_text.split("# ")
            text_lines = [line for line in text_lines if line.strip()]
            all_chunks.extend(text_lines)

    return all_chunks

import uuid
import numpy as np
from openai import OpenAI
from fastembed import TextEmbedding


class RAG:
    """RAG (Retrieval-Augmented Generation) class built upon FastEmbed, DeepSeek and openGauss."""

    def __init__(self, deepseek_client: OpenAI, opengauss_client: OpenGaussClient):
        self._prepare_deepseek(deepseek_client)
        self._prepare_opengauss(opengauss_client)

    def _embedder(self, text: str) -> List[float]:
        """Embed text into embedding vector."""
        response = list(self.embedding_model.embed([text]))
        result = response[0]
        return result.tolist() if isinstance(result, np.ndarray) else result

    def _prepare_deepseek(self, deepseek_client: OpenAI, llm_model: str = "deepseek-v4-pro"):
        """Prepare DeepSeek client and embedding model."""
        self.deepseek_client = deepseek_client
        self.llm_model = llm_model
        self.embedding_model = TextEmbedding(model_name="BAAI/bge-small-zh-v1.5")
        self.SYSTEM_PROMPT = """你是一个智能问答助手。
你的任务是根据提供的参考资料（Context）回答用户问题。
请遵循以下规则：
1. 仅根据参考资料回答，不要编造信息。
2. 如果资料中没有答案，请直接说"根据提供的资料，无法回答该问题"。
3. 回答要简洁、准确。"""
        self.USER_PROMPT = """请根据以下参考资料回答问题。

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

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

请开始回答："""

    def _prepare_opengauss(self, opengauss_client: OpenGaussClient, table_name: str = "bentoml_rag"):
        """Prepare openGauss client."""
        self.opengauss_client = opengauss_client
        self.table_name = table_name

        self.opengauss_client.drop_table(self.table_name)
        dim = len(self._embedder("get_dim"))
        self.opengauss_client.create_table(self.table_name, dim)

    def load(self, urls: List[str]):
        """Load the urls."""
        rows: list[Any] = []
        texts = chunk_source(urls)
        for text in texts:
            rows.append(
                (
                    str(uuid.uuid4()),
                    text,
                    self._embedder(text),
                )
            )

        self.opengauss_client.insert_into_table(table_name=self.table_name, rows=rows)

    def retrieve(self, question: str, top_k: int = 3):
        """Retrieve the top k results."""
        question_vec = self._embedder(question)
        retrieve_texts = self.opengauss_client.query(self.table_name, question_vec, top_k)
        return retrieve_texts

    def answer(self, question: str, top_k: int = 3):
        """Answer the question."""
        retrieve_texts = self.retrieve(question, top_k)

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

        completion = self.deepseek_client.chat.completions.create(
            model=self.llm_model,
            messages=[
                {"role": "system", "content": self.SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": self.USER_PROMPT.format(
                        contexts_text=retrieve_texts_str, question=question
                    ),
                },
            ],
            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:
                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 build_rag():
    """Build Rag."""
    deepseek_client = OpenAI(
        api_key=os.getenv("DEEPSEEK_API_KEY"),
        base_url="https://api.deepseek.com",
    )
    opengauss_client = OpenGaussClient()
    rag = RAG(deepseek_client=deepseek_client, opengauss_client=opengauss_client)
    return rag

rag = build_rag()

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

rag.load(urls)

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

response = rag.answer(question=question, top_k=3)

print(response)

import json
import os

import bentoml
import pydantic
from openai import OpenAI
from starlette.responses import Response

import rag as rag_module


def json_response(data: dict) -> Response:
    """Return a JSON response with unicode preserved."""
    return Response(
        json.dumps(data, ensure_ascii=False, indent=2),
        media_type="application/json",
    )


class IngestRequest(pydantic.BaseModel):
    """Request schema for the /ingest endpoint."""

    sources: list[str] = pydantic.Field(default_factory=list)
    table_name: str = pydantic.Field(default="bentoml_rag")


class QueryRequest(pydantic.BaseModel):
    """Request schema for the /query endpoint."""

    question: str
    table_name: str = pydantic.Field(default="bentoml_rag")
    top_k: int = pydantic.Field(default=5, ge=1)


@bentoml.service(
    name="opengauss_rag",
    traffic={"timeout": 300},
)
class OpenGaussRAG:
    """BentoML service that wraps a RAG pipeline backed by openGauss DataVec."""

    def __init__(self) -> None:
        deepseek_client = OpenAI(
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            base_url="https://api.deepseek.com",
        )
        self.client = rag_module.OpenGaussClient()
        self.rag = rag_module.RAG(
            deepseek_client=deepseek_client, opengauss_client=self.client
        )

    @bentoml.api
    async def ingest(self, request: IngestRequest):
        """Ingest documents: chunk, embed, and store into openGauss."""
        if not request.sources:
            return json_response({"status": "ok", "skipped": True})

        self.rag.table_name = request.table_name
        self.rag._prepare_opengauss(self.client, request.table_name)
        self.rag.load(request.sources)

        return json_response({"status": "ok", "table_name": request.table_name})

    @bentoml.api
    async def query(self, request: QueryRequest):
        """Query the RAG pipeline with a question."""
        answer = self.rag.answer(request.question, top_k=request.top_k)
        return json_response({
            "question": request.question,
            "answer": answer,
        })

    @bentoml.api
    async def health(self):
        """Health check endpoint."""
        return json_response({"status": "healthy"})

import json
import bentoml

client = bentoml.SyncHTTPClient("http://localhost:3000")

ingest_result = client.call(
    "ingest",
    request={
        "sources": [
            "https://raw.gitcode.com/opengauss/docs/raw/master/docs/zh/datavec/datavec_overview.md",
        ],
        "table_name": "bentoml_rag",
    },
)
print(json.dumps(ingest_result, indent=2, ensure_ascii=False))

query_result = client.call(
    "query",
    request={
        "question": "openGauss支持哪些向量数据类型",
        "table_name": "bentoml_rag",
        "top_k": 5,
    }
)
print(json.dumps(query_result, indent=2, ensure_ascii=False))

health_result = client.call(
    "health",
)
print(json.dumps(health_result, indent=2, ensure_ascii=False))