143 lines
3.8 KiB
Bash
143 lines
3.8 KiB
Bash
#!/bin/bash
|
|
set -e
|
|
|
|
POSTGRES_VERSION="${POSTGRES_VERSION:-16}"
|
|
MEMORY_MODEL="${MEMORY_MODEL:-all-MiniLM-L6-v2}"
|
|
|
|
# Install Python if not present
|
|
install_python() {
|
|
if command -v python3 &>/dev/null; then
|
|
echo "Python already installed: $(python3 --version)"
|
|
return 0
|
|
fi
|
|
|
|
echo "Installing Python..."
|
|
apt-get update
|
|
apt-get install -y python3 python3-pip python3-venv
|
|
|
|
echo "Python installed: $(python3 --version)"
|
|
}
|
|
|
|
# Install pgvector extension
|
|
install_pgvector() {
|
|
if [ -f "/usr/share/postgresql/${POSTGRES_VERSION}/extension/vector.control" ]; then
|
|
echo "pgvector already installed"
|
|
return 0
|
|
fi
|
|
|
|
echo "Installing pgvector..."
|
|
apt-get update
|
|
apt-get install -y postgresql-${POSTGRES_VERSION}-pgvector
|
|
|
|
echo "pgvector installed"
|
|
}
|
|
|
|
# Setup Python virtual environment and dependencies
|
|
setup_python_env() {
|
|
local skill_dir="$(dirname "$(dirname "$0")")"
|
|
local venv_dir="$skill_dir/.venv"
|
|
|
|
if [ -d "$venv_dir" ]; then
|
|
echo "Python venv already exists"
|
|
return 0
|
|
fi
|
|
|
|
echo "Creating Python virtual environment..."
|
|
python3 -m venv "$venv_dir"
|
|
|
|
echo "Installing Python dependencies..."
|
|
"$venv_dir/bin/pip" install --upgrade pip
|
|
"$venv_dir/bin/pip" install \
|
|
fastapi==0.109.0 \
|
|
uvicorn==0.27.0 \
|
|
sentence-transformers==2.3.1 \
|
|
psycopg2-binary==2.9.9 \
|
|
python-ulid==2.2.0
|
|
|
|
echo "Python environment ready"
|
|
}
|
|
|
|
# Pre-download embedding model
|
|
download_model() {
|
|
local skill_dir="$(dirname "$(dirname "$0")")"
|
|
local venv_dir="$skill_dir/.venv"
|
|
|
|
echo "Pre-downloading embedding model: $MEMORY_MODEL..."
|
|
"$venv_dir/bin/python" -c "
|
|
from sentence_transformers import SentenceTransformer
|
|
model = SentenceTransformer('$MEMORY_MODEL')
|
|
print(f'Model loaded: {model.get_sentence_embedding_dimension()} dimensions')
|
|
"
|
|
echo "Model downloaded"
|
|
}
|
|
|
|
# Initialize database schema
|
|
init_database() {
|
|
# Wait for postgres to be available
|
|
local retries=30
|
|
while [ $retries -gt 0 ]; do
|
|
if pg_isready -q 2>/dev/null; then
|
|
break
|
|
fi
|
|
echo "Waiting for PostgreSQL..."
|
|
sleep 1
|
|
retries=$((retries - 1))
|
|
done
|
|
|
|
if [ $retries -eq 0 ]; then
|
|
echo "PostgreSQL not available - schema will be created on first run"
|
|
return 0
|
|
fi
|
|
|
|
# Source postgres env if available
|
|
[ -f /run/vibestack/postgres.env ] && source /run/vibestack/postgres.env
|
|
|
|
local pg_user="${POSTGRES_USER:-vibestack}"
|
|
local pg_db="${POSTGRES_DB:-vibestack}"
|
|
|
|
echo "Initializing memory schema..."
|
|
psql -U "$pg_user" -d "$pg_db" << 'EOF'
|
|
-- Enable pgvector
|
|
CREATE EXTENSION IF NOT EXISTS vector;
|
|
|
|
-- Memories table
|
|
CREATE TABLE IF NOT EXISTS memories (
|
|
id TEXT PRIMARY KEY,
|
|
type TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
embedding vector(384),
|
|
metadata JSONB DEFAULT '{}',
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
|
|
-- Indexes
|
|
CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(type);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_created ON memories(created_at DESC);
|
|
|
|
-- Vector index (IVFFlat for approximate nearest neighbor)
|
|
-- Only create if we have enough rows, otherwise exact search is faster
|
|
DO $$
|
|
BEGIN
|
|
IF NOT EXISTS (SELECT 1 FROM pg_indexes WHERE indexname = 'idx_memories_embedding') THEN
|
|
CREATE INDEX idx_memories_embedding ON memories
|
|
USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
|
|
END IF;
|
|
EXCEPTION WHEN others THEN
|
|
-- Index creation might fail if not enough rows, that's ok
|
|
RAISE NOTICE 'Vector index not created (will use exact search)';
|
|
END $$;
|
|
|
|
SELECT 'Memory schema initialized' as status;
|
|
EOF
|
|
|
|
echo "Database schema ready"
|
|
}
|
|
|
|
install_python
|
|
install_pgvector
|
|
setup_python_env
|
|
download_model
|
|
init_database
|
|
|
|
echo "Memory skill setup complete"
|