Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 2a7942222d stricter mcp lib 2025-12-12 18:40:39 +01:00
Nicolò Boschi 339bd166b1 switch to pg0-embedded 2025-12-12 18:24:16 +01:00
Nicolò Boschi 9fd8c79ee8 switch to pg0-embedded 2025-12-12 18:18:56 +01:00
7 changed files with 80 additions and 415 deletions
-11
View File
@@ -1,11 +0,0 @@
name: 'Setup pg0'
description: 'Install pg0 embedded PostgreSQL'
runs:
using: 'composite'
steps:
- name: Install pg0
shell: bash
run: |
curl -fsSL https://raw.githubusercontent.com/vectorize-io/pg0/main/install.sh | bash
echo "$HOME/.pg0/bin" >> $GITHUB_PATH
-12
View File
@@ -125,9 +125,6 @@ jobs:
with:
python-version-file: ".python-version"
- name: Install pg0
uses: ./.github/actions/setup-pg0
- name: Build API
working-directory: ./hindsight-api
run: uv build
@@ -185,9 +182,6 @@ jobs:
with:
python-version-file: ".python-version"
- name: Install pg0
uses: ./.github/actions/setup-pg0
- name: Build API
working-directory: ./hindsight-api
run: uv build
@@ -269,9 +263,6 @@ jobs:
with:
node-version: '20'
- name: Install pg0
uses: ./.github/actions/setup-pg0
- name: Build API
working-directory: ./hindsight-api
run: uv build
@@ -360,9 +351,6 @@ jobs:
hindsight-clients/rust/target
key: ${{ runner.os }}-cargo-client-${{ hashFiles('hindsight-clients/rust/Cargo.lock') }}
- name: Install pg0
uses: ./.github/actions/setup-pg0
- name: Build API
working-directory: ./hindsight-api
run: uv build
+11 -60
View File
@@ -130,38 +130,10 @@ RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
USER hindsight
# Set PATH for hindsight user
ENV PATH="/home/hindsight/.hindsight/bin:/app/api/.venv/bin:${PATH}"
ENV PATH="/app/api/.venv/bin:${PATH}"
# Install pg0 binary
RUN mkdir -p /home/hindsight/.hindsight/bin && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
PG0_BINARY="pg0-linux-aarch64-gnu"; \
elif [ "$ARCH" = "x86_64" ]; then \
PG0_BINARY="pg0-linux-x86_64-gnu"; \
else \
echo "Unsupported architecture: $ARCH" && exit 1; \
fi && \
echo "Installing pg0 binary: $PG0_BINARY" && \
for i in 1 2 3 4 5; do \
curl -fsSL -o /home/hindsight/.hindsight/bin/pg0 \
"https://github.com/vectorize-io/pg0/releases/latest/download/$PG0_BINARY" && \
chmod +x /home/hindsight/.hindsight/bin/pg0 && \
ls -lh /home/hindsight/.hindsight/bin/pg0 && \
file /home/hindsight/.hindsight/bin/pg0 && \
ldd /home/hindsight/.hindsight/bin/pg0 2>&1 || true && \
break || (echo "Retry $i failed, waiting..." && sleep 10); \
done && \
echo "Testing pg0 binary..." && \
/home/hindsight/.hindsight/bin/pg0 --version || (echo "pg0 --version failed with exit code $?"; ldd /home/hindsight/.hindsight/bin/pg0; exit 1)
# Pre-download PostgreSQL binaries
# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
ENV PG0_HOME=/home/hindsight/.pg0-cache
RUN pg0 start --help && \
(pg0 start --name hindsight --port 5555 --username hindsight --password hindsight --database hindsight && \
sleep 2 && \
pg0 stop --name hindsight && \
echo "PostgreSQL pre-cached to $PG0_HOME") || echo "Pre-download skipped"
ENV PG0_HOME=/home/hindsight/.pg0
@@ -267,38 +239,17 @@ RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
USER hindsight
# Set PATH for hindsight user
ENV PATH="/home/hindsight/.hindsight/bin:/app/api/.venv/bin:${PATH}"
ENV PATH="/app/api/.venv/bin:${PATH}"
# Install pg0 binary
RUN mkdir -p /home/hindsight/.hindsight/bin && \
ARCH=$(uname -m) && \
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then \
PG0_BINARY="pg0-linux-aarch64-gnu"; \
elif [ "$ARCH" = "x86_64" ]; then \
PG0_BINARY="pg0-linux-x86_64-gnu"; \
else \
echo "Unsupported architecture: $ARCH" && exit 1; \
fi && \
echo "Installing pg0 binary: $PG0_BINARY" && \
for i in 1 2 3 4 5; do \
curl -fsSL -o /home/hindsight/.hindsight/bin/pg0 \
"https://github.com/vectorize-io/pg0/releases/latest/download/$PG0_BINARY" && \
chmod +x /home/hindsight/.hindsight/bin/pg0 && \
ls -lh /home/hindsight/.hindsight/bin/pg0 && \
file /home/hindsight/.hindsight/bin/pg0 && \
ldd /home/hindsight/.hindsight/bin/pg0 2>&1 || true && \
break || (echo "Retry $i failed, waiting..." && sleep 10); \
done && \
echo "Testing pg0 binary..." && \
/home/hindsight/.hindsight/bin/pg0 --version || (echo "pg0 --version failed with exit code $?"; ldd /home/hindsight/.hindsight/bin/pg0; exit 1)
# Pre-download PostgreSQL binaries
# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
ENV PG0_HOME=/home/hindsight/.pg0-cache
RUN pg0 start --help && \
(pg0 start --name hindsight --port 5555 --username hindsight --password hindsight --database hindsight && \
sleep 2 && \
pg0 stop --name hindsight && \
echo "PostgreSQL pre-cached to $PG0_HOME") || echo "Pre-download skipped"
RUN /app/api/.venv/bin/python -c "\
from pg0 import Pg0; \
print('Pre-caching PostgreSQL binaries...'); \
pg = Pg0(name='hindsight', port=5555, username='hindsight', password='hindsight', database='hindsight'); \
pg.start(); \
pg.stop(); \
print('PostgreSQL pre-cached to PG0_HOME')" || echo "Pre-download skipped"
ENV PG0_HOME=/home/hindsight/.pg0
+1 -5
View File
@@ -121,11 +121,7 @@ class MCPMiddleware:
self.app = app
self.memory = memory
self.mcp_server = create_mcp_server(memory)
# Use sse_app - http_app requires lifespan management that's complex with middleware
import warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore", DeprecationWarning)
self.mcp_app = self.mcp_server.sse_app()
self.mcp_app = self.mcp_server.http_app()
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
+52 -326
View File
@@ -1,373 +1,114 @@
import asyncio
import json
import logging
import os
import platform
import re
import shutil
import stat
import subprocess
from pathlib import Path
from typing import Optional
import httpx
from pg0 import Pg0
logger = logging.getLogger(__name__)
# pg0 configuration
BINARY_NAME = "pg0"
DEFAULT_PORT = 5555
DEFAULT_USERNAME = "hindsight"
DEFAULT_PASSWORD = "hindsight"
DEFAULT_DATABASE = "hindsight"
def get_platform_binary_name() -> str:
"""Get the appropriate binary name for the current platform.
Supported platforms:
- macOS ARM64 (darwin-aarch64)
- Linux x86_64 (gnu)
- Linux ARM64 (gnu)
- Windows x86_64
"""
system = platform.system().lower()
machine = platform.machine().lower()
# Normalize architecture names
if machine in ("x86_64", "amd64"):
arch = "x86_64"
elif machine in ("arm64", "aarch64"):
arch = "aarch64"
else:
raise RuntimeError(
f"Embedded PostgreSQL is not supported on architecture: {machine}. "
f"Supported architectures: x86_64/amd64 (Linux, Windows), aarch64/arm64 (macOS, Linux)"
)
if system == "darwin" and arch == "aarch64":
return "pg0-darwin-aarch64"
elif system == "linux" and arch == "x86_64":
return "pg0-linux-x86_64-gnu"
elif system == "linux" and arch == "aarch64":
return "pg0-linux-aarch64-gnu"
elif system == "windows" and arch == "x86_64":
return "pg0-windows-x86_64.exe"
else:
raise RuntimeError(
f"Embedded PostgreSQL is not supported on {system}-{arch}. "
f"Supported platforms: darwin-aarch64 (macOS ARM), linux-x86_64-gnu, linux-aarch64-gnu, windows-x86_64"
)
def get_download_url(
version: str = "latest",
repo: str = "vectorize-io/pg0",
) -> str:
"""Get the download URL for pg0 binary."""
binary_name = get_platform_binary_name()
if version == "latest":
return f"https://github.com/{repo}/releases/latest/download/{binary_name}"
else:
return f"https://github.com/{repo}/releases/download/{version}/{binary_name}"
def _find_pg0_binary() -> Optional[Path]:
"""Find pg0 binary in PATH or default install location."""
# First check PATH
pg0_in_path = shutil.which("pg0")
if pg0_in_path:
return Path(pg0_in_path)
# Fall back to default install location
default_path = Path.home() / ".hindsight" / "bin" / "pg0"
if default_path.exists() and os.access(default_path, os.X_OK):
return default_path
return None
class EmbeddedPostgres:
"""
Manages an embedded PostgreSQL server instance using pg0.
This class handles:
- Finding or downloading the pg0 CLI
- Starting/stopping the PostgreSQL server
- Getting the connection URI
Example:
pg = EmbeddedPostgres()
await pg.ensure_installed()
await pg.start()
uri = await pg.get_uri()
# ... use uri with asyncpg ...
await pg.stop()
"""
"""Manages an embedded PostgreSQL server instance using pg0-embedded."""
def __init__(
self,
version: str = "latest",
port: int = DEFAULT_PORT,
username: str = DEFAULT_USERNAME,
password: str = DEFAULT_PASSWORD,
database: str = DEFAULT_DATABASE,
name: str = "hindsight",
**kwargs,
):
"""
Initialize the embedded PostgreSQL manager.
Args:
version: Version of pg0 to download if not found. Defaults to "latest"
port: Port to listen on. Defaults to 5555
username: Username for the database. Defaults to "hindsight"
password: Password for the database. Defaults to "hindsight"
database: Database name to create. Defaults to "hindsight"
name: Instance name for pg0. Defaults to "hindsight"
"""
self.version = version
self.port = port
self.username = username
self.password = password
self.database = database
self.name = name
self._pg0: Optional[Pg0] = None
# Will be set when binary is found/installed
self._binary_path: Optional[Path] = _find_pg0_binary()
@property
def binary_path(self) -> Path:
"""Get the path to the pg0 binary."""
if self._binary_path is None:
# Default install location
return Path.home() / ".hindsight" / "bin" / "pg0"
return self._binary_path
def is_installed(self) -> bool:
"""Check if pg0 is available (in PATH or installed)."""
self._binary_path = _find_pg0_binary()
return self._binary_path is not None
async def ensure_installed(self) -> None:
"""
Ensure pg0 is available.
Checks PATH and default location. If not found, raises an error
instructing the user to install pg0 manually.
"""
if self.is_installed():
logger.debug(f"pg0 found at {self._binary_path}")
return
raise RuntimeError(
"pg0 is not installed. Please install it manually:\n"
" curl -fsSL https://github.com/vectorize-io/pg0/releases/latest/download/pg0-linux-amd64 -o ~/.local/bin/pg0 && chmod +x ~/.local/bin/pg0\n"
"Or visit: https://github.com/vectorize-io/pg0/releases"
)
def _run_command(self, *args: str, capture_output: bool = True) -> subprocess.CompletedProcess:
"""Run a pg0 command synchronously."""
cmd = [str(self.binary_path), *args]
return subprocess.run(cmd, capture_output=capture_output, text=True)
async def _run_command_async(self, *args: str, timeout: int = 120) -> tuple[int, str, str]:
"""Run a pg0 command asynchronously."""
cmd = [str(self.binary_path), *args]
def run_sync():
try:
result = subprocess.run(
cmd,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
timeout=timeout,
)
return result.returncode, result.stdout, result.stderr
except subprocess.TimeoutExpired:
return 1, "", "Command timed out"
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, run_sync)
def _extract_uri_from_output(self, output: str) -> Optional[str]:
"""Extract the PostgreSQL URI from pg0 start output."""
match = re.search(r"Connection URI:\s*(postgresql://[^\s]+)", output)
if match:
return match.group(1)
return None
async def _get_version(self) -> str:
"""Get the pg0 version."""
returncode, stdout, stderr = await self._run_command_async("--version", timeout=10)
if returncode == 0 and stdout:
return stdout.strip()
return "unknown"
def _get_pg0(self) -> Pg0:
if self._pg0 is None:
self._pg0 = Pg0(
name=self.name,
port=self.port,
username=self.username,
password=self.password,
database=self.database,
)
return self._pg0
async def start(self, max_retries: int = 3, retry_delay: float = 2.0) -> str:
"""
Start the PostgreSQL server with retry logic.
Args:
max_retries: Maximum number of start attempts (default: 3)
retry_delay: Initial delay between retries in seconds (default: 2.0)
Returns:
The connection URI for the started server.
Raises:
RuntimeError: If the server fails to start after all retries.
"""
if not self.is_installed():
raise RuntimeError("pg0 is not installed. Call ensure_installed() first.")
# Log pg0 version
version = await self._get_version()
logger.info(f"Starting embedded PostgreSQL with pg0 {version} (name: {self.name}, port: {self.port})...")
"""Start the PostgreSQL server with retry logic."""
logger.info(f"Starting embedded PostgreSQL (name: {self.name}, port: {self.port})...")
pg0 = self._get_pg0()
last_error = None
for attempt in range(1, max_retries + 1):
returncode, stdout, stderr = await self._run_command_async(
"start",
"--name", self.name,
"--port", str(self.port),
"--username", self.username,
"--password", self.password,
"--database", self.database,
timeout=300,
)
# Try to extract URI from output
uri = self._extract_uri_from_output(stdout)
if uri:
logger.info(f"PostgreSQL started on port {self.port}")
return uri
# Check if pg0 info can find the running instance
try:
uri = await self.get_uri()
loop = asyncio.get_event_loop()
info = await loop.run_in_executor(None, pg0.start)
logger.info(f"PostgreSQL started on port {self.port}")
return uri
except RuntimeError:
pass
return info.uri
except Exception as e:
last_error = str(e)
if attempt < max_retries:
delay = retry_delay * (2 ** (attempt - 1))
logger.debug(f"pg0 start attempt {attempt}/{max_retries} failed: {last_error}")
logger.debug(f"Retrying in {delay:.1f}s...")
await asyncio.sleep(delay)
else:
logger.debug(f"pg0 start attempt {attempt}/{max_retries} failed: {last_error}")
# Start failed, log and retry
last_error = stderr or f"pg0 start returned exit code {returncode}"
if attempt < max_retries:
delay = retry_delay * (2 ** (attempt - 1))
logger.debug(f"pg0 start attempt {attempt}/{max_retries} failed: {last_error.strip()}")
logger.debug(f"Retrying in {delay:.1f}s...")
await asyncio.sleep(delay)
else:
logger.debug(f"pg0 start attempt {attempt}/{max_retries} failed: {last_error.strip()}")
# All retries exhausted - fail
raise RuntimeError(
f"Failed to start embedded PostgreSQL after {max_retries} attempts. "
f"Last error: {last_error.strip() if last_error else 'unknown'}"
f"Last error: {last_error}"
)
async def stop(self) -> None:
"""Stop the PostgreSQL server."""
if not self.is_installed():
return
pg0 = self._get_pg0()
logger.info(f"Stopping embedded PostgreSQL (name: {self.name})...")
returncode, stdout, stderr = await self._run_command_async("stop", "--name", self.name)
if returncode != 0:
if "not running" in stderr.lower():
return
raise RuntimeError(f"Failed to stop PostgreSQL: {stderr}")
logger.info("Embedded PostgreSQL stopped")
async def _get_info(self) -> dict:
"""Get info from pg0 using the `info -o json` command."""
if not self.is_installed():
raise RuntimeError("pg0 is not installed.")
returncode, stdout, stderr = await self._run_command_async(
"info", "--name", self.name, "-o", "json"
)
if returncode != 0:
raise RuntimeError(f"Failed to get PostgreSQL info: {stderr}")
try:
return json.loads(stdout.strip())
except json.JSONDecodeError as e:
raise RuntimeError(f"Failed to parse pg0 info output: {e}")
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, pg0.stop)
logger.info("Embedded PostgreSQL stopped")
except Exception as e:
if "not running" in str(e).lower():
return
raise RuntimeError(f"Failed to stop PostgreSQL: {e}")
async def get_uri(self) -> str:
"""Get the connection URI for the PostgreSQL server."""
info = await self._get_info()
uri = info.get("uri")
if not uri:
pg0 = self._get_pg0()
loop = asyncio.get_event_loop()
info = await loop.run_in_executor(None, pg0.info)
if info is None or not info.running:
raise RuntimeError("PostgreSQL server is not running or URI not available")
return uri
async def status(self) -> dict:
"""Get the status of the PostgreSQL server."""
if not self.is_installed():
return {"installed": False, "running": False}
try:
info = await self._get_info()
return {
"installed": True,
"running": info.get("running", False),
"uri": info.get("uri"),
}
except RuntimeError:
return {"installed": True, "running": False}
return info.uri
async def is_running(self) -> bool:
"""Check if the PostgreSQL server is currently running."""
if not self.is_installed():
return False
try:
info = await self._get_info()
return info.get("running", False)
except RuntimeError:
pg0 = self._get_pg0()
loop = asyncio.get_event_loop()
info = await loop.run_in_executor(None, pg0.info)
return info is not None and info.running
except Exception:
return False
async def ensure_running(self) -> str:
"""
Ensure the PostgreSQL server is running.
Installs if needed, starts if not running.
Returns:
The connection URI.
"""
await self.ensure_installed()
"""Ensure the PostgreSQL server is running, starting it if needed."""
if await self.is_running():
return await self.get_uri()
return await self.start()
def uninstall(self) -> None:
"""Remove the pg0 binary (only if we installed it)."""
default_path = Path.home() / ".hindsight" / "bin" / "pg0"
if default_path.exists():
default_path.unlink()
logger.info(f"Removed {default_path}")
def clear_data(self) -> None:
"""Remove all PostgreSQL data (destructive!)."""
result = self._run_command("drop", "--name", self.name, "--force")
if result.returncode == 0:
logger.info(f"Dropped pg0 instance {self.name}")
else:
logger.warning(f"Failed to drop pg0 instance {self.name}: {result.stderr}")
# Convenience functions
_default_instance: Optional[EmbeddedPostgres] = None
@@ -375,33 +116,18 @@ _default_instance: Optional[EmbeddedPostgres] = None
def get_embedded_postgres() -> EmbeddedPostgres:
"""Get or create the default EmbeddedPostgres instance."""
global _default_instance
if _default_instance is None:
_default_instance = EmbeddedPostgres()
return _default_instance
async def start_embedded_postgres() -> str:
"""
Quick start function for embedded PostgreSQL.
Downloads, installs, and starts PostgreSQL in one call.
Returns:
Connection URI string
Example:
db_url = await start_embedded_postgres()
conn = await asyncpg.connect(db_url)
"""
pg = get_embedded_postgres()
return await pg.ensure_running()
"""Quick start function for embedded PostgreSQL."""
return await get_embedded_postgres().ensure_running()
async def stop_embedded_postgres() -> None:
"""Stop the default embedded PostgreSQL instance."""
global _default_instance
if _default_instance:
await _default_instance.stop()
+2 -1
View File
@@ -28,7 +28,8 @@ dependencies = [
"torch>=2.0.0,<2.6.0",
"tiktoken>=0.12.0",
"httpx>=0.27.0",
"fastmcp>=2.0.0",
"fastmcp>=2.3.0",
"pg0-embedded>=0.1.0",
"python-dateutil>=2.8.0",
"opentelemetry-api>=1.20.0",
"opentelemetry-sdk>=1.20.0",
Generated
+14
View File
@@ -1182,6 +1182,7 @@ dependencies = [
{ name = "opentelemetry-exporter-prometheus" },
{ name = "opentelemetry-instrumentation-fastapi" },
{ name = "opentelemetry-sdk" },
{ name = "pg0-embedded" },
{ name = "pgvector" },
{ name = "psycopg2-binary" },
{ name = "pydantic" },
@@ -1233,6 +1234,7 @@ requires-dist = [
{ name = "opentelemetry-exporter-prometheus", specifier = ">=0.41b0" },
{ name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.41b0" },
{ name = "opentelemetry-sdk", specifier = ">=1.20.0" },
{ name = "pg0-embedded", specifier = ">=0.1.0" },
{ name = "pgvector", specifier = ">=0.4.1" },
{ name = "psycopg2-binary", specifier = ">=2.9.11" },
{ name = "pydantic", specifier = ">=2.0.0" },
@@ -2502,6 +2504,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305 },
]
[[package]]
name = "pg0-embedded"
version = "0.10.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/5d/2a/26aed143a5bc4396321c5016c3e7574d596b94122e4de555ec21ebd1f135/pg0_embedded-0.10.1.tar.gz", hash = "sha256:afbfa9e050bec48587d55410e2a93694390c8fb50e1bbab2ac22a36a7eec146d", size = 17619 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/08/61/03f089d9d812e782db200f330e7cf35903834d7f106a2d650bb92c73c8d7/pg0_embedded-0.10.1-py3-none-macosx_14_0_arm64.whl", hash = "sha256:913010ad1a2321367f47cdc907a1639537af36597cd9b61549a341d8da3f249b", size = 13073670 },
{ url = "https://files.pythonhosted.org/packages/72/03/d6e64310c59da880cda4931216f328f72a014c19722fa754b1a0d6422cbb/pg0_embedded-0.10.1-py3-none-manylinux_2_35_aarch64.whl", hash = "sha256:b6f2fc089e844a67dbc1b16899f582ca857744bdd7842b4166a6eecce807a5af", size = 14785516 },
{ url = "https://files.pythonhosted.org/packages/7b/25/a2f84a1c142b48c2a41f14765721650076e2be56762b5ad7cd72ae32e5e4/pg0_embedded-0.10.1-py3-none-manylinux_2_35_x86_64.whl", hash = "sha256:f2ae4ed1ce0aa42a310f20b1ea47dd6091f0f74106b7ded39e9c089e7f80ab25", size = 15224456 },
{ url = "https://files.pythonhosted.org/packages/d0/a8/64963aef0d6ae720b88068441ebdede95b727218c26927e0b8c17d91cf2f/pg0_embedded-0.10.1-py3-none-win_amd64.whl", hash = "sha256:39516c952edc050fbb9e24c35d28c9e013f108879b9c8e91091325231cbdb5a1", size = 54977766 },
]
[[package]]
name = "pgvector"
version = "0.4.1"