使用 Jina AI 和 openGauss 构建 RAG ​

import os
import urllib.request

JINA_API_KEY = os.getenv("JINA_API_KEY", "")

def jina_read(url: str) -> str:
    """Convert a web page to clean markdown via Jina Reader API."""
    reader_url = f"https://r.jina.ai/{url}"
    req = urllib.request.Request(reader_url)
    req.add_header("User-Agent", "Mozilla/5.0")
    if JINA_API_KEY:
        req.add_header("Authorization", f"Bearer {JINA_API_KEY}")
    req.add_header("X-Return-Format", "markdown")
    with urllib.request.urlopen(req) as resp:
        return resp.read().decode("utf-8")

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

file=jina_read(url)

print(file)

from typing import List

def text_chunks(source: str) -> List[str]:
    """Fetch and chunk a document into text chunks."""
    all_chunks: List[str] = []

    if source.startswith("http://") or source.startswith("https://"):
        text = jina_read(source)
        text_lines = text.split("# ")
        text_lines = [line for line in text_lines if line.strip()]
        all_chunks.extend(text_lines)
    else:
        with open(source, "r", encoding="utf-8") as f:
            text = f.read()
        text_lines = text.split("# ")
        text_lines = [line for line in text_lines if line.strip()]
        all_chunks.extend(text_lines)

    return all_chunks

chunk_texts = text_chunks(url)

for i, chunk in enumerate(chunk_texts[:3]):
    print(f"[Chunk {i}]: {chunk[:80]}...")

from openai import OpenAI

jina_client = OpenAI(
    api_key=os.getenv("JINA_API_KEY", ""),
    base_url="https://api.jina.ai/v1",
)

def batch_embedder(texts: List[str], task: str = "retrieval.passage") -> List[List[float]]:
    """Generate embedding vectors via Jina Embeddings API."""
    if not texts:
        return []
    vectors: List[List[float]] = []
    max_batch_size = 16
    for i in range(0, len(texts), max_batch_size):
        batch = texts[i : i + max_batch_size]
        response = jina_client.embeddings.create(
            input=batch,
            model="jina-embeddings-v5-text-small",
            extra_body={"task": task, "truncate": True},
        )
        vectors.extend(item.embedding for item in response.data)
    return vectors

import uuid
from typing import Any

vectors = batch_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 = "jina_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 1;
        """
    )
    .format(table_name=sql.Identifier(table_name))
    .as_string(cursor), 
)

print(cursor.fetchall())

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

question_vec = batch_embedder(question, task="retrieval.query")[0]

print(question_vec)

candidate_k = 20

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, candidate_k),
)

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

print(len(candidates))
print(candidates[0])

import json

def rerank(query: str, documents: List[str], top_n: int = 5) -> List[dict]:
    """Rerank documents against a query via Jina Reranker API."""
    if not documents:
        return []
    payload = json.dumps({
        "model": "jina-reranker-v3",
        "query": query,
        "documents": documents,
        "top_n": top_n,
    }).encode("utf-8")
    req = urllib.request.Request("https://api.jina.ai/v1/rerank", data=payload)
    req.add_header("User-Agent", "Mozilla/5.0")
    req.add_header("Content-Type", "application/json")
    if JINA_API_KEY:
        req.add_header("Authorization", f"Bearer {JINA_API_KEY}")
    with urllib.request.urlopen(req) as resp:
        body = json.loads(resp.read().decode("utf-8"))
    return body.get("results", [])

results = rerank(question, candidates, top_n=5)
reranked_contexts = [candidates[r["index"]] for r in results]

print("精排后的文档顺序及相关性得分:")
for r in results:
    idx = r["index"]
    score = r["relevance_score"]
    print(f"  [{idx}] score={score:.4f}  {candidates[idx][:80]}...")

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

print(contexts_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=contexts_text, question=question
        ),
    },
]

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)

#!/usr/bin/env python3
"""RAG module based on Jina AI (Reader + Embeddings + Reranker) and openGauss."""

import os
import json
import hashlib
import urllib.request
from typing import List, Any
from tqdm import tqdm

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

import psycopg2
from psycopg2 import sql


JINA_API_KEY = os.getenv("JINA_API_KEY", "")
JINA_READER_URL = "https://r.jina.ai/"
JINA_RERANKER_URL = "https://api.jina.ai/v1/rerank"

jina_client = OpenAI(
    api_key=JINA_API_KEY,
    base_url="https://api.jina.ai/v1",
)

deepseek_client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    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", "jinaai")


def jina_read(url: str) -> str:
    """Convert a web page to clean markdown via Jina Reader API."""
    reader_url = f"{JINA_READER_URL}{url}"
    req = urllib.request.Request(reader_url)
    req.add_header("User-Agent", "Mozilla/5.0")
    if JINA_API_KEY:
        req.add_header("Authorization", f"Bearer {JINA_API_KEY}")
    req.add_header("X-Return-Format", "markdown")
    with urllib.request.urlopen(req) as resp:
        return resp.read().decode("utf-8")


def batch_embedder(
    texts: List[str], task: str = "retrieval.passage"
) -> List[List[float]]:
    """Generate embedding vectors via Jina Embeddings API."""
    if not texts:
        return []
    vectors: List[List[float]] = []
    max_batch_size = 16
    for i in range(0, len(texts), max_batch_size):
        batch = texts[i : i + max_batch_size]
        response = jina_client.embeddings.create(
            input=batch,
            model="jina-embeddings-v5-text-small",
            extra_body={"task": task, "truncate": True},
        )
        vectors.extend(item.embedding for item in response.data)
    return vectors


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


def rerank(query: str, documents: List[str], top_n: int = 5) -> List[dict]:
    """Rerank documents against a query via Jina Reranker API."""
    if not documents:
        return []
    payload = json.dumps(
        {
            "model": "jina-reranker-v3",
            "query": query,
            "documents": documents,
            "top_n": top_n,
        }
    ).encode("utf-8")
    req = urllib.request.Request(JINA_RERANKER_URL, data=payload)
    req.add_header("User-Agent", "Mozilla/5.0")
    req.add_header("Content-Type", "application/json")
    if JINA_API_KEY:
        req.add_header("Authorization", f"Bearer {JINA_API_KEY}")
    with urllib.request.urlopen(req) as resp:
        body = json.loads(resp.read().decode("utf-8"))
    return body.get("results", [])


def text_chunks(source: str) -> List[str]:
    """Fetch and chunk a document into text chunks."""
    all_chunks: List[str] = []

    if source.startswith("http://") or source.startswith("https://"):
        text = jina_read(source)
        text_lines = text.split("# ")
        text_lines = [line for line in text_lines if line.strip()]
        all_chunks.extend(text_lines)
    else:
        with open(source, "r", encoding="utf-8") as f:
            text = f.read()
        text_lines = text.split("# ")
        text_lines = [line for line in text_lines if line.strip()]
        all_chunks.extend(text_lines)

    return all_chunks


def batch_text_chunks(sources: List[str]) -> List[tuple]:
    """Convert multiple document sources into (source, chunk) tuples."""
    items: List[Any] = []
    for source in tqdm(sources, desc="Reading & chunking"):
        for text in text_chunks(source):
            items.append((source, 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, task="retrieval.passage")
    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, task="retrieval.query")
    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 retrieve_with_rerank(
    cursor, table_name: str, question: str, top_k: int = 5, candidate_k: int = 20
) -> List[str]:
    """Two-stage retrieval: coarse vector search + Jina Reranker fine ranking."""
    candidates = retrieve(cursor, table_name, question, top_k=candidate_k)
    if len(candidates) <= top_k:
        return candidates

    results = rerank(question, candidates, top_n=top_k)
    reranked = [candidates[r["index"]] for r in results]
    return reranked


def reasoning(messages: list[ChatCompletionMessageParam]) -> str:
    """Send messages to DeepSeek 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)
    print()

    return "".join(answer_chunks)


def rag(  # pylint: disable=too-many-positional-arguments
    cursor,
    table_name: str,
    question: str,
    top_k: int = 5,
    use_rerank: bool = True,
    candidate_k: int = 20,
) -> str:
    """Run the full RAG pipeline: retrieve + optional rerank + generate answer."""
    if use_rerank:
        contexts = retrieve_with_rerank(
            cursor, table_name, question, top_k=top_k, candidate_k=candidate_k
        )
    else:
        contexts = retrieve(cursor, table_name, question, top_k=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://docs.opengauss.org/zh/docs/latest/datavec/datavec_overview.html"]
    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支持哪些向量数据类型"
        print("问题:", QUESTION)
        print("\n--- 未使用 Reranker ---")
        rag(cursor, table_name, QUESTION, top_k=5, use_rerank=False)
        print("\n--- 使用 Jina Reranker 精排后 ---")
        rag(cursor, table_name, QUESTION, top_k=5, use_rerank=True)
    finally:
        close_connection(conn, cursor)