OGAI Usage Guide ​

CREATE EXTENSION ogai;

-- View all model configurations for the current user
SELECT model_key, model_name, model_provider, url FROM ogai.model_sources;

-- Register a new model
INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, api_key, owner_name)
VALUES ('my_model', 'text-embedding-v3', 'Qwen',
        'https://dashscope.aliyuncs.com/compatible-mode/v1',
        'sk-xxx', CURRENT_USER);

-- View all vectorization tasks for the current user
SELECT task_name, type, src_table, src_col, method, dim, enable_bm25
FROM ogai.vectorize_tasks;

-- View queue status statistics
SELECT status, COUNT(*) FROM ogai.vectorize_queue GROUP BY status;

-- View failed tasks
SELECT msg_id, task_id, pk_value, fail_reason, retry_count
FROM ogai.vectorize_queue WHERE status = 'failed';

SELECT ogai_embedding('openGauss is an open-source database', 'my_embed_model', 1536);

SELECT ogai_generate('What is a vector database?', 'qwen_chat');

SELECT * FROM ogai_rerank(
    'Database optimization',
    ARRAY['Index optimization', 'Backup strategy', 'Query plan'],
    'rerank_model'
) ORDER BY rerank_score DESC;

SELECT * FROM ogai_chunk('Long text...', 500, 100);

SELECT load_onnx_model('local_bge');

SELECT unload_onnx_model('local_bge');

ogai.ai_vectorize(
    p_task_name TEXT,
    p_task_type TEXT,
    p_index_type TEXT,
    p_embed_model TEXT,
    p_src_schema TEXT,
    p_src_table TEXT,
    p_src_col TEXT,
    p_primary_key TEXT,
    p_table_method TEXT,
    p_dim INTEGER,
    p_max_chunk_size INTEGER DEFAULT 1000,
    p_max_chunk_overlap INTEGER DEFAULT 200,
    p_enable_bm25 BOOLEAN DEFAULT true
) RETURNS TABLE(task_id INT, success BOOLEAN, processed_count INT, message TEXT)

-- Synchronous mode with the append method
SELECT * FROM ogai.ai_vectorize(
    'simple_task', 'sync', 'cosine', 'my_embed_model',
    'public', 'articles', 'content', 'id',
    'append', 1536, 0, 0, true
);

-- Asynchronous mode with the join method and chunking enabled
SELECT * FROM ogai.ai_vectorize(
    'chunked_task', 'async', 'cosine', 'my_embed_model',
    'public', 'documents', 'content', 'id',
    'join', 1536, 500, 100, true
);

-- Basic search
SELECT * FROM ogai.search('my_task', 'Database optimization methods');

-- Specify the columns and number of results to return
SELECT * FROM ogai.search('my_task', 'Database optimization', 'id, title, content', 5, '');

-- Set the hybrid search weight
SET ogai.hybrid_search_ratio = 0.7;

-- Run a hybrid search
SELECT * FROM ogai.hybrid_search('my_task', 'openGauss database optimization', 'id, title', 10, '');

SELECT ogai.rag(
    'What vector index types does openGauss support?',
    'knowledge_base_task',
    'rerank_model',
    'chat_model',
    5,
    20
);

SELECT * FROM ogai.ai_unvectorize('my_task');

INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, api_key, owner_name)
VALUES ('openai_embed', 'text-embedding-ada-002', 'openai',
        'https://api.openai.com/v1', 'sk-xxx', CURRENT_USER);

INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, api_key, owner_name)
VALUES ('qwen_embed', 'text-embedding-v3', 'Qwen',
        'https://dashscope.aliyuncs.com/compatible-mode/v1', 'sk-xxx', CURRENT_USER);

INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, owner_name)
VALUES ('ollama_embed', 'nomic-embed-text', 'ollama', 'http://localhost:11434', CURRENT_USER);

INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, owner_name)
VALUES ('onnx_bge', 'bge-small-zh', 'onnx', '/data/models/bge-small-zh.onnx', CURRENT_USER);

-- 1. Register models
INSERT INTO ogai.model_sources (model_key, model_name, model_provider, url, api_key, owner_name)
VALUES
    ('embed_model', 'text-embedding-v3', 'Qwen',
     'https://dashscope.aliyuncs.com/compatible-mode/v1', 'your-api-key', CURRENT_USER),
    ('chat_model', 'qwen-turbo', 'Qwen',
     'https://dashscope.aliyuncs.com/compatible-mode/v1', 'your-api-key', CURRENT_USER),
    ('rerank_model', 'gte-rerank', 'Qwen',
     'https://dashscope.aliyuncs.com/compatible-mode/v1', 'your-api-key', CURRENT_USER);

-- 2. Create a knowledge base table
CREATE TABLE knowledge_base (
    id SERIAL PRIMARY KEY,
    title TEXT,
    content TEXT,
    category TEXT
);

-- 3. Insert data
INSERT INTO knowledge_base (title, content, category) VALUES
('Introduction to openGauss', 'openGauss is an open-source relational database...', 'Product'),
('Vector indexes', 'openGauss supports Hierarchical Navigable Small World (HNSW), Inverted File Flat (IVFFlat), and other vector indexes...', 'Technology');

-- 4. Create a vectorization task
SELECT * FROM ogai.ai_vectorize(
    'kb_task', 'sync', 'cosine', 'embed_model',
    'public', 'knowledge_base', 'content', 'id',
    'join', 1024, 500, 100, true
);

-- 5. Perform a vector search
SELECT * FROM ogai.search('kb_task', 'What is openGauss', 'id, title', 5, '');

-- 6. Perform a hybrid search
SET ogai.hybrid_search_ratio = 0.7;
SELECT * FROM ogai.hybrid_search('kb_task', 'openGauss vector indexes', '', 5, '');

-- 7. Perform RAG question answering
SELECT ogai.rag('What indexes does openGauss support?', 'kb_task', 'rerank_model', 'chat_model', 5, 20);

-- 8. Clean up
SELECT * FROM ogai.ai_unvectorize('kb_task');