AI: RAG 检索增强

from openai import OpenAI

client = OpenAI()

def get_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

text = "The company travel allowance is $45 per day."
vector = get_embedding(text)

print(f"Dimension: {len(vector)}")
print(f"First 5 values: {vector[:5]}")

Dimension: 1536
First 5 values: [0.0123, -0.0045, 0.0378, -0.0211, 0.0056]

import numpy as np

def cosine_similarity(a: list[float], b: list[float]) -> float:
    a_arr = np.array(a)
    b_arr = np.array(b)
    dot = np.dot(a_arr, b_arr)
    norm_a = np.linalg.norm(a_arr)
    norm_b = np.linalg.norm(b_arr)
    return float(dot / (norm_a * norm_b))

sentences = {
    "travel_allowance": get_embedding("The company travel allowance is $45 per day."),
    "reimbursement": get_embedding("How do I claim travel reimbursement?"),
    "unrelated": get_embedding("The weather is sunny today."),
}

sim1 = cosine_similarity(sentences["travel_allowance"], sentences["reimbursement"])
sim2 = cosine_similarity(sentences["travel_allowance"], sentences["unrelated"])

print(f"Related sentences:   {sim1:.4f}")
print(f"Unrelated sentences: {sim2:.4f}")

Related sentences:   0.8234
Unrelated sentences: 0.3012

def top_k_search(
    query_embedding: list[float],
    doc_embeddings: list[list[float]],
    doc_texts: list[str],
    k: int = 3,
) -> list[tuple[str, float]]:
    similarities = [
        (doc_texts[i], cosine_similarity(query_embedding, doc_embeddings[i]))
        for i in range(len(doc_texts))
    ]
    similarities.sort(key=lambda x: x[1], reverse=True)
    return similarities[:k]

documents = [
    "Travel allowance is $45 per day for domestic trips.",
    "International travel requires VP approval in advance.",
    "Employees must submit receipts within 30 days.",
    "Remote work policy allows up to 2 days per week at home.",
    "The office kitchen is stocked with free coffee and snacks.",
]

doc_embeddings = [get_embedding(doc) for doc in documents]
query = "How much can I spend on travel per day?"
query_embedding = get_embedding(query)

results = top_k_search(query_embedding, doc_embeddings, documents, k=3)

for text, score in results:
    print(f"[{score:.4f}] {text}")

[0.8912] Travel allowance is $45 per day for domestic trips.
[0.7634] Employees must submit receipts within 30 days.
[0.7102] International travel requires VP approval in advance.

from openai import OpenAI

client = OpenAI()

def rag_query(question: str, documents: list[str], k: int = 3) -> str:
    query_emb = get_embedding(question)
    doc_embs = [get_embedding(doc) for doc in documents]
    top_docs = top_k_search(query_emb, doc_embs, documents, k=k)

    context = "\n\n".join([f"[Doc {i+1}] {text}" for i, (text, _) in enumerate(top_docs)])

    prompt = f"""Answer the question based ONLY on the following context.
If the context does not contain the answer, say "I don't know."

Context:
{context}

Question: {question}
Answer:"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    return response.choices[0].message.content

docs = [
    "Travel allowance is $45 per day for domestic trips.",
    "International travel requires VP approval at least 2 weeks in advance.",
    "Employees must submit expense reports within 30 days of the trip.",
    "The maximum hotel reimbursement is $200 per night.",
    "Flight bookings must use the company travel portal.",
]

answer = rag_query("What is the daily travel allowance?", docs)
print(answer)

The daily travel allowance is $45 for domestic trips.

def ask_without_rag(question: str) -> str:
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": question}],
        temperature=0,
    )
    return response.choices[0].message.content

question = "What is Acme Corp's daily travel allowance?"

print("=== Without RAG ===")
print(ask_without_rag(question))

print("\n=== With RAG ===")
print(rag_query(question, docs))

=== Without RAG ===
I don't have specific information about Acme Corp's daily travel allowance.
Company travel allowances vary, but a common range is $50-$100 per day.
Please check your company's travel policy for the exact amount.

=== With RAG ===
The daily travel allowance is $45 for domestic trips.

from langchain.text_splitter import RecursiveCharacterTextSplitter

text = """Acme Corp Travel Policy

1. Domestic Travel
The daily travel allowance is $45. This covers meals and incidental expenses.
Hotel reimbursement is capped at $200 per night.

2. International Travel
All international trips require VP approval at least 2 weeks in advance.
The daily allowance for international travel is $75.
Hotel reimbursement is capped at $300 per night.

