색인: 분할 조각 임베딩 및 저장

from openai import OpenAI

client = OpenAI()

def embed_batch(texts, model='text-embedding-3-small'):
    response = client.embeddings.create(
        model=model,
        input=texts  # up to 2048 texts per call
    )
    return [item.embedding for item in response.data]

# Embed one batch of 100 chunk texts
texts = [chunk['text'] for chunk in chunks[:100]]
vectors = embed_batch(texts)
print(f'Embedding dimension: {len(vectors[0])}')
print(f'Embedded {len(vectors)} chunks')

def embed_all_chunks(chunks, batch_size=200):
    embedded = []
    total = len(chunks)
    for i in range(0, total, batch_size):
        batch = chunks[i:i+batch_size]
        texts = [c['text'] for c in batch]
        vectors = embed_batch(texts)
        for chunk, vector in zip(batch, vectors):
            embedded.append({
                **chunk,
                'embedding': vector
            })
        if (i // batch_size) % 10 == 0:
            print(f'Progress: {min(i+batch_size, total)}/{total}')
    return embedded

import time
import random
from openai import RateLimitError

def embed_batch_with_retry(texts, max_retries=5):
    for attempt in range(max_retries):
        try:
            return embed_batch(texts)
        except RateLimitError:
            if attempt == max_retries - 1:
                raise
            wait = (2 ** attempt) + random.uniform(0, 1)
            print(f'Rate limited. Waiting {wait:.1f}s...')
            time.sleep(wait)
    return []

import pinecone

pc = pinecone.Pinecone(api_key='YOUR_KEY')
index = pc.Index('rag-index')

def upsert_to_pinecone(embedded_chunks, batch_size=100):
    for i in range(0, len(embedded_chunks), batch_size):
        batch = embedded_chunks[i:i+batch_size]
        vectors = [
            (
                chunk['id'],
                chunk['embedding'],
                {
                    'text': chunk['text'],
                    'source': chunk['metadata']['source'],
                    'page': chunk['metadata'].get('page', 0)
                }
            )
            for chunk in batch
        ]
        index.upsert(vectors=vectors)
        print(f'Upserted {min(i+batch_size, len(embedded_chunks))}/{len(embedded_chunks)}')

import psycopg2
from psycopg2.extras import execute_values

def upsert_to_pgvector(conn, embedded_chunks):
    with conn.cursor() as cur:
        records = [
            (
                chunk['id'],
                chunk['text'],
                chunk['metadata']['source'],
                chunk['metadata'].get('page', 0),
                chunk['embedding']   # list of floats
            )
            for chunk in embedded_chunks
        ]
        execute_values(cur, '''
            INSERT INTO document_chunks (id, text, source, page, embedding)
            VALUES %s
            ON CONFLICT (id) DO UPDATE
            SET text = EXCLUDED.text, embedding = EXCLUDED.embedding
        ''', records)
    conn.commit()

-- Build HNSW index after bulk insertion
CREATE INDEX CONCURRENTLY ON document_chunks
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Set ef_search at query time to trade recall vs speed
SET hnsw.ef_search = 100;

-- Verify index was created
SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = 'document_chunks';

# Well-designed metadata schema
METADATA_SCHEMA = {
    # For filtering at retrieval time
    'department': 'HR',         # string
    'doc_type': 'policy',       # string
    'year': 2025,               # integer
    'is_active': True,          # boolean

    # For display in citations
    'title': 'Employee Handbook 2025',
    'author': 'HR Team',
    'page': 12,
    'source': 's3://docs/handbook_2025.pdf',

    # For debugging and updates
    'chunk_index': 3,
    'total_chunks': 24,
    'indexed_at': '2025-09-01T10:00:00Z'
}

import json
from pathlib import Path

CHECKPOINT_FILE = '/tmp/index_checkpoint.json'

def load_checkpoint():
    if Path(CHECKPOINT_FILE).exists():
        return set(json.loads(Path(CHECKPOINT_FILE).read_text()))
    return set()

def save_checkpoint(indexed_ids):
    Path(CHECKPOINT_FILE).write_text(json.dumps(list(indexed_ids)))

def index_with_checkpoint(chunks, index):
    done = load_checkpoint()
    remaining = [c for c in chunks if c['id'] not in done]
    print(f'Resuming: {len(done)} done, {len(remaining)} remaining')
    for chunk in remaining:
        upsert_to_pinecone([chunk], index)
        done.add(chunk['id'])
        save_checkpoint(done)

def verify_index(index, chunks, sample_size=10):
    index_stats = index.describe_index_stats()
    total_vectors = index_stats.total_vector_count
    expected = len(chunks)
    print(f'Index vectors: {total_vectors}, Expected: {expected}')
    if total_vectors != expected:
        print('WARNING: mismatch — some chunks may not have been indexed')

    # Spot-check retrieval
    import random
    sample = random.sample(chunks, sample_size)
    for chunk in sample:
        vec = embed_batch([chunk['text']])[0]
        results = index.query(vector=vec, top_k=1, include_metadata=True)
        top_id = results.matches[0].id if results.matches else None
        if top_id != chunk['id']:
            print(f'WARNING: expected {chunk["id"]}, got {top_id}')

from sentence_transformers import SentenceTransformer

# Load once at startup
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

def embed_locally(texts, batch_size=64):
    # encode() handles batching internally
    embeddings = model.encode(
        texts,
        batch_size=batch_size,
        show_progress_bar=True,
        convert_to_numpy=True
    )
    return embeddings.tolist()  # convert numpy array to Python list

vectors = embed_locally([c['text'] for c in chunks])