-- Load the Iceberg extension
INSTALL iceberg;
LOAD iceberg;

-- Mount the warehouse
ATTACH 's3://my-warehouse' AS iceberg (
    TYPE iceberg,
    S3_ENDPOINT 's3.amazonaws.com',
    AWS_REGION 'us-east-1'
);

-- Create schema
CREATE SCHEMA iceberg.default.rag_knowledge;

CREATE TABLE iceberg.default.rag_knowledge.documents (
    id BIGINT PRIMARY KEY,
    title VARCHAR,
    content VARCHAR,
    embedding DOUBLE[1536],        -- OpenAI ada-002 embedding
    category VARCHAR,
    created_at TIMESTAMP,
    updated_at TIMESTAMP
);

-- Load the Lance extension
INSTALL lance;
LOAD lance;

-- Export from Iceberg to Lance format
COPY (
    SELECT id, title, content, embedding, category, created_at
    FROM iceberg.default.rag_knowledge.documents
) TO 's3://my-lance-bucket/documents.lance'
    (FORMAT lance, mode 'overwrite');

-- Create an IVF_FLAT vector index on the Lance dataset
CREATE INDEX embedding_idx ON 's3://my-lance-bucket/documents.lance' (embedding)
USING IVF_FLAT WITH (
    num_partitions = 64,
    num_iterations = 10,
    metric_type = 'cosine'    -- Cosine similarity, ideal for OpenAI embeddings
);

SELECT 
    id, 
    title, 
    content, 
    category,
    _hybrid_score,     -- Combined score
    _distance,         -- Vector distance
    _score             -- Vector similarity score
FROM lance_hybrid_search(
    's3://my-lance-bucket/documents.lance',
    'embedding',                      -- Vector column name
    [0.1, 0.2, 0.3, ...]::DOUBLE[1536],  -- Query vector
    'content',                        -- Text column name
    'data analytics automated reports', -- Keywords
    k = 20,                           -- Return top-K
    alpha = 0.7                       -- 0=keywords only, 1=vectors only
)
ORDER BY _hybrid_score DESC;

-- Assume new document data comes from an external system
WITH new_docs AS (
    SELECT 
        1001 AS id,
        'DuckDB Memory Optimization Guide' AS title,
        'DuckDB memory management strategies...' AS content,
        ARRAY[0.05, 0.12, ...]::DOUBLE[1536] AS embedding,
        'Technical Tutorial' AS category,
        CURRENT_TIMESTAMP AS ts
)
MERGE INTO iceberg.default.rag_knowledge.documents AS target
USING new_docs AS source
ON target.id = source.id
WHEN MATCHED THEN
    UPDATE SET 
        title = source.title,
        content = source.content,
        embedding = source.embedding,
        category = source.category,
        updated_at = source.ts
WHEN NOT MATCHED THEN
    INSERT (id, title, content, embedding, category, created_at, updated_at)
    VALUES (
        source.id, source.title, source.content, 
        source.embedding, source.category, source.ts, source.ts
    );

MERGE INTO iceberg.default.rag_knowledge.documents AS target
USING (
    SELECT * FROM read_json_auto('s3://etl-output/daily-docs.json')
) AS source
ON target.id = source.id
WHEN MATCHED THEN UPDATE SET
    title = source.title,
    content = source.content,
    embedding = source.embedding
WHEN NOT MATCHED THEN INSERT ALL
    VALUES (
        source.id, source.title, source.content, 
        source.embedding, source.category, source.created_at, source.updated_at
    );

import duckdb
import openai
from datetime import datetime

# Connect to DuckDB
con = duckdb.connect('ai-data-lake.duckdb')

# 1. Load Iceberg extension
con.execute("INSTALL iceberg; LOAD iceberg;")
con.execute("ATTACH 's3://my-warehouse' AS iceberg (TYPE iceberg);")

# 2. Read raw data from S3
raw_data = con.execute("""
    SELECT * FROM read_json_auto('s3://raw-data/knowledge-base/*.json')
""").fetchdf()

# 3. Generate embeddings (OpenAI API)
client = openai.OpenAI(api_key="sk-xxx")
embeddings = []
for _, row in raw_data.iterrows():
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=row["content"]
    )
    embeddings.append(resp.data[0].embedding)
raw_data["embedding"] = embeddings

# 4. Upsert to Iceberg
for _, row in raw_data.iterrows():
    con.execute("""
        MERGE INTO iceberg.default.rag_knowledge.documents AS t
        USING (SELECT ? AS id, ? AS title, ? AS content, 
                    ?::DOUBLE[1536] AS embedding, ? AS category) AS s
        ON t.id = s.id
        WHEN MATCHED THEN UPDATE SET content = s.content, embedding = s.embedding
        WHEN NOT MATCHED THEN INSERT VALUES (s.id, s.title, s.content, s.embedding, s.category)
    """, [row["id"], row["title"], row["content"], 
          row["embedding"], row.get("category", "General")])

# 5. Export Lance for vector search
con.execute("""
    COPY (SELECT * FROM iceberg.default.rag_knowledge.documents) 
    TO 's3://my-lance-bucket/documents.lance' 
    (FORMAT lance)
""")

print(f"✅ Synced {len(raw_data)} documents to AI Data Lake")

from fastapi import FastAPI
import duckdb
import openai

app = FastAPI()
con = duckdb.connect()
client = openai.OpenAI()

@app.post("/search")
def search(query: str):
    # Generate query embedding
    emb = client.embeddings.create(
        model="text-embedding-3-small",
        input=query
    ).data[0].embedding
    
    # DuckDB hybrid search
    results = con.execute("""
        SELECT id, title, content, category, _hybrid_score
        FROM lance_hybrid_search(
            's3://my-lance-bucket/documents.lance',
            'embedding', ?, 'content', ?, k=5, alpha=0.7
        )
    """, [str(emb), query]).fetchdf()
    
    return {"results": results.to_dict(orient="records")}