3. Expense Reporting
All expense reports must be submitted within 30 days of trip completion.
Receipts are required for all expenses over $25."""

splitter = RecursiveCharacterTextSplitter(
    chunk_size=200,
    chunk_overlap=50,
    separators=["\n\n", "\n", ". ", " "],
)

chunks = splitter.split_text(text)
for i, chunk in enumerate(chunks):
    print(f"--- Chunk {i+1} ---")
    print(chunk)
    print()

--- Chunk 1 ---
Acme Corp Travel Policy

1. Domestic Travel
The daily travel allowance is $45. This covers meals and incidental expenses.
Hotel reimbursement is capped at $200 per night.

--- Chunk 2 ---
2. International Travel
All international trips require VP approval at least 2 weeks in advance.
The daily allowance for international travel is $75.

--- Chunk 3 ---
international travel is $75.
Hotel reimbursement is capped at $300 per night.

3. Expense Reporting
All expense reports must be submitted within 30 days of trip completion.

import numpy as np
from openai import OpenAI

client = OpenAI()

# --- Step 1: Prepare documents ---
documents = [
    {
        "title": "Travel Policy",
        "content": (
            "Acme Corp Travel Policy\n\n"
            "1. Domestic Travel\n"
            "The daily travel allowance is $45 for domestic trips. "
            "This covers meals and incidental expenses. "
            "Hotel reimbursement is capped at $200 per night.\n\n"
            "2. International Travel\n"
            "All international trips require VP approval at least 2 weeks in advance. "
            "The daily allowance for international travel is $75. "
            "Hotel reimbursement is capped at $300 per night."
        ),
    },
    {
        "title": "Leave Policy",
        "content": (
            "Acme Corp Leave Policy\n\n"
            "1. Annual Leave\n"
            "Full-time employees receive 15 days of paid annual leave per year. "
            "Unused leave can be carried over to the next year, up to a maximum of 5 days.\n\n"
            "2. Sick Leave\n"
            "Employees receive 10 days of paid sick leave per year. "
            "A doctor's note is required for absences exceeding 3 consecutive days."
        ),
    },
    {
        "title": "IT Security Policy",
        "content": (
            "Acme Corp IT Security Policy\n\n"
            "1. Password Requirements\n"
            "All passwords must be at least 12 characters long and include uppercase, "
            "lowercase, numbers, and special characters. "
            "Passwords must be changed every 90 days.\n\n"
            "2. Device Policy\n"
            "Personal devices may not connect to the corporate network. "
            "All company devices must have approved antivirus software installed."
        ),
    },
]

# --- Step 2: Chunk documents ---
def chunk_text(text: str, chunk_size: int = 200, overlap: int = 50) -> list[str]:
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunks.append(" ".join(words[start:end]))
        start += chunk_size - overlap
    return chunks

all_chunks: list[str] = []
chunk_sources: list[str] = []

for doc in documents:
    chunks = chunk_text(doc["content"])
    for chunk in chunks:
        all_chunks.append(chunk)
        chunk_sources.append(doc["title"])

print(f"Total chunks: {len(all_chunks)}")

# --- Step 3: Generate embeddings ---
def get_embedding(text: str, model: str = "text-embedding-3-small") -> list[float]:
    response = client.embeddings.create(input=text, model=model)
    return response.data[0].embedding

chunk_embeddings = [get_embedding(chunk) for chunk in all_chunks]
print(f"Embedding dimension: {len(chunk_embeddings[0])}")

# --- Step 4: Vector search ---
def cosine_similarity(a: list[float], b: list[float]) -> float:
    a_arr = np.array(a)
    b_arr = np.array(b)
    return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))

def search(query: str, k: int = 3) -> list[tuple[str, str, float]]:
    query_emb = get_embedding(query)
    scores = [
        (all_chunks[i], chunk_sources[i], cosine_similarity(query_emb, chunk_embeddings[i]))
        for i in range(len(all_chunks))
    ]
    scores.sort(key=lambda x: x[2], reverse=True)
    return scores[:k]

# --- Step 5: RAG query ---
def rag_answer(question: str, k: int = 3) -> str:
    results = search(question, k=k)

    context_parts = []
    for i, (chunk, source, score) in enumerate(results):
        context_parts.append(f"[Source: {source}, Relevance: {score:.2f}]\n{chunk}")
    context = "\n\n---\n\n".join(context_parts)

    prompt = f"""Answer the question based ONLY on the following context.
If the context does not contain enough information, say "I don't have enough information."
Always cite which source document your answer comes from.

Context:
{context}

Question: {question}
Answer:"""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    return response.choices[0].message.content

# --- Step 6: Test ---
questions = [
    "What is the daily travel allowance for domestic trips?",
    "How many days of annual leave do employees get?",
    "What are the password requirements?",
    "What is the company's remote work policy?",
]

for q in questions:
    print(f"Q: {q}")
    print(f"A: {rag_answer(q)}")
    print()

Total chunks: 9
Embedding dimension: 1536

Q: What is the daily travel allowance for domestic trips?
A: The daily travel allowance for domestic trips is $45, according to the Travel Policy.

Q: How many days of annual leave do employees get?
A: Full-time employees receive 15 days of paid annual leave per year, according to the Leave Policy.

Q: What are the password requirements?
A: Passwords must be at least 12 characters long and include uppercase, lowercase, numbers, and special characters. They must be changed every 90 days, according to the IT Security Policy.

Q: What is the company's remote work policy?
A: I don't have enough information. The provided context does not contain any remote work policy.

texts = [
    "Cats are popular pets known for their independence.",
    "Dogs are loyal companions that love to play fetch.",
    "The stock market rose 3% today on strong earnings.",
    "Felines enjoy climbing and sleeping in sunny spots.",
    "Interest rates were raised by the central bank.",
]

def vector_search(
    query: str,
    corpus: list[str],
    k: int = 5,
    embedding_model: str = "text-embedding-3-small",
) -> list[dict]:
    """
    Returns top-k most similar documents with scores.
    Each result is a dict: {"text": ..., "score": ..., "rank": ...}
    """
    # Your implementation here
    pass