chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:10:23 +08:00
commit fbab2c6005
567 changed files with 114434 additions and 0 deletions
+81
View File
@@ -0,0 +1,81 @@
# Backend Rules (api/ + open_notebook/ + commands/ + prompts/)
Normative rules for working on the Python backend. Architecture and design rationale live in [docs/7-DEVELOPMENT/](../docs/7-DEVELOPMENT/index.md) — this file is only what you must know before changing code. Project-wide rules are in the root [AGENTS.md](../AGENTS.md).
## Commands
- Run API: `uv run uvicorn api.main:app --port 5055` (Swagger at http://localhost:5055/docs)
- Background jobs need the worker: `make worker-start` (`surreal-commands-worker --import-modules commands`)
- Tests: `uv run pytest tests/`
- Lint/typecheck: `ruff check . --fix` and `uv run python -m mypy .`
## API layer (`api/`)
- Structure is routes → services → models. Routers stay thin; business logic goes in `*_service.py`.
- Provider metadata (env vars, modalities, test models, discovery URLs, docs links) lives in the registry: `open_notebook/ai/provider_registry.py` `PROVIDERS`. `TEST_MODELS`, `PROVIDER_ENV_CONFIG`, `PROVIDER_MODALITIES` and `OPENAI_COMPAT_PROVIDERS` are derived from it, and `GET /api/providers` exposes it. Adding a provider = add it to the registry, plus **one** manual copy: the `SupportedProvider` Literal in `api/models.py` (typing can't be derived at runtime) — enforced by `tests/test_credential_provider_validation.py`. The frontend consumes `GET /api/providers` at runtime (`useProviders()`), so it needs no edit; the registry declaration order is the display order.
- NEVER return API key values from any endpoint — metadata only.
- Every user-supplied URL field must go through `validate_url()` (`open_notebook/utils/url_validation.py`, async) for SSRF protection. Private IPs/localhost are intentionally allowed (self-hosted Ollama, LM Studio).
- Errors: raise typed exceptions from `open_notebook.exceptions` — global handlers map them to HTTP status codes (`NotFoundError`→404, `InvalidInputError`→400, `AuthenticationError`→401, `RateLimitError`→429, `ConfigurationError`→422, `NetworkError`/`ExternalServiceError`→502, `OpenNotebookError`→500). Don't raise bare `HTTPException` for domain errors.
- Requests over `OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB` (default 100) are rejected by `MaxBodySizeMiddleware` before auth/routing.
- CORS is open by default (`CORS_ORIGINS`); `allow_credentials` flips to `True` only when origins are explicit. No rate limiting built in.
## AI / model provisioning (`open_notebook/ai/`)
- All LLM calls in graph nodes go through `provision_langchain_model()` — never instantiate provider clients directly. It auto-upgrades to `large_context_model` above 105,000 tokens (hard-coded threshold).
- Missing/unconfigured model → raise `ConfigurationError` (not `ValueError`) so the API returns 422.
- Credential-linked models are preferred; `provision_provider_keys()` is the env-var fallback and **mutates `os.environ`** — be aware in tests.
- `DefaultModels.get_instance()` intentionally bypasses the singleton cache (fresh DB fetch each call).
## Graphs (`open_notebook/graphs/`)
- Sync nodes that need async calls use the `asyncio.new_event_loop()` / ThreadPool workaround (see `chat.py`) — fragile, follow the existing pattern exactly.
- Every node wraps LLM calls with `classify_error()`:
```python
except Exception as e:
exc_class, message = classify_error(e)
raise exc_class(message) from e
```
- Strip extended-thinking output with `clean_thinking_content()` before using model responses.
- Chat checkpoints (SqliteSaver) live at the path in `LANGGRAPH_CHECKPOINT_FILE`.
## Domain (`open_notebook/domain/`)
- `Source.save()` does **NOT** auto-embed — call `source.vectorize()` explicitly (fire-and-forget, returns a command id). `Note.save()` DOES auto-submit `embed_note`.
- `ObjectModel.get()` is polymorphic via ID prefix — the subclass must be imported first or resolution fails.
- `RecordModel` subclasses are singletons — call `clear_instance()` in tests.
- Relationship strings passed to `relate()` must match the schema (`reference`, `artifact`, `refers_to`).
## Database (`open_notebook/database/`)
- New migration = new file `open_notebook/database/migrations/N.surrealql` (+ `N_down.surrealql`) **and** an edit to `AsyncMigrationManager` — migrations are hard-coded, not auto-discovered. They run automatically on API startup.
- No connection pooling — each `repo_*` call opens/closes a connection.
- Transaction-conflict `RuntimeError`s are retriable and logged at DEBUG (don't "fix" the missing stack trace).
- Read the `snl-development:surrealdb-queries` skill notes / SurrealDB docs before writing SurrealQL.
## Background commands (`commands/`)
- Retry config uses a blocklist: `stop_on: [ValueError]` — raise `ValueError` for permanent failures (no retry, job marked `failed`); any other exception auto-retries.
- Submission is fire-and-forget via `submit_command()`; commands must be idempotent-ish under retry.
- Podcast generation uses `max_attempts: 1` on purpose (prevents duplicate episodes); retry is the explicit `POST /podcasts/episodes/{id}/retry` endpoint.
## Prompts (`prompts/`)
- Template path syntax: `Prompter(prompt_template="ask/entry")` → `prompts/ask/entry.jinja` (forward slashes, no extension).
- Data is passed as `data=dict`; dict keys must match template variable names exactly.
- With a `PydanticOutputParser`, Prompter auto-injects `format_instructions` — the template must contain `{{ format_instructions }}` or the parser is silently ignored.
- No template inheritance/composition; templates are flat by design.
- Templates are cached — restart the app after editing.
## Environment knobs to know
| Variable | Meaning |
|---|---|
| `OPEN_NOTEBOOK_ENCRYPTION_KEY` (or `_FILE`) | Required for credential storage; any string, no default |
| `OPEN_NOTEBOOK_CHUNK_SIZE` / `_CHUNK_OVERLAP` | Token-based (default 400 / 15%); restart required |
| `OPEN_NOTEBOOK_MAX_UPLOAD_SIZE_MB` | Upload cap (default 100) |
| `LANGGRAPH_CHECKPOINT_FILE` | Chat history SQLite path |
| `CORS_ORIGINS` | Restrict before production |
## Deep dives
[architecture](../docs/7-DEVELOPMENT/architecture.md) · [credentials](../docs/7-DEVELOPMENT/credentials.md) · [content processing](../docs/7-DEVELOPMENT/content-processing.md) · [podcasts](../docs/7-DEVELOPMENT/podcasts.md) · [prompts](../docs/7-DEVELOPMENT/prompts.md) · [change playbooks](../docs/7-DEVELOPMENT/change-playbooks.md) · [testing](../docs/7-DEVELOPMENT/testing.md)
+1
View File
@@ -0,0 +1 @@
@AGENTS.md
View File
+2
View File
@@ -0,0 +1,2 @@
# AI infrastructure module
# Contains model configuration, provisioning, and management
Binary file not shown.
+484
View File
@@ -0,0 +1,484 @@
"""
Connection testing for AI providers.
This module provides functionality to test if a provider's API key is valid
by making minimal API calls to each provider, and to test individual model
configurations end-to-end.
"""
import io
import json
import os
import struct
from typing import Dict, Optional, Tuple
import httpx
from esperanto import (
EmbeddingModel,
LanguageModel,
SpeechToTextModel,
TextToSpeechModel,
)
from esperanto.common_types import ChatCompletion
from loguru import logger
from open_notebook.ai.provider_registry import PROVIDERS
from open_notebook.utils.url_validation import validate_url
def _is_vertex_credentials_file_error(exc: Exception) -> bool:
"""
True if `exc` came from loading a Vertex service-account file
(credentials_path - free text, no path validation; see
open_notebook/ai/key_provider.py).
Google's auth library raises distinguishable exceptions for "file
missing" (FileNotFoundError, an OSError), "not valid JSON"
(json.JSONDecodeError), and "valid JSON but wrong shape"
(google.auth.exceptions.GoogleAuthError) - confirmed by direct
reproduction. Echoing any of these back to an API caller turns
credential/model testing into a filesystem oracle: an attacker who can
create/test a Vertex credential could probe for the existence and
contents-shape of arbitrary files on the server. Callers should catch
these and return one generic message instead of the raw exception text.
Network failures are excluded even though they'd otherwise match
(ConnectionError/TimeoutError are OSError subclasses, TransportError a
GoogleAuthError subclass): they say nothing about the credentials file,
and classifying them here would tell a user with a blocked network to
go debug their file path. Letting them fall through reveals only the
error's category ("connection error"), which keeps the oracle closed.
"""
from google.auth.exceptions import GoogleAuthError, TransportError
if isinstance(exc, (ConnectionError, TimeoutError, TransportError)):
return False
return isinstance(exc, (OSError, json.JSONDecodeError, GoogleAuthError))
# Test models for each provider - uses minimal/cheapest models for testing.
# Derived from the provider registry (the source of truth for test models).
# Format: (model_name, model_type); None model = dynamic (first available).
#
# Prefer a provider-maintained floating alias where one exists, so a model
# retirement doesn't silently break the connection test (see #970: Google
# hard-shuts-down Gemini model ids on a schedule). `gemini-flash-latest`
# is Google's alias for the current stable Flash model and moves forward on
# its own. The provider test also no longer treats a model-level failure as
# a connection failure (see `_connection_failure_reason`), so even if an
# alias ever breaks, the test still reports the credentials correctly.
TEST_MODELS: Dict[str, Tuple[Optional[str], str]] = {
name: (spec.test_model, spec.test_model_type) for name, spec in PROVIDERS.items()
}
async def _test_azure_connection(
endpoint: Optional[str] = None,
api_key: Optional[str] = None,
api_version: Optional[str] = None,
) -> Tuple[bool, str]:
"""
Test Azure OpenAI connectivity by listing models.
Azure requires deployment names which vary per user, so instead of
invoking a model, we list available models to validate credentials.
"""
test_endpoint = endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
test_api_key = api_key or os.environ.get("AZURE_OPENAI_API_KEY")
test_api_version = api_version or os.environ.get("AZURE_OPENAI_API_VERSION", "2024-10-21")
if not test_endpoint:
return False, "No Azure endpoint configured"
if not test_api_key:
return False, "No Azure API key configured"
# Strip trailing slash to avoid double-slash in URL
test_endpoint = test_endpoint.rstrip("/")
try:
# Re-validate at request time: the endpoint may have been saved
# against a hostname that only later resolved to an internal
# address (DNS rebinding), so a save-time check alone isn't enough.
await validate_url(test_endpoint, "azure")
async with httpx.AsyncClient(timeout=10.0) as client:
response = await client.get(
f"{test_endpoint}/openai/models?api-version={test_api_version}",
headers={"api-key": test_api_key},
)
if response.status_code == 200:
data = response.json()
models = data.get("data", [])
count = len(models)
if count > 0:
names = [m.get("id", "unknown") for m in models[:3]]
name_list = ", ".join(names)
if count > 3:
name_list += f" (+{count - 3} more)"
return True, f"Connected. {count} models: {name_list}"
else:
return True, "Connected successfully (no models found)"
elif response.status_code == 401:
return False, "Invalid API key"
elif response.status_code == 403:
return False, "API key lacks required permissions"
else:
return False, f"Azure returned status {response.status_code}"
except ValueError as e:
return False, str(e)
except httpx.ConnectError:
return False, "Cannot connect to Azure endpoint. Check the URL."
except httpx.TimeoutException:
return False, "Connection timed out. Check the endpoint URL."
except Exception as e:
return False, f"Connection error: {str(e)[:100]}"
async def _test_ollama_connection(base_url: str) -> Tuple[bool, str]:
"""Test Ollama server connectivity."""
try:
# Re-validate at request time (see _test_azure_connection for why).
await validate_url(base_url, "ollama")
async with httpx.AsyncClient(timeout=10.0) as client:
# Try /api/tags endpoint (standard Ollama)
response = await client.get(f"{base_url}/api/tags")
if response.status_code == 200:
data = response.json()
models = data.get("models", [])
model_count = len(models)
if model_count > 0:
model_names = [m.get("name", "unknown") for m in models[:3]]
model_list = ", ".join(model_names)
if model_count > 3:
model_list += f" (+{model_count - 3} more)"
return True, f"Connected. {model_count} models available: {model_list}"
else:
return True, "Connected successfully (no models listed)"
elif response.status_code == 401:
return False, "Invalid API key"
elif response.status_code == 403:
return False, "API key lacks required permissions"
else:
return False, f"Server returned status {response.status_code}"
except ValueError as e:
return False, str(e)
except httpx.ConnectError:
return False, "Cannot connect to Ollama. Check if Ollama server is running."
except httpx.TimeoutException:
return False, "Connection timed out. Check if Ollama server is accessible."
except Exception as e:
return False, f"Connection error: {str(e)[:100]}"
async def _test_openai_compatible_connection(base_url: str, api_key: Optional[str] = None) -> Tuple[bool, str]:
"""Test OpenAI-compatible server connectivity."""
try:
# Re-validate at request time (see _test_azure_connection for why).
await validate_url(base_url, "openai_compatible")
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
async with httpx.AsyncClient(timeout=10.0) as client:
# Try /models endpoint (standard OpenAI-compatible)
response = await client.get(f"{base_url}/models", headers=headers)
if response.status_code == 200:
data = response.json()
models = data.get("data", [])
model_count = len(models)
if model_count > 0:
model_names = [m.get("id", "unknown") for m in models[:3]]
model_list = ", ".join(model_names)
if model_count > 3:
model_list += f" (+{model_count - 3} more)"
return True, f"Connected. {model_count} models available: {model_list}"
else:
return True, "Connected successfully (no models listed)"
elif response.status_code == 401:
return False, "Invalid API key"
elif response.status_code == 403:
return False, "API key lacks required permissions"
else:
return False, f"Server returned status {response.status_code}"
except ValueError as e:
return False, str(e)
except httpx.ConnectError:
return False, "Cannot connect to server. Check the URL is correct."
except httpx.TimeoutException:
return False, "Connection timed out. Check if server is accessible."
except Exception as e:
return False, f"Connection error: {str(e)[:100]}"
# Default voices for TTS testing per provider
# ElevenLabs and Mistral excluded: voices looked up dynamically via available_voices
DEFAULT_TEST_VOICES = {
"openai": "alloy",
"azure": "alloy",
"google": "Kore",
"vertex": "Kore",
"openai_compatible": "alloy",
"deepgram": "aura-2-thalia-en",
"xai": "eve",
}
def _generate_test_wav() -> io.BytesIO:
"""Generate a minimal 0.5s silence WAV file in memory (16kHz, 16-bit mono)."""
sample_rate = 16000
num_samples = sample_rate // 2 # 0.5 seconds
bits_per_sample = 16
num_channels = 1
byte_rate = sample_rate * num_channels * bits_per_sample // 8
block_align = num_channels * bits_per_sample // 8
data_size = num_samples * block_align
buf = io.BytesIO()
# RIFF header
buf.write(b"RIFF")
buf.write(struct.pack("<I", 36 + data_size))
buf.write(b"WAVE")
# fmt chunk
buf.write(b"fmt ")
buf.write(struct.pack("<I", 16)) # chunk size
buf.write(struct.pack("<H", 1)) # PCM format
buf.write(struct.pack("<H", num_channels))
buf.write(struct.pack("<I", sample_rate))
buf.write(struct.pack("<I", byte_rate))
buf.write(struct.pack("<H", block_align))
buf.write(struct.pack("<H", bits_per_sample))
# data chunk
buf.write(b"data")
buf.write(struct.pack("<I", data_size))
buf.write(b"\x00" * data_size) # silence
buf.seek(0)
buf.name = "test.wav"
return buf
# A short bundled clip of speech ("Hello there") used to validate STT models.
# Real speech (vs. silence) makes the test transcription non-empty, so a passing
# test visibly returns text instead of a blank result.
_TEST_SPEECH_PATH = os.path.join(os.path.dirname(__file__), "assets", "test_speech.mp3")
def _get_test_audio() -> io.BytesIO:
"""Return a short speech clip for STT testing, or silence as a fallback."""
try:
with open(_TEST_SPEECH_PATH, "rb") as f:
buf = io.BytesIO(f.read())
buf.name = "test_speech.mp3"
buf.seek(0)
return buf
except OSError:
# Fall back to a silent WAV if the bundled clip is missing
return _generate_test_wav()
def _connection_failure_reason(error_msg: str) -> Optional[str]:
"""Classify whether an error means the provider is genuinely unreachable
or the credentials are rejected.
Returns a user-facing failure message for the only errors that actually
disprove a working provider connection — bad key (401), insufficient
permissions (403), and network/timeout failures. Returns None for
anything the provider itself returned *after* authenticating (a missing
or retired model, an unsupported request, a rate limit): reaching the
model layer at all proves the credentials and endpoint work, so those
are not connection failures. This is what keeps a retired test model
(see #970) from being misreported as a broken provider connection.
"""
lower = error_msg.lower()
if "401" in error_msg or "unauthorized" in lower:
return "Invalid API key"
if "403" in error_msg or "forbidden" in lower:
return "API key lacks required permissions"
if "timeout" in lower or "timed out" in lower:
return "Connection timed out - check network/endpoint"
if (
"connection" in lower
or "network" in lower
or "getaddrinfo" in lower
or "name resolution" in lower
or "failed to establish" in lower
):
return "Connection error - check network/endpoint"
return None
def _is_rate_limit(error_msg: str) -> bool:
"""True if the error is a throttling/quota response. Being rate-limited
proves the request authenticated, so callers treat this as connection-OK.
Covers the common phrasings across providers (429, quota, resource
exhausted) rather than just the literal words "rate limit"."""
lower = error_msg.lower()
return (
("rate" in lower and "limit" in lower)
or "429" in error_msg
or "quota" in lower
or "resource has been exhausted" in lower
or "resource exhausted" in lower
)
def _normalize_error_message(error_msg: str) -> Tuple[bool, str]:
"""Normalize common error patterns into user-friendly messages.
Used by the *individual model* test, where the user is validating one
specific registered model — so a missing model IS a failure (unlike the
provider-level test, which only cares that the credentials work).
"""
reason = _connection_failure_reason(error_msg)
if reason:
return False, reason
if _is_rate_limit(error_msg):
return True, "Rate limited - but connection works"
lower = error_msg.lower()
if "not found" in lower and "model" in lower:
return False, "Model not found on this provider"
return False, error_msg
# Substrings that indicate the provider answered but the *test model* is
# missing/retired/unsupported - proof the credentials and endpoint work.
# Only consulted for fixed-endpoint API-key providers (URL-based providers
# are tested via their own handlers), so a "not found" here is about the
# model, never a user-supplied base URL.
_MODEL_UNAVAILABLE_MARKERS = (
"not found",
"not supported",
"does not exist",
"deprecated",
"unavailable",
"no longer available",
)
def classify_provider_test_error(error_msg: str) -> Tuple[bool, str]:
"""Classify a provider connection-test exception into (success, message).
The provider test only asks "do these credentials reach a working
provider?" - so the sole real failures are a rejected key (401),
insufficient permissions (403), and an unreachable endpoint. Anything
the provider returned after authenticating - a rate limit, or a
missing/retired/unsupported test model - still proves the connection
works, so it's reported as success. This is the durable half of the
#970 fix: even if the hard-coded test model is retired, a valid key is
never misreported as a broken connection.
"""
reason = _connection_failure_reason(error_msg)
if reason:
return False, reason
if _is_rate_limit(error_msg):
return True, "Rate limited - but connection works"
lower = error_msg.lower()
if any(marker in lower for marker in _MODEL_UNAVAILABLE_MARKERS):
return True, "API key valid (test model unavailable)"
truncated = error_msg[:100] + "..." if len(error_msg) > 100 else error_msg
return False, f"Error: {truncated}"
async def test_individual_model(model) -> Tuple[bool, str]:
"""
Test a specific model configuration end-to-end by making a real API call.
Args:
model: A Model instance (from open_notebook.ai.models)
Returns:
Tuple of (success: bool, message: str)
"""
from open_notebook.ai.models import ModelManager
try:
manager = ModelManager()
esp_model = await manager.get_model(model.id)
if esp_model is None:
return False, "Could not create model instance"
if model.type == "language":
if not isinstance(esp_model, LanguageModel):
return False, f"Model type mismatch: expected a language model, got {type(esp_model).__name__}"
response = await esp_model.achat_complete(
messages=[{"role": "user", "content": "Hi!"}]
)
if not isinstance(response, ChatCompletion):
# Non-streaming call; a streaming response would be a bug upstream.
return True, "Connection successful (streaming response)"
text = response.content[:100] if response.content else "(empty response)"
return True, f"Response: {text}"
elif model.type == "embedding":
if not isinstance(esp_model, EmbeddingModel):
return False, f"Model type mismatch: expected an embedding model, got {type(esp_model).__name__}"
result = await esp_model.aembed(["This is a test."])
if result and len(result) > 0:
dims = len(result[0])
return True, f"Embedding dimensions: {dims}"
return True, "Embedding successful"
elif model.type == "text_to_speech":
if not isinstance(esp_model, TextToSpeechModel):
return False, f"Model type mismatch: expected a text-to-speech model, got {type(esp_model).__name__}"
# For ElevenLabs, look up first available voice (API uses voice_id, not name)
voice = DEFAULT_TEST_VOICES.get(model.provider)
if not voice and hasattr(esp_model, "available_voices"):
try:
voices = esp_model.available_voices
if voices:
voice = next(iter(voices.keys()))
except Exception:
pass
if not voice:
voice = "alloy" # fallback
audio = await esp_model.agenerate_speech(
text="Hello from Open Notebook", voice=voice
)
if audio and hasattr(audio, "content"):
size = len(audio.content)
return True, f"Audio generated: {size} bytes"
return True, "Speech generation successful"
elif model.type == "speech_to_text":
if not isinstance(esp_model, SpeechToTextModel):
return False, f"Model type mismatch: expected a speech-to-text model, got {type(esp_model).__name__}"
audio_file = _get_test_audio()
transcription = await esp_model.atranscribe(
audio_file=audio_file, language="en"
)
text = (
str(transcription.text).strip()
if hasattr(transcription, "text")
else str(transcription).strip()
)
if not text:
return True, "Connection successful (test clip produced no transcription)"
return True, f"Transcription: {text[:100]}"
else:
return False, f"Unsupported model type: {model.type}"
except Exception as e:
if model.provider == "vertex" and _is_vertex_credentials_file_error(e):
logger.debug(f"Vertex credentials file error for model {model.id}: {e}")
return False, "Invalid or inaccessible credentials file"
error_msg = str(e)
success, normalized = _normalize_error_message(error_msg)
if success:
return True, normalized
logger.debug(f"Test individual model error for {model.id}: {e}")
return False, normalized
+307
View File
@@ -0,0 +1,307 @@
"""
API Key Provider - Database-first with environment fallback.
This module provides a unified interface for retrieving API keys and provider
configuration. It reads from Credential records (individual per-provider
credentials) and falls back to environment variables for backward compatibility.
Usage:
from open_notebook.ai.key_provider import provision_provider_keys
# Call before model provisioning to set env vars from DB
await provision_provider_keys("openai")
"""
import os
from typing import Optional
from loguru import logger
from open_notebook.domain.credential import Credential
# =============================================================================
# Provider Configuration Mapping
# =============================================================================
# Maps provider names to their environment variable names.
# This is the single source of truth for provider-to-env-var mapping.
PROVIDER_CONFIG = {
# Simple providers (just API key)
"openai": {
"env_var": "OPENAI_API_KEY",
},
"anthropic": {
"env_var": "ANTHROPIC_API_KEY",
},
"google": {
"env_var": "GOOGLE_API_KEY",
},
"groq": {
"env_var": "GROQ_API_KEY",
},
"mistral": {
"env_var": "MISTRAL_API_KEY",
},
"deepseek": {
"env_var": "DEEPSEEK_API_KEY",
},
"xai": {
"env_var": "XAI_API_KEY",
},
"openrouter": {
"env_var": "OPENROUTER_API_KEY",
},
"voyage": {
"env_var": "VOYAGE_API_KEY",
},
"elevenlabs": {
"env_var": "ELEVENLABS_API_KEY",
},
"deepgram": {
"env_var": "DEEPGRAM_API_KEY",
},
# URL-based providers
"ollama": {
"env_var": "OLLAMA_API_BASE",
},
"dashscope": {
"env_var": "DASHSCOPE_API_KEY",
},
"minimax": {
"env_var": "MINIMAX_API_KEY",
},
}
async def _get_default_credential(provider: str) -> Optional[Credential]:
"""Get the first credential for a provider from the database."""
try:
credentials = await Credential.get_by_provider(provider)
if credentials:
return credentials[0]
except Exception as e:
logger.debug(f"Could not load credential from database for {provider}: {e}")
return None
async def get_api_key(provider: str) -> Optional[str]:
"""
Get API key for a provider. Checks database first, then env var.
Args:
provider: Provider name (openai, anthropic, etc.)
Returns:
API key string or None if not configured
"""
cred = await _get_default_credential(provider)
if cred and cred.api_key:
logger.debug(f"Using {provider} API key from Credential")
return cred.api_key.get_secret_value()
# Fall back to environment variable
config_info = PROVIDER_CONFIG.get(provider.lower())
if config_info:
env_value = os.environ.get(config_info["env_var"])
if env_value:
logger.debug(f"Using {provider} API key from environment variable")
return env_value
return None
async def _provision_simple_provider(provider: str) -> bool:
"""
Set environment variable for a simple provider from DB config.
Returns:
True if key was set from database, False otherwise
"""
provider_lower = provider.lower()
config_info = PROVIDER_CONFIG.get(provider_lower)
if not config_info:
return False
env_var = config_info["env_var"]
cred = await _get_default_credential(provider_lower)
if not cred:
return False
# Set API key / primary env var
if cred.api_key:
os.environ[env_var] = cred.api_key.get_secret_value()
logger.debug(f"Set {env_var} from Credential")
# Set base URL if present
if cred.base_url:
provider_upper = provider_lower.upper()
os.environ[f"{provider_upper}_API_BASE"] = cred.base_url
logger.debug(f"Set {provider_upper}_API_BASE from Credential")
return True
async def _provision_vertex() -> bool:
"""
Set environment variables for Google Vertex AI from DB config.
Returns:
True if any keys were set from database
"""
any_set = False
cred = await _get_default_credential("vertex")
if not cred:
return False
if cred.project:
os.environ["VERTEX_PROJECT"] = cred.project
logger.debug("Set VERTEX_PROJECT from Credential")
any_set = True
if cred.location:
os.environ["VERTEX_LOCATION"] = cred.location
logger.debug("Set VERTEX_LOCATION from Credential")
any_set = True
if cred.credentials_path:
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = cred.credentials_path
logger.debug("Set GOOGLE_APPLICATION_CREDENTIALS from Credential")
any_set = True
return any_set
async def _provision_azure() -> bool:
"""
Set environment variables for Azure OpenAI from DB config.
Returns:
True if any keys were set from database
"""
any_set = False
cred = await _get_default_credential("azure")
if not cred:
return False
if cred.api_key:
os.environ["AZURE_OPENAI_API_KEY"] = cred.api_key.get_secret_value()
logger.debug("Set AZURE_OPENAI_API_KEY from Credential")
any_set = True
if cred.api_version:
os.environ["AZURE_OPENAI_API_VERSION"] = cred.api_version
logger.debug("Set AZURE_OPENAI_API_VERSION from Credential")
any_set = True
# For Azure, base_url from the UI form maps to endpoint
azure_endpoint = cred.endpoint or cred.base_url
if azure_endpoint:
os.environ["AZURE_OPENAI_ENDPOINT"] = azure_endpoint
logger.debug("Set AZURE_OPENAI_ENDPOINT from Credential")
any_set = True
if cred.endpoint_llm:
os.environ["AZURE_OPENAI_ENDPOINT_LLM"] = cred.endpoint_llm
logger.debug("Set AZURE_OPENAI_ENDPOINT_LLM from Credential")
any_set = True
if cred.endpoint_embedding:
os.environ["AZURE_OPENAI_ENDPOINT_EMBEDDING"] = cred.endpoint_embedding
logger.debug("Set AZURE_OPENAI_ENDPOINT_EMBEDDING from Credential")
any_set = True
if cred.endpoint_stt:
os.environ["AZURE_OPENAI_ENDPOINT_STT"] = cred.endpoint_stt
logger.debug("Set AZURE_OPENAI_ENDPOINT_STT from Credential")
any_set = True
if cred.endpoint_tts:
os.environ["AZURE_OPENAI_ENDPOINT_TTS"] = cred.endpoint_tts
logger.debug("Set AZURE_OPENAI_ENDPOINT_TTS from Credential")
any_set = True
return any_set
async def _provision_openai_compatible() -> bool:
"""
Set environment variables for OpenAI-Compatible providers from DB config.
Returns:
True if any keys were set from database
"""
any_set = False
cred = await _get_default_credential("openai_compatible")
if not cred:
return False
if cred.api_key:
os.environ["OPENAI_COMPATIBLE_API_KEY"] = cred.api_key.get_secret_value()
logger.debug("Set OPENAI_COMPATIBLE_API_KEY from Credential")
any_set = True
if cred.base_url:
os.environ["OPENAI_COMPATIBLE_BASE_URL"] = cred.base_url
logger.debug("Set OPENAI_COMPATIBLE_BASE_URL from Credential")
any_set = True
return any_set
async def provision_provider_keys(provider: str) -> bool:
"""
Provision environment variables from database for a specific provider.
This function checks if the provider has a Credential record stored in the
database and sets the corresponding environment variables. If the database
doesn't have the configuration, existing environment variables remain unchanged.
This is the main entry point for the DB->Env fallback mechanism.
Args:
provider: Provider name (openai, anthropic, azure, vertex,
openai-compatible, etc.)
Returns:
True if any keys were set from database, False otherwise
Example:
# Before provisioning a model, ensure DB keys are in env vars
await provision_provider_keys("openai")
model = AIFactory.create_language(model_name="gpt-4", provider="openai")
"""
# Normalize provider name
provider_lower = provider.lower()
# Handle complex providers with multiple config fields
if provider_lower == "vertex":
return await _provision_vertex()
elif provider_lower == "azure":
return await _provision_azure()
elif provider_lower in ("openai-compatible", "openai_compatible"):
return await _provision_openai_compatible()
# Handle simple providers
return await _provision_simple_provider(provider_lower)
async def provision_all_keys() -> dict[str, bool]:
"""
Provision environment variables from database for all providers.
NOTE: This function is deprecated for request-time use because it can leave
stale env vars after key deletion. Keys should only be provisioned at startup
or via provision_provider_keys() for specific providers.
Useful at application startup to load all DB-stored keys into environment.
Returns:
Dict mapping provider names to whether keys were set from DB
"""
results: dict[str, bool] = {}
# Simple providers
for provider in PROVIDER_CONFIG.keys():
results[provider] = await provision_provider_keys(provider)
# Complex providers
results["vertex"] = await provision_provider_keys("vertex")
results["azure"] = await provision_provider_keys("azure")
results["openai_compatible"] = await provision_provider_keys("openai_compatible")
return results
+768
View File
@@ -0,0 +1,768 @@
"""
Model Discovery - Automatic model fetching from AI providers.
This module provides functionality to discover available models from configured
AI providers and automatically register them in the database.
"""
import asyncio
import os
from dataclasses import dataclass
from typing import Awaitable, Callable, Dict, List, Optional, Tuple
import httpx
from loguru import logger
from open_notebook.ai.models import Model
from open_notebook.ai.provider_registry import PROVIDERS
from open_notebook.database.repository import repo_query
from open_notebook.domain.credential import Credential
@dataclass
class DiscoveredModel:
"""Represents a model discovered from a provider."""
name: str
provider: str
model_type: str # language, embedding, speech_to_text, text_to_speech
description: Optional[str] = None
# =============================================================================
# Provider-Specific Model Type Classification
# =============================================================================
# These mappings help classify models by their capabilities based on naming patterns
OPENAI_MODEL_TYPES = {
"language": [
"gpt-4",
"gpt-3.5",
"o1",
"o3",
"chatgpt",
"text-davinci",
"davinci",
"curie",
"babbage",
"ada",
],
"embedding": ["text-embedding", "embedding"],
"speech_to_text": ["whisper"],
"text_to_speech": ["tts"],
}
# Fallback list used only when Anthropic's model listing API
# (GET https://api.anthropic.com/v1/models) is unreachable or errors.
ANTHROPIC_FALLBACK_MODELS = [
"claude-opus-4-8",
"claude-opus-4-7",
"claude-opus-4-6",
"claude-opus-4-5",
"claude-sonnet-5",
"claude-sonnet-4-6",
"claude-sonnet-4-5",
"claude-haiku-4-5",
]
GOOGLE_MODEL_TYPES = {
"language": ["gemini", "palm", "bison", "chat"],
"embedding": ["embedding", "textembedding"],
# Gemini TTS preview models carry "tts" in the name (checked before language).
# Google STT reuses plain Gemini names and can't be told apart by name, so it
# has no pattern here — users assign the speech_to_text type manually.
"text_to_speech": ["tts"],
}
OLLAMA_MODEL_TYPES = {
# Ollama models can do multiple things, classify by common names
"language": [
"llama",
"mistral",
"mixtral",
"codellama",
"phi",
"gemma",
"qwen",
"deepseek",
"vicuna",
"falcon",
"orca",
"neural",
"dolphin",
"openchat",
"starling",
"solar",
"yi",
"nous",
"wizard",
"zephyr",
"tinyllama",
],
"embedding": ["nomic-embed", "mxbai-embed", "all-minilm", "bge-", "e5-"],
}
MISTRAL_MODEL_TYPES = {
"language": [
"mistral",
"mixtral",
"codestral",
"ministral",
"pixtral",
"open-mistral",
"open-mixtral",
],
"embedding": ["mistral-embed"],
# Voxtral. TTS first by specificity: the "-tts" model must not be caught by
# the broader STT names. classify_model_type checks speech_to_text before
# text_to_speech, so STT patterns are the explicit non-tts model names.
"text_to_speech": ["voxtral-mini-tts", "voxtral-tts"],
"speech_to_text": ["voxtral-mini-latest", "voxtral-small-latest"],
}
GROQ_MODEL_TYPES = {
"language": ["llama", "mixtral", "gemma", "whisper"],
"speech_to_text": ["whisper"],
}
DEEPSEEK_MODEL_TYPES = {
"language": ["deepseek-chat", "deepseek-reasoner", "deepseek-coder"],
}
XAI_MODEL_TYPES = {
"language": ["grok"],
}
VOYAGE_MODEL_TYPES = {
"embedding": ["voyage"],
}
ELEVENLABS_MODEL_TYPES = {
"text_to_speech": ["eleven"],
"speech_to_text": ["scribe"],
}
DEEPGRAM_MODEL_TYPES = {
"text_to_speech": ["aura"],
}
DASHSCOPE_MODEL_TYPES = {
"language": ["qwen"],
}
MINIMAX_MODEL_TYPES = {
"language": ["minimax", "abab"],
}
def classify_model_type(model_name: str, provider: str) -> str:
"""
Classify a model into a type based on its name and provider.
Returns one of: language, embedding, speech_to_text, text_to_speech
"""
name_lower = model_name.lower()
type_mappings = {
"openai": OPENAI_MODEL_TYPES,
"google": GOOGLE_MODEL_TYPES,
"ollama": OLLAMA_MODEL_TYPES,
"mistral": MISTRAL_MODEL_TYPES,
"groq": GROQ_MODEL_TYPES,
"deepseek": DEEPSEEK_MODEL_TYPES,
"xai": XAI_MODEL_TYPES,
"voyage": VOYAGE_MODEL_TYPES,
"elevenlabs": ELEVENLABS_MODEL_TYPES,
"deepgram": DEEPGRAM_MODEL_TYPES,
"dashscope": DASHSCOPE_MODEL_TYPES,
"minimax": MINIMAX_MODEL_TYPES,
}
mapping = type_mappings.get(provider, {})
# Check each type in order of specificity
for model_type in ["speech_to_text", "text_to_speech", "embedding", "language"]:
patterns = mapping.get(model_type, [])
for pattern in patterns:
if pattern in name_lower:
return model_type
# Default to language for unknown models
return "language"
# =============================================================================
# OpenAI-Compatible Provider Discovery (table-driven)
# =============================================================================
# All of these providers expose the same endpoint shape:
# GET {url} with "Authorization: Bearer {key}" -> {"data": [{"id": ...}, ...]}
# Only the URL, the env var holding the key, and small per-provider quirks
# differ, so they share one generic discovery function driven by this table.
def _classify_mistral(model: dict) -> str:
"""Mistral quirk: trust the capabilities flag over name-based patterns."""
if model.get("capabilities", {}).get("completion_chat"):
return "language"
return classify_model_type(model.get("id", ""), "mistral")
@dataclass(frozen=True)
class ProviderDiscoverySpec:
"""Spec for a provider with an OpenAI-compatible /models endpoint."""
url: str
env_var: str
# Optional quirk hooks; defaults are classify_model_type(id, provider)
# and no description.
classify: Optional[Callable[[dict], str]] = None
description: Optional[Callable[[dict], Optional[str]]] = None
# Per-provider quirk hooks that can't live in the (pure data) registry.
_COMPAT_CLASSIFY: Dict[str, Callable[[dict], str]] = {
"mistral": _classify_mistral,
# OpenRouter models are typically language models
"openrouter": lambda model: "language",
}
_COMPAT_DESCRIPTION: Dict[str, Callable[[dict], Optional[str]]] = {
"openrouter": lambda model: model.get("name"),
}
# Built from the provider registry: every provider with an
# `openai_compat_discovery_url` gets a discovery spec. The API key env var
# is the provider's (single) required env var from the registry.
OPENAI_COMPAT_PROVIDERS: Dict[str, ProviderDiscoverySpec] = {
name: ProviderDiscoverySpec(
url=spec.openai_compat_discovery_url,
env_var=spec.required_env[0],
classify=_COMPAT_CLASSIFY.get(name),
description=_COMPAT_DESCRIPTION.get(name),
)
for name, spec in PROVIDERS.items()
if spec.openai_compat_discovery_url
}
async def discover_openai_compatible_provider(provider: str) -> List[DiscoveredModel]:
"""Fetch available models from a provider with an OpenAI-compatible API."""
spec = OPENAI_COMPAT_PROVIDERS[provider]
api_key = os.environ.get(spec.env_var)
if not api_key:
return []
models = []
try:
async with httpx.AsyncClient() as client:
response = await client.get(
spec.url,
headers={"Authorization": f"Bearer {api_key}"},
timeout=30.0,
)
response.raise_for_status()
data = response.json()
for model in data.get("data", []):
model_id = model.get("id", "")
if not model_id:
continue
if spec.classify is not None:
model_type = spec.classify(model)
else:
model_type = classify_model_type(model_id, provider)
description = (
spec.description(model) if spec.description is not None else None
)
models.append(
DiscoveredModel(
name=model_id,
provider=provider,
model_type=model_type,
description=description,
)
)
except Exception as e:
logger.warning(f"Failed to discover {provider} models: {e}")
return models
def _make_openai_compat_discoverer(
provider: str,
) -> Callable[[], Awaitable[List[DiscoveredModel]]]:
async def _discover() -> List[DiscoveredModel]:
return await discover_openai_compatible_provider(provider)
_discover.__name__ = f"discover_{provider}_models"
_discover.__doc__ = f"Fetch available models from the {provider} API."
return _discover
# Kept as module-level names so existing imports/patches keep working.
discover_openai_models = _make_openai_compat_discoverer("openai")
discover_groq_models = _make_openai_compat_discoverer("groq")
discover_mistral_models = _make_openai_compat_discoverer("mistral")
discover_deepseek_models = _make_openai_compat_discoverer("deepseek")
discover_xai_models = _make_openai_compat_discoverer("xai")
discover_openrouter_models = _make_openai_compat_discoverer("openrouter")
discover_dashscope_models = _make_openai_compat_discoverer("dashscope")
discover_minimax_models = _make_openai_compat_discoverer("minimax")
# =============================================================================
# Bespoke Provider Discovery Functions
# =============================================================================
async def fetch_anthropic_model_ids(api_key: str) -> List[str]:
"""
Fetch model ids from Anthropic's model listing API.
Uses GET https://api.anthropic.com/v1/models with pagination
(after_id/has_more cursors). Raises on any HTTP or network error —
callers decide whether to fall back to ANTHROPIC_FALLBACK_MODELS.
"""
model_ids: List[str] = []
params: Dict[str, str] = {"limit": "100"}
async with httpx.AsyncClient() as client:
# Hard page cap as a safety net against a misbehaving cursor.
for _ in range(20):
response = await client.get(
"https://api.anthropic.com/v1/models",
headers={
"x-api-key": api_key,
"anthropic-version": "2023-06-01",
},
params=params,
timeout=30.0,
)
response.raise_for_status()
data = response.json()
for model in data.get("data", []):
model_id = model.get("id", "")
if model_id:
model_ids.append(model_id)
if not data.get("has_more") or not data.get("last_id"):
break
params["after_id"] = data["last_id"]
return model_ids
async def discover_anthropic_models() -> List[DiscoveredModel]:
"""Fetch available models from Anthropic's model listing API."""
api_key = os.environ.get("ANTHROPIC_API_KEY")
if not api_key:
return []
try:
model_names = await fetch_anthropic_model_ids(api_key)
except Exception as e:
logger.warning(
f"Failed to discover Anthropic models, using static fallback: {e}"
)
model_names = list(ANTHROPIC_FALLBACK_MODELS)
return [
DiscoveredModel(
name=model_name,
provider="anthropic",
model_type=classify_model_type(model_name, "anthropic"),
)
for model_name in model_names
]
async def discover_google_models() -> List[DiscoveredModel]:
"""Fetch available models from Google Gemini API."""
api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
if not api_key:
return []
models = []
try:
async with httpx.AsyncClient() as client:
# Build URL without logging the key to avoid exposure
url = "https://generativelanguage.googleapis.com/v1/models"
headers = {"X-Goog-Api-Key": api_key}
response = await client.get(url, headers=headers, timeout=30.0)
response.raise_for_status()
data = response.json()
for model in data.get("models", []):
# Google returns full path like "models/gemini-2.5-flash"
model_name = model.get("name", "").replace("models/", "")
if model_name:
model_type = classify_model_type(model_name, "google")
# Check supported generation methods for better classification
methods = model.get("supportedGenerationMethods", [])
if "embedContent" in methods:
model_type = "embedding"
elif "generateContent" in methods:
model_type = "language"
models.append(
DiscoveredModel(
name=model_name,
provider="google",
model_type=model_type,
description=model.get("displayName"),
)
)
except Exception as e:
# Log without exposing the API key in the message
logger.warning(f"Failed to discover Google models: {type(e).__name__}")
return models
async def discover_ollama_models() -> List[DiscoveredModel]:
"""Fetch available models from local Ollama instance."""
base_url = os.environ.get("OLLAMA_API_BASE", "http://localhost:11434")
if not base_url:
return []
models = []
try:
async with httpx.AsyncClient() as client:
response = await client.get(
f"{base_url}/api/tags",
timeout=10.0,
)
response.raise_for_status()
data = response.json()
for model in data.get("models", []):
model_name = model.get("name", "")
if model_name:
model_type = classify_model_type(model_name, "ollama")
models.append(
DiscoveredModel(
name=model_name,
provider="ollama",
model_type=model_type,
)
)
except Exception as e:
logger.warning(f"Failed to discover Ollama models: {e}")
return models
async def discover_voyage_models() -> List[DiscoveredModel]:
"""Return static list of Voyage AI models (embedding only)."""
api_key = os.environ.get("VOYAGE_API_KEY")
if not api_key:
return []
# Voyage AI specializes in embeddings
voyage_models = [
"voyage-3",
"voyage-3-lite",
"voyage-code-3",
"voyage-finance-2",
"voyage-law-2",
"voyage-multilingual-2",
]
return [
DiscoveredModel(name=m, provider="voyage", model_type="embedding")
for m in voyage_models
]
async def discover_elevenlabs_models() -> List[DiscoveredModel]:
"""Return static list of ElevenLabs TTS models."""
api_key = os.environ.get("ELEVENLABS_API_KEY")
if not api_key:
return []
# ElevenLabs TTS models + the Scribe STT model
elevenlabs_models = [
"eleven_multilingual_v2",
"eleven_turbo_v2_5",
"eleven_turbo_v2",
"eleven_monolingual_v1",
"eleven_multilingual_v1",
]
discovered = [
DiscoveredModel(name=m, provider="elevenlabs", model_type="text_to_speech")
for m in elevenlabs_models
]
discovered.append(
DiscoveredModel(
name="scribe_v1", provider="elevenlabs", model_type="speech_to_text"
)
)
return discovered
async def discover_deepgram_models() -> List[DiscoveredModel]:
"""Return a curated static list of Deepgram Aura TTS voices.
Deepgram has no model-listing API and treats each voice as a model id.
This is a representative subset of the Aura-2 English catalog; users can
add any other voice manually via the custom-model input.
"""
api_key = os.environ.get("DEEPGRAM_API_KEY")
if not api_key:
return []
deepgram_voices = [
"aura-2-thalia-en",
"aura-2-andromeda-en",
"aura-2-helena-en",
"aura-2-apollo-en",
"aura-2-arcas-en",
"aura-2-asteria-en",
"aura-2-athena-en",
"aura-2-hera-en",
"aura-2-hermes-en",
"aura-2-atlas-en",
]
return [
DiscoveredModel(name=m, provider="deepgram", model_type="text_to_speech")
for m in deepgram_voices
]
async def discover_openai_compatible_models() -> List[DiscoveredModel]:
"""
Fetch available models from an OpenAI-compatible API endpoint.
Uses the configured base_url from the database or environment variable.
"""
api_key = None
base_url = None
# Try to get config from Credential database first
try:
credentials = await Credential.get_by_provider("openai_compatible")
if credentials:
cred = credentials[0]
config = cred.to_esperanto_config()
api_key = config.get("api_key")
base_url = config.get("base_url", "").rstrip("/")
except Exception as e:
logger.warning(f"Failed to read openai_compatible config from Credential: {e}")
# Fall back to environment variables
if not api_key:
api_key = os.environ.get("OPENAI_COMPATIBLE_API_KEY")
if not base_url:
base_url = os.environ.get("OPENAI_COMPATIBLE_BASE_URL", "").rstrip("/")
if not base_url:
logger.warning("No base_url configured for openai_compatible provider")
return []
models = []
try:
async with httpx.AsyncClient() as client:
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
response = await client.get(
f"{base_url}/models",
headers=headers,
timeout=30.0,
)
response.raise_for_status()
data = response.json()
for model in data.get("data", []):
model_id = model.get("id", "")
if model_id:
# Classify based on model name patterns
model_type = classify_model_type(model_id, "openai")
models.append(
DiscoveredModel(
name=model_id,
provider="openai_compatible",
model_type=model_type,
)
)
except httpx.HTTPStatusError as e:
logger.warning(f"Failed to discover openai_compatible models: HTTP {e.response.status_code}")
except Exception as e:
logger.warning(f"Failed to discover openai_compatible models: {e}")
return models
# =============================================================================
# Main Discovery Functions
# =============================================================================
# Map provider names to their discovery functions
PROVIDER_DISCOVERY_FUNCTIONS = {
"openai": discover_openai_models,
"anthropic": discover_anthropic_models,
"google": discover_google_models,
"ollama": discover_ollama_models,
"groq": discover_groq_models,
"mistral": discover_mistral_models,
"deepseek": discover_deepseek_models,
"xai": discover_xai_models,
"openrouter": discover_openrouter_models,
"voyage": discover_voyage_models,
"elevenlabs": discover_elevenlabs_models,
"deepgram": discover_deepgram_models,
"openai_compatible": discover_openai_compatible_models,
"dashscope": discover_dashscope_models,
"minimax": discover_minimax_models,
"azure": None, # Azure requires credential-based discovery (different auth)
"vertex": None, # Vertex requires credential-based discovery (service account)
}
async def discover_provider_models(provider: str) -> List[DiscoveredModel]:
"""
Discover available models for a specific provider.
Args:
provider: Provider name (openai, anthropic, etc.)
Returns:
List of discovered models
"""
discover_func = PROVIDER_DISCOVERY_FUNCTIONS.get(provider)
if discover_func is None:
if provider in PROVIDER_DISCOVERY_FUNCTIONS:
logger.info(
f"Provider '{provider}' requires credential-based discovery. "
f"Use the /credentials/{{id}}/discover endpoint instead."
)
else:
logger.warning(f"No discovery function for provider: {provider}")
return []
return await discover_func()
async def sync_provider_models(
provider: str, auto_register: bool = True
) -> Tuple[int, int, int]:
"""
Sync models for a provider: discover and optionally register in database.
Args:
provider: Provider name
auto_register: If True, automatically create Model records in database
Returns:
Tuple of (discovered_count, new_count, existing_count)
"""
discovered = await discover_provider_models(provider)
discovered_count = len(discovered)
new_count = 0
existing_count = 0
if not auto_register:
return discovered_count, 0, 0
if not discovered:
return 0, 0, 0
# Batch fetch existing models to avoid N+1 query pattern
try:
existing_models = await repo_query(
"SELECT string::lowercase(name) as name, string::lowercase(type) as type FROM model "
"WHERE string::lowercase(provider) = $provider",
{"provider": provider.lower()},
)
# Create a set of (name, type) tuples for O(1) lookup
existing_keys = set()
for m in existing_models:
existing_keys.add((m.get("name", ""), m.get("type", "")))
except Exception as e:
logger.warning(f"Failed to fetch existing models for {provider}: {e}")
existing_keys = set()
for model in discovered:
model_key = (model.name.lower(), model.model_type.lower())
# Check if model already exists using pre-fetched data
if model_key in existing_keys:
existing_count += 1
continue
# Create new model
try:
new_model = Model(
name=model.name,
provider=model.provider,
type=model.model_type,
)
await new_model.save()
new_count += 1
logger.info(f"Registered new model: {model.provider}/{model.name} ({model.model_type})")
except Exception as e:
logger.warning(f"Failed to register model {model.name}: {e}")
logger.info(
f"Synced {provider}: {discovered_count} discovered, "
f"{new_count} new, {existing_count} existing"
)
return discovered_count, new_count, existing_count
async def sync_all_providers() -> Dict[str, Tuple[int, int, int]]:
"""
Sync models for all configured providers.
Returns:
Dict mapping provider names to (discovered, new, existing) tuples
"""
results = {}
# Run discovery for all providers in parallel
tasks = []
providers = list(PROVIDER_DISCOVERY_FUNCTIONS.keys())
for provider in providers:
tasks.append(sync_provider_models(provider, auto_register=True))
task_results = await asyncio.gather(*tasks, return_exceptions=True)
for provider, result in zip(providers, task_results):
if isinstance(result, BaseException):
logger.error(f"Error syncing {provider}: {result}")
results[provider] = (0, 0, 0)
else:
results[provider] = result
return results
async def get_provider_model_count(provider: str) -> Dict[str, int]:
"""
Get count of registered models for a provider, grouped by type.
Args:
provider: Provider name (case-insensitive)
Returns:
Dict mapping model type to count
"""
# Use case-insensitive comparison by lowercasing the provider
result = await repo_query(
"SELECT type, count() as count FROM model WHERE string::lowercase(provider) = string::lowercase($provider) GROUP BY type",
{"provider": provider},
)
counts = {
"language": 0,
"embedding": 0,
"speech_to_text": 0,
"text_to_speech": 0,
}
for row in result:
model_type = row.get("type")
count = row.get("count", 0)
if model_type in counts:
counts[model_type] = count
return counts
+341
View File
@@ -0,0 +1,341 @@
from typing import Any, ClassVar, Dict, Optional, Sequence, Union
from esperanto import (
AIFactory,
EmbeddingModel,
LanguageModel,
SpeechToTextModel,
TextToSpeechModel,
)
from loguru import logger
from surrealdb import RecordID
from open_notebook.database.repository import ensure_record_id, repo_query
from open_notebook.domain.base import ObjectModel, RecordModel
from open_notebook.exceptions import ConfigurationError
from open_notebook.utils.url_validation import validate_url
ModelType = Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel]
# Config keys from Credential.to_esperanto_config() that may carry a
# user-configured URL (ollama/azure/openai_compatible/vertex).
_URL_CONFIG_KEYS = (
"base_url",
"endpoint",
"endpoint_llm",
"endpoint_embedding",
"endpoint_stt",
"endpoint_tts",
)
async def _revalidate_config_urls(config: dict, provider: str) -> None:
"""
Re-validate a credential's URL fields immediately before they're used for
a real request.
validate_url() is also enforced when a credential is created/updated, but
that alone leaves a DNS-rebinding TOCTOU window: a hostname that resolved
to a public IP at save time can later be repointed to an internal/
metadata address, and Esperanto/httpx re-resolve DNS fresh on every
connection. Re-checking here narrows that window to "this call", instead
of "any time after the credential was saved".
"""
for key in _URL_CONFIG_KEYS:
value = config.get(key)
if value:
try:
await validate_url(value, provider)
except ValueError as e:
raise ConfigurationError(str(e)) from e
class Model(ObjectModel):
table_name: ClassVar[str] = "model"
nullable_fields: ClassVar[set[str]] = {"credential"}
name: str
provider: str
type: str
credential: Optional[str] = None
@classmethod
async def get_models_by_type(cls, model_type):
models = await repo_query(
"SELECT * FROM model WHERE type=$model_type;", {"model_type": model_type}
)
return [Model(**model) for model in models]
@classmethod
async def get_display_info_for_ids(
cls, model_ids: Sequence[Union[str, RecordID]]
) -> Dict[str, Dict[str, str]]:
"""
Batch-fetch {provider, name} display info for many model IDs in one
query.
Episode listing resolves the model references stored in the
denormalized episode/speaker profile snapshots (outline_llm,
transcript_llm, voice_model) into human-readable display fields.
Doing that with Model.get() would cost one round trip per reference
per episode (no connection pooling in the repository layer) - this
collects the distinct IDs and resolves them in a single query,
mirroring PodcastEpisode.get_job_details_for_commands().
Unresolvable IDs (deleted models) are simply absent from the result;
a total query failure returns an empty dict so display resolution
degrades gracefully instead of breaking the caller.
"""
ids = sorted({str(mid) for mid in model_ids if mid})
grouped: Dict[str, Dict[str, str]] = {}
if not ids:
return grouped
try:
result = await repo_query(
"SELECT id, name, provider FROM model WHERE id IN $model_ids",
{"model_ids": [ensure_record_id(mid) for mid in ids]},
)
except Exception as e:
logger.error(f"Error batch-fetching model display info: {e}")
return grouped
for row in result:
grouped[str(row.get("id"))] = {
"provider": row.get("provider", ""),
"name": row.get("name", ""),
}
return grouped
@classmethod
async def get_by_credential(cls, credential_id: str):
"""Get all models linked to a specific credential."""
models = await repo_query(
"SELECT * FROM model WHERE credential=$cred_id;",
{"cred_id": ensure_record_id(credential_id)},
)
return [Model(**model) for model in models]
def _prepare_save_data(self) -> Dict[str, Any]:
data = super()._prepare_save_data()
if data.get("credential"):
data["credential"] = ensure_record_id(data["credential"])
return data
async def get_credential_obj(self):
"""Get the Credential object linked to this model, if any."""
if not self.credential:
return None
from open_notebook.domain.credential import Credential
try:
return await Credential.get(self.credential)
except Exception:
logger.warning(f"Could not load credential {self.credential} for model {self.id}")
return None
class DefaultModels(RecordModel):
record_id: ClassVar[str] = "open_notebook:default_models"
default_chat_model: Optional[str] = None
default_transformation_model: Optional[str] = None
large_context_model: Optional[str] = None
default_text_to_speech_model: Optional[str] = None
default_speech_to_text_model: Optional[str] = None
# default_vision_model: Optional[str]
default_embedding_model: Optional[str] = None
default_tools_model: Optional[str] = None
@classmethod
async def get_instance(cls) -> "DefaultModels":
"""Always fetch fresh defaults from database (override parent caching behavior)"""
result = await repo_query(
"SELECT * FROM ONLY $record_id",
{"record_id": ensure_record_id(cls.record_id)},
)
if result:
if isinstance(result, list) and len(result) > 0:
data = result[0]
elif isinstance(result, dict):
data = result
else:
data = {}
else:
data = {}
# Create new instance with fresh data (bypass singleton cache)
instance = object.__new__(cls)
object.__setattr__(instance, "__dict__", {})
super(RecordModel, instance).__init__(**data)
return instance
class ModelManager:
def __init__(self):
pass # No caching needed
async def get_model(self, model_id: str, **kwargs) -> Optional[ModelType]:
"""Get a model by ID. Esperanto will cache the actual model instance."""
if not model_id:
return None
try:
model: Model = await Model.get(model_id)
except Exception:
raise ConfigurationError(f"Model with ID {model_id} not found")
if not model.type or model.type not in [
"language",
"embedding",
"speech_to_text",
"text_to_speech",
]:
raise ConfigurationError(f"Invalid model type: {model.type}")
# Build config from credential if linked, otherwise fall back to env vars
config: dict = {}
if model.credential:
credential = await model.get_credential_obj()
if credential:
config = credential.to_esperanto_config()
await _revalidate_config_urls(config, model.provider)
logger.debug(
f"Using credential '{credential.name}' for model {model.name}"
)
else:
logger.warning(
f"Model {model.id} has credential {model.credential} but it could not be loaded. "
f"Falling back to env vars."
)
# Fall back to env var provisioning
from open_notebook.ai.key_provider import provision_provider_keys
await provision_provider_keys(model.provider)
else:
# No credential linked - use env var fallback
from open_notebook.ai.key_provider import provision_provider_keys
await provision_provider_keys(model.provider)
# Merge any additional kwargs (e.g. temperature)
config.update(kwargs)
# Normalize provider name: DB stores underscores but Esperanto expects hyphens
provider = model.provider.replace("_", "-")
# Create model based on type (Esperanto will cache the instance)
if model.type == "language":
return AIFactory.create_language(
model_name=model.name,
provider=provider,
config=config,
)
elif model.type == "embedding":
return AIFactory.create_embedding(
model_name=model.name,
provider=provider,
config=config,
)
elif model.type == "speech_to_text":
return AIFactory.create_speech_to_text(
model_name=model.name,
provider=provider,
config=config,
)
elif model.type == "text_to_speech":
return AIFactory.create_text_to_speech(
model_name=model.name,
provider=provider,
config=config,
)
else:
raise ConfigurationError(f"Invalid model type: {model.type}")
async def get_defaults(self) -> DefaultModels:
"""Get the default models configuration from database"""
defaults = await DefaultModels.get_instance()
if not defaults:
raise RuntimeError("Failed to load default models configuration")
return defaults
async def get_speech_to_text(self, **kwargs) -> Optional[SpeechToTextModel]:
"""Get the default speech-to-text model"""
defaults = await self.get_defaults()
model_id = defaults.default_speech_to_text_model
if not model_id:
return None
model = await self.get_model(model_id, **kwargs)
assert model is None or isinstance(model, SpeechToTextModel), (
f"Expected SpeechToTextModel but got {type(model)}"
)
return model
async def get_text_to_speech(self, **kwargs) -> Optional[TextToSpeechModel]:
"""Get the default text-to-speech model"""
defaults = await self.get_defaults()
model_id = defaults.default_text_to_speech_model
if not model_id:
return None
model = await self.get_model(model_id, **kwargs)
assert model is None or isinstance(model, TextToSpeechModel), (
f"Expected TextToSpeechModel but got {type(model)}"
)
return model
async def get_embedding_model(self, **kwargs) -> Optional[EmbeddingModel]:
"""Get the default embedding model"""
defaults = await self.get_defaults()
model_id = defaults.default_embedding_model
if not model_id:
return None
model = await self.get_model(model_id, **kwargs)
assert model is None or isinstance(model, EmbeddingModel), (
f"Expected EmbeddingModel but got {type(model)}"
)
return model
async def get_default_model(self, model_type: str, **kwargs) -> Optional[ModelType]:
"""
Get the default model for a specific type.
Args:
model_type: The type of model to retrieve (e.g., 'chat', 'embedding', etc.)
**kwargs: Additional arguments to pass to the model constructor
"""
defaults = await self.get_defaults()
model_id = None
if model_type == "chat":
model_id = defaults.default_chat_model
elif model_type == "transformation":
model_id = (
defaults.default_transformation_model or defaults.default_chat_model
)
elif model_type == "tools":
model_id = defaults.default_tools_model or defaults.default_chat_model
elif model_type == "embedding":
model_id = defaults.default_embedding_model
elif model_type == "text_to_speech":
model_id = defaults.default_text_to_speech_model
elif model_type == "speech_to_text":
model_id = defaults.default_speech_to_text_model
elif model_type == "large_context":
model_id = defaults.large_context_model
if not model_id:
logger.warning(
f"No default model configured for type '{model_type}'. "
f"Please go to Settings → Models and set a default model."
)
return None
try:
return await self.get_model(model_id, **kwargs)
except (ValueError, ConfigurationError) as e:
logger.error(
f"Failed to load default model for type '{model_type}': {e}. "
f"The configured model_id '{model_id}' may have been deleted or misconfigured. "
f"Please go to Settings → Models and reconfigure the default model."
)
return None
model_manager = ModelManager()
+249
View File
@@ -0,0 +1,249 @@
"""
Provider Registry — the single source of truth for AI provider metadata.
Adding a provider used to require keeping ~6 independent dicts in sync
(env config, modalities, test models, discovery table, API Literal,
frontend table). Now the backend surfaces are all derived from the
`PROVIDERS` registry below:
- `api/credentials_service.py` `PROVIDER_ENV_CONFIG` / `PROVIDER_MODALITIES`
- `open_notebook/ai/connection_tester.py` `TEST_MODELS`
- `open_notebook/ai/model_discovery.py` `OPENAI_COMPAT_PROVIDERS`
- `GET /api/providers` (api/routers/providers.py)
One place still needs a manual edit when adding a provider — enforced by
tests (tests/test_credential_provider_validation.py): the
`SupportedProvider` Literal in `api/models.py` (typing can't be built at
runtime from this dict). The frontend consumes `GET /api/providers` at
runtime, so it needs no edit.
The declaration order below is the display order the frontend renders —
`GET /api/providers` returns `PROVIDERS.values()` as declared.
This module is pure data: it must not import anything from the rest of
the project so it stays importable from anywhere without cycles.
"""
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
@dataclass(frozen=True)
class ProviderSpec:
"""Everything the backend needs to know about one AI provider."""
name: str
display_name: str
# Default modalities offered when creating a credential for this provider.
modalities: Tuple[str, ...]
# Env var configuration for env-based setup/migration:
# - required_env: ALL must be set for the provider to count as configured.
# - required_any_env: at least ONE must be set.
# - optional_env: read during migration but not required.
required_env: Tuple[str, ...] = ()
required_any_env: Tuple[str, ...] = ()
optional_env: Tuple[str, ...] = ()
# Cheapest model used by the provider connection test. None means the
# test is dynamic (first available model) or handled by a bespoke tester.
test_model: Optional[str] = None
test_model_type: str = "language"
# Where users get an API key / set the provider up.
docs_url: Optional[str] = None
# For providers exposing an OpenAI-compatible GET /models endpoint,
# the discovery URL. Drives OPENAI_COMPAT_PROVIDERS in model_discovery.
openai_compat_discovery_url: Optional[str] = None
def env_config(self) -> Dict[str, List[str]]:
"""Env var config in the legacy PROVIDER_ENV_CONFIG dict shape."""
config: Dict[str, List[str]] = {}
if self.required_env:
config["required"] = list(self.required_env)
if self.required_any_env:
config["required_any"] = list(self.required_any_env)
if self.optional_env:
config["optional"] = list(self.optional_env)
return config
_LANGUAGE_ONLY = ("language",)
_ALL_MODALITIES = ("language", "embedding", "speech_to_text", "text_to_speech")
_PROVIDER_SPECS: Tuple[ProviderSpec, ...] = (
ProviderSpec(
name="openai",
display_name="OpenAI",
modalities=_ALL_MODALITIES,
required_env=("OPENAI_API_KEY",),
test_model="gpt-3.5-turbo",
docs_url="https://platform.openai.com/api-keys",
openai_compat_discovery_url="https://api.openai.com/v1/models",
),
ProviderSpec(
name="anthropic",
display_name="Anthropic",
modalities=_LANGUAGE_ONLY,
required_env=("ANTHROPIC_API_KEY",),
test_model="claude-3-haiku-20240307",
docs_url="https://console.anthropic.com/settings/keys",
),
ProviderSpec(
name="google",
display_name="Google AI",
modalities=_ALL_MODALITIES,
required_any_env=("GOOGLE_API_KEY", "GEMINI_API_KEY"),
test_model="gemini-flash-latest",
docs_url="https://aistudio.google.com/app/apikey",
),
ProviderSpec(
name="groq",
display_name="Groq",
modalities=("language", "speech_to_text"),
required_env=("GROQ_API_KEY",),
test_model="llama-3.1-8b-instant",
docs_url="https://console.groq.com/keys",
openai_compat_discovery_url="https://api.groq.com/openai/v1/models",
),
ProviderSpec(
name="mistral",
display_name="Mistral AI",
modalities=("language", "embedding", "speech_to_text", "text_to_speech"),
required_env=("MISTRAL_API_KEY",),
test_model="mistral-small-latest",
docs_url="https://console.mistral.ai/api-keys/",
openai_compat_discovery_url="https://api.mistral.ai/v1/models",
),
ProviderSpec(
name="deepseek",
display_name="DeepSeek",
modalities=_LANGUAGE_ONLY,
required_env=("DEEPSEEK_API_KEY",),
test_model="deepseek-chat",
docs_url="https://platform.deepseek.com/api_keys",
openai_compat_discovery_url="https://api.deepseek.com/models",
),
ProviderSpec(
name="xai",
display_name="xAI (Grok)",
modalities=("language", "text_to_speech"),
required_env=("XAI_API_KEY",),
test_model="grok-beta",
docs_url="https://console.x.ai/",
openai_compat_discovery_url="https://api.x.ai/v1/models",
),
ProviderSpec(
name="openrouter",
display_name="OpenRouter",
modalities=("language", "embedding"),
required_env=("OPENROUTER_API_KEY",),
test_model="openai/gpt-3.5-turbo",
docs_url="https://openrouter.ai/keys",
openai_compat_discovery_url="https://openrouter.ai/api/v1/models",
),
ProviderSpec(
name="dashscope",
display_name="DashScope (Qwen)",
modalities=_LANGUAGE_ONLY,
required_env=("DASHSCOPE_API_KEY",),
test_model="qwen-plus",
docs_url="https://help.aliyun.com/zh/model-studio/getting-started/",
openai_compat_discovery_url="https://dashscope.aliyuncs.com/compatible-mode/v1/models",
),
ProviderSpec(
name="minimax",
display_name="MiniMax",
modalities=_LANGUAGE_ONLY,
required_env=("MINIMAX_API_KEY",),
test_model="MiniMax-M2.5",
docs_url="https://platform.minimaxi.com/document/Guides",
openai_compat_discovery_url="https://api.minimax.io/v1/models",
),
ProviderSpec(
name="voyage",
display_name="Voyage AI",
modalities=("embedding",),
required_env=("VOYAGE_API_KEY",),
test_model="voyage-3-lite",
test_model_type="embedding",
docs_url="https://dash.voyageai.com/api-keys",
),
ProviderSpec(
name="elevenlabs",
display_name="ElevenLabs",
modalities=("text_to_speech", "speech_to_text"),
required_env=("ELEVENLABS_API_KEY",),
test_model="eleven_multilingual_v2",
test_model_type="text_to_speech",
docs_url="https://elevenlabs.io/app/settings/api-keys",
),
ProviderSpec(
name="deepgram",
display_name="Deepgram",
modalities=("text_to_speech",),
required_env=("DEEPGRAM_API_KEY",),
test_model="aura-2-thalia-en",
test_model_type="text_to_speech",
docs_url="https://console.deepgram.com/",
),
ProviderSpec(
name="ollama",
display_name="Ollama",
modalities=("language", "embedding"),
required_env=("OLLAMA_API_BASE",),
test_model=None, # Dynamic - uses first available model
),
ProviderSpec(
name="azure",
display_name="Azure OpenAI",
modalities=_ALL_MODALITIES,
required_env=(
"AZURE_OPENAI_API_KEY",
"AZURE_OPENAI_ENDPOINT",
"AZURE_OPENAI_API_VERSION",
),
optional_env=(
"AZURE_OPENAI_ENDPOINT_LLM",
"AZURE_OPENAI_ENDPOINT_EMBEDDING",
"AZURE_OPENAI_ENDPOINT_STT",
"AZURE_OPENAI_ENDPOINT_TTS",
),
test_model="gpt-35-turbo", # Azure OpenAI deployment name
docs_url="https://portal.azure.com/#view/Microsoft_Azure_ProjectOxford/CognitiveServicesHub/~/OpenAI",
),
ProviderSpec(
name="vertex",
display_name="Google Vertex AI",
modalities=("language", "embedding", "text_to_speech"),
required_env=("VERTEX_PROJECT", "VERTEX_LOCATION"),
optional_env=("GOOGLE_APPLICATION_CREDENTIALS",),
test_model="gemini-flash-latest", # Uses Google Vertex AI
docs_url="https://cloud.google.com/vertex-ai/docs/start/cloud-environment",
),
ProviderSpec(
name="openai_compatible",
display_name="OpenAI Compatible",
modalities=_ALL_MODALITIES,
required_any_env=("OPENAI_COMPATIBLE_BASE_URL", "OPENAI_COMPATIBLE_API_KEY"),
test_model=None, # Dynamic - uses first available model
docs_url="https://github.com/lfnovo/open-notebook/blob/main/docs/5-CONFIGURATION/openai-compatible.md",
),
)
def _build_registry(specs: Tuple[ProviderSpec, ...]) -> Dict[str, ProviderSpec]:
"""Build the name -> spec map, refusing duplicate names at import time.
A plain dict comprehension would silently drop the earlier spec on a
name collision; fail loudly instead.
"""
registry: Dict[str, ProviderSpec] = {}
for spec in specs:
if spec.name in registry:
raise ValueError(
f"Duplicate provider name in registry: {spec.name!r}"
)
registry[spec.name] = spec
return registry
PROVIDERS: Dict[str, ProviderSpec] = _build_registry(_PROVIDER_SPECS)
+61
View File
@@ -0,0 +1,61 @@
from esperanto import LanguageModel
from langchain_core.language_models.chat_models import BaseChatModel
from loguru import logger
from open_notebook.ai.models import model_manager
from open_notebook.exceptions import ConfigurationError
from open_notebook.utils import token_count
async def provision_langchain_model(
content, model_id, default_type, **kwargs
) -> BaseChatModel:
"""
Returns the best model to use based on the context size and on whether there is a specific model being requested in Config.
If context > 105_000, returns the large_context_model
If model_id is specified in Config, returns that model
Otherwise, returns the default model for the given type
"""
tokens = token_count(content)
model = None
selection_reason = ""
if tokens > 105_000:
selection_reason = f"large_context (content has {tokens} tokens)"
logger.debug(
f"Using large context model because the content has {tokens} tokens"
)
model = await model_manager.get_default_model("large_context", **kwargs)
elif model_id:
selection_reason = f"explicit model_id={model_id}"
model = await model_manager.get_model(model_id, **kwargs)
else:
selection_reason = f"default for type={default_type}"
model = await model_manager.get_default_model(default_type, **kwargs)
logger.debug(f"Using model: {model}")
if model is None:
logger.error(
f"Model provisioning failed: No model found. "
f"Selection reason: {selection_reason}. "
f"model_id={model_id}, default_type={default_type}. "
f"Please check Settings → Models and ensure a default model is configured for '{default_type}'."
)
raise ConfigurationError(
f"No model configured for {selection_reason}. "
f"Please go to Settings → Models and configure a default model for '{default_type}'."
)
if not isinstance(model, LanguageModel):
logger.error(
f"Model type mismatch: Expected LanguageModel but got {type(model).__name__}. "
f"Selection reason: {selection_reason}. "
f"model_id={model_id}, default_type={default_type}."
)
raise ConfigurationError(
f"Model is not a LanguageModel: {model}. "
f"Please check that the model configured for '{default_type}' is a language model, not an embedding or speech model."
)
return model.to_langchain()
+26
View File
@@ -0,0 +1,26 @@
import os
# ROOT DATA FOLDER
DATA_FOLDER = "./data"
# LANGGRAPH CHECKPOINT FILE
sqlite_folder = f"{DATA_FOLDER}/sqlite-db"
os.makedirs(sqlite_folder, exist_ok=True)
LANGGRAPH_CHECKPOINT_FILE = f"{sqlite_folder}/checkpoints.sqlite"
# UPLOADS FOLDER
UPLOADS_FOLDER = f"{DATA_FOLDER}/uploads"
os.makedirs(UPLOADS_FOLDER, exist_ok=True)
# PODCASTS FOLDER
# Matches the root that build_episode_output_dir() (commands/podcast_commands.py)
# creates episode directories under when called with DATA_FOLDER in production.
PODCASTS_FOLDER = f"{DATA_FOLDER}/podcasts"
os.makedirs(PODCASTS_FOLDER, exist_ok=True)
# TIKTOKEN CACHE FOLDER
# Reads TIKTOKEN_CACHE_DIR from the environment so Docker can redirect the cache
# to a path outside /data/ (which is typically volume-mounted and would hide the
# pre-baked encoding baked into the image at build time).
TIKTOKEN_CACHE_DIR = os.environ.get("TIKTOKEN_CACHE_DIR", "").strip() or f"{DATA_FOLDER}/tiktoken-cache"
os.makedirs(TIKTOKEN_CACHE_DIR, exist_ok=True)
+295
View File
@@ -0,0 +1,295 @@
"""
Async migration system for SurrealDB using the official Python client.
Based on patterns from sblpy migration system.
"""
from typing import List
from loguru import logger
from .repository import db_connection, repo_query
class AsyncMigration:
"""
Handles individual migration operations with async support.
"""
def __init__(self, sql: str) -> None:
"""Initialize migration with SQL content."""
self.sql = sql
@classmethod
def from_file(cls, file_path: str) -> "AsyncMigration":
"""Create migration from SQL file."""
with open(file_path, "r", encoding="utf-8") as file:
raw_content = file.read()
# Clean up SQL content
lines = []
for line in raw_content.split("\n"):
line = line.strip()
if line and not line.startswith("--"):
lines.append(line)
sql = " ".join(lines)
return cls(sql)
async def run(self, bump: bool = True) -> None:
"""Run the migration."""
try:
async with db_connection() as connection:
await connection.query(self.sql)
if bump:
await bump_version()
else:
await lower_version()
except Exception as e:
logger.error(f"Migration failed: {str(e)}")
raise
class AsyncMigrationRunner:
"""
Handles running multiple migrations in sequence.
"""
def __init__(
self,
up_migrations: List[AsyncMigration],
down_migrations: List[AsyncMigration],
) -> None:
"""Initialize runner with migration lists."""
self.up_migrations = up_migrations
self.down_migrations = down_migrations
async def run_all(self) -> None:
"""Run all pending up migrations."""
current_version = await get_latest_version()
for i in range(current_version, len(self.up_migrations)):
logger.info(f"Running migration {i + 1}")
await self.up_migrations[i].run(bump=True)
async def run_one_up(self) -> None:
"""Run one up migration."""
current_version = await get_latest_version()
if current_version < len(self.up_migrations):
logger.info(f"Running migration {current_version + 1}")
await self.up_migrations[current_version].run(bump=True)
async def run_one_down(self) -> None:
"""Run one down migration."""
current_version = await get_latest_version()
if current_version > 0:
logger.info(f"Rolling back migration {current_version}")
await self.down_migrations[current_version - 1].run(bump=False)
class AsyncMigrationManager:
"""
Main migration manager with async support.
"""
def __init__(self):
"""Initialize migration manager."""
self.up_migrations = [
AsyncMigration.from_file("open_notebook/database/migrations/1.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/2.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/3.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/4.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/5.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/6.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/7.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/8.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/9.surrealql"),
AsyncMigration.from_file("open_notebook/database/migrations/10.surrealql"),
AsyncMigration.from_file(
"open_notebook/database/migrations/11.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/12.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/13.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/14.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/15.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/16.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/17.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/18.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/19.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/20.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/21.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/22.surrealql"
),
]
self.down_migrations = [
AsyncMigration.from_file(
"open_notebook/database/migrations/1_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/2_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/3_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/4_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/5_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/6_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/7_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/8_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/9_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/10_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/11_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/12_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/13_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/14_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/15_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/16_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/17_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/18_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/19_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/20_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/21_down.surrealql"
),
AsyncMigration.from_file(
"open_notebook/database/migrations/22_down.surrealql"
),
]
self.runner = AsyncMigrationRunner(
up_migrations=self.up_migrations,
down_migrations=self.down_migrations,
)
async def get_current_version(self) -> int:
"""Get current database version."""
return await get_latest_version()
async def ping(self) -> None:
"""Check whether SurrealDB is reachable for migration startup."""
async with db_connection() as connection:
await connection.query("RETURN true;")
# Also exercise the migration version path. get_current_version() already
# treats a missing migrations table as version 0 for fresh databases.
await self.get_current_version()
async def needs_migration(self) -> bool:
"""Check if migration is needed."""
current_version = await self.get_current_version()
return current_version < len(self.up_migrations)
async def run_migration_up(self):
"""Run all pending migrations."""
current_version = await self.get_current_version()
logger.info(f"Current version before migration: {current_version}")
if await self.needs_migration():
try:
await self.runner.run_all()
new_version = await self.get_current_version()
logger.info(f"Migration successful. New version: {new_version}")
except Exception as e:
logger.error(f"Migration failed: {str(e)}")
raise
else:
logger.info("Database is already at the latest version")
# Database version management functions
async def get_latest_version() -> int:
"""Get the latest version from the migrations table."""
try:
versions = await get_all_versions()
if not versions:
return 0
return max(version["version"] for version in versions)
except Exception:
# If migrations table doesn't exist, we're at version 0
return 0
async def get_all_versions() -> List[dict]:
"""Get all versions from the migrations table."""
try:
result = await repo_query("SELECT * FROM _sbl_migrations ORDER BY version;")
return result
except Exception:
# If table doesn't exist, return empty list
return []
async def bump_version() -> None:
"""Bump the version by adding a new entry to migrations table."""
current_version = await get_latest_version()
new_version = current_version + 1
await repo_query(
"CREATE type::thing('_sbl_migrations', $version) SET version = $version, applied_at = time::now();",
{"version": new_version},
)
async def lower_version() -> None:
"""Lower the version by removing the latest entry from migrations table."""
current_version = await get_latest_version()
if current_version > 0:
await repo_query(
"DELETE type::thing('_sbl_migrations', $version);",
{"version": current_version},
)
+26
View File
@@ -0,0 +1,26 @@
import asyncio
from .async_migrate import AsyncMigrationManager
class MigrationManager:
"""
Synchronous wrapper around AsyncMigrationManager for backward compatibility.
"""
def __init__(self):
"""Initialize with async migration manager."""
self._async_manager = AsyncMigrationManager()
def get_current_version(self) -> int:
"""Get current database version (sync wrapper)."""
return asyncio.run(self._async_manager.get_current_version())
@property
def needs_migration(self) -> bool:
"""Check if migration is needed (sync wrapper)."""
return asyncio.run(self._async_manager.needs_migration())
def run_migration_up(self):
"""Run migrations (sync wrapper)."""
asyncio.run(self._async_manager.run_migration_up())
@@ -0,0 +1,178 @@
DEFINE TABLE IF NOT EXISTS source SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS
asset
ON TABLE source
FLEXIBLE TYPE option<object>;
DEFINE FIELD IF NOT EXISTS title ON TABLE source TYPE option<string>;
DEFINE FIELD IF NOT EXISTS topics ON TABLE source TYPE option<array<string>>;
DEFINE FIELD IF NOT EXISTS full_text ON TABLE source TYPE option<string>;
DEFINE FIELD IF NOT EXISTS created ON source DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON source DEFAULT time::now() VALUE time::now();
DEFINE TABLE IF NOT EXISTS source_embedding SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS source ON TABLE source_embedding TYPE record<source>;
DEFINE FIELD IF NOT EXISTS order ON TABLE source_embedding TYPE int;
DEFINE FIELD IF NOT EXISTS content ON TABLE source_embedding TYPE string;
DEFINE FIELD IF NOT EXISTS embedding ON TABLE source_embedding TYPE array<float>;
DEFINE TABLE IF NOT EXISTS source_insight SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS source ON TABLE source_insight TYPE record<source>;
DEFINE FIELD IF NOT EXISTS insight_type ON TABLE source_insight TYPE string;
DEFINE FIELD IF NOT EXISTS content ON TABLE source_insight TYPE string;
DEFINE FIELD IF NOT EXISTS embedding ON TABLE source_insight TYPE array<float>;
DEFINE EVENT IF NOT EXISTS source_delete ON TABLE source WHEN ($after == NONE) THEN {
delete source_embedding where source == $before.id;
delete source_insight where source == $before.id;
};
DEFINE TABLE IF NOT EXISTS note SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS title ON TABLE note TYPE option<string>;
DEFINE FIELD IF NOT EXISTS summary ON TABLE note TYPE option<string>;
DEFINE FIELD IF NOT EXISTS content ON TABLE note TYPE option<string>;
DEFINE FIELD IF NOT EXISTS embedding ON TABLE note TYPE array<float>;
DEFINE FIELD IF NOT EXISTS created ON note DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON note DEFAULT time::now() VALUE time::now();
DEFINE TABLE IF NOT EXISTS notebook SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS name ON TABLE notebook TYPE option<string>;
DEFINE FIELD IF NOT EXISTS description ON TABLE notebook TYPE option<string>;
DEFINE FIELD IF NOT EXISTS archived ON TABLE notebook TYPE option<bool> DEFAULT False;
DEFINE FIELD IF NOT EXISTS created ON notebook DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON notebook DEFAULT time::now() VALUE time::now();
DEFINE TABLE IF NOT EXISTS reference
TYPE RELATION
FROM source TO notebook;
DEFINE TABLE IF NOT EXISTS artifact
TYPE RELATION
FROM note TO notebook;
DEFINE TABLE IF NOT EXISTS podcast_config SCHEMALESS;
-- entender o analyzer
DEFINE ANALYZER IF NOT EXISTS my_analyzer TOKENIZERS blank,class,camel,punct FILTERS snowball(english), lowercase;
DEFINE INDEX IF NOT EXISTS idx_source_title ON TABLE source COLUMNS title SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX IF NOT EXISTS idx_source_full_text ON TABLE source COLUMNS full_text SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX IF NOT EXISTS idx_source_embed_chunk ON TABLE source_embedding COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX IF NOT EXISTS idx_source_insight ON TABLE source_insight COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX IF NOT EXISTS idx_note ON TABLE note COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX IF NOT EXISTS idx_note_title ON TABLE note COLUMNS title SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_title_search =
IF $sources {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM source
WHERE title @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_embedding_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source_embedding
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_full_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source
WHERE full_text @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source_insight
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $note_title_search =
IF $show_notes {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM note
WHERE title @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM note
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
let $source_asset_results = array::union($source_title_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_asset_results );
let $note_results = array::union($note_title_search, $note_content_search );
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT item_id, math::max(relevance) as relevance from $final_results
group by item_id ORDER BY relevance DESC LIMIT $match_count);
};
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_embedding_search =
IF $sources {(
SELECT source as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding LIMIT $match_count)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT source as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight LIMIT $match_count)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM note LIMIT $match_count)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_insight_search);
let $note_results = $note_content_search;
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT item_id, math::max(similarity) as similarity from $final_results
group by item_id ORDER BY similarity DESC LIMIT $match_count);
};
IF array::len(select * from open_notebook:default_models) == 0 THEN
CREATE open_notebook:default_models SET
default_chat_model= ""
END;
@@ -0,0 +1,13 @@
-- Migration 10: Add indexes for source_insight and source_embedding source field
-- These indexes significantly improve performance of source listing queries
-- that count insights and check embedding existence per source
DEFINE INDEX IF NOT EXISTS idx_source_insight_source ON source_insight FIELDS source CONCURRENTLY;
DEFINE INDEX IF NOT EXISTS idx_source_embedding_source ON source_embedding FIELDS source CONCURRENTLY;
DEFINE FIELD OVERWRITE embedding ON TABLE source_insight TYPE option<array<float>>;
DEFINE FIELD OVERWRITE embedding ON TABLE note TYPE option<array<float>>;
-- delete orphan records
DELETE from source_embedding WHERE source.id=NONE;
DELETE from source_insight WHERE source.id=NONE;
@@ -0,0 +1,4 @@
-- Rollback Migration 10: Remove source field indexes
REMOVE INDEX IF EXISTS idx_source_insight_source ON TABLE source_insight;
REMOVE INDEX IF EXISTS idx_source_embedding_source ON TABLE source_embedding;
@@ -0,0 +1,10 @@
-- Migration 11: Create provider configuration singleton record
-- This record stores multiple API key configurations per provider
-- The data is managed by the ProviderConfig RecordModel class
-- Create the provider configs singleton record for multi-config support
-- This record stores multiple API key configurations per provider
-- The data is managed by the ProviderConfig RecordModel class
UPSERT open_notebook:provider_configs CONTENT {
credentials: {}
};
@@ -0,0 +1,4 @@
-- Rollback Migration 11: Remove provider configuration records
-- Remove provider configs singleton (if exists)
DELETE open_notebook:provider_configs;
@@ -0,0 +1,29 @@
-- Migration 12: Create credential table and add credential link to model table
-- Individual credential records replace the ProviderConfig singleton
-- Each credential stores API key and provider-specific configuration
DEFINE TABLE credential SCHEMAFULL;
DEFINE FIELD name ON credential TYPE string;
DEFINE FIELD provider ON credential TYPE string;
DEFINE FIELD modalities ON credential TYPE array DEFAULT [];
DEFINE FIELD modalities.* ON credential TYPE string;
DEFINE FIELD api_key ON credential TYPE option<string>;
DEFINE FIELD base_url ON credential TYPE option<string>;
DEFINE FIELD endpoint ON credential TYPE option<string>;
DEFINE FIELD api_version ON credential TYPE option<string>;
DEFINE FIELD endpoint_llm ON credential TYPE option<string>;
DEFINE FIELD endpoint_embedding ON credential TYPE option<string>;
DEFINE FIELD endpoint_stt ON credential TYPE option<string>;
DEFINE FIELD endpoint_tts ON credential TYPE option<string>;
DEFINE FIELD project ON credential TYPE option<string>;
DEFINE FIELD location ON credential TYPE option<string>;
DEFINE FIELD credentials_path ON credential TYPE option<string>;
DEFINE FIELD created ON credential TYPE option<datetime> DEFAULT time::now();
DEFINE FIELD updated ON credential TYPE option<datetime> DEFAULT time::now();
-- Index for fast provider lookups
DEFINE INDEX idx_credential_provider ON credential FIELDS provider;
-- Add optional credential link to model table
DEFINE FIELD credential ON model TYPE option<record<credential>>;
@@ -0,0 +1,5 @@
-- Rollback Migration 12: Remove credential table and credential field from model
REMOVE FIELD credential ON TABLE model;
REMOVE INDEX idx_credential_provider ON credential;
REMOVE TABLE credential;
@@ -0,0 +1,3 @@
DEFINE FIELD OVERWRITE embedding ON TABLE source_insight TYPE option<array<float>>;
DEFINE FIELD OVERWRITE embedding ON TABLE note TYPE option<array<float>>;
@@ -0,0 +1,2 @@
DEFINE FIELD OVERWRITE embedding ON TABLE source_insight TYPE array<float>;
DEFINE FIELD OVERWRITE embedding ON TABLE note TYPE array<float>;
@@ -0,0 +1,27 @@
-- Migration 14: Podcast profiles model registry integration
-- Adds record<model> references to replace loose provider/model strings
-- Adds language field to episode_profile
-- Adds per-speaker TTS override support
-- EPISODE PROFILE
-- Legacy fields: make optional (app ignores, preserved for data migration)
DEFINE FIELD OVERWRITE outline_provider ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE outline_model ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE transcript_provider ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE transcript_model ON TABLE episode_profile TYPE option<string>;
-- New fields: reference to Model registry
DEFINE FIELD IF NOT EXISTS outline_llm ON TABLE episode_profile TYPE option<record<model>>;
DEFINE FIELD IF NOT EXISTS transcript_llm ON TABLE episode_profile TYPE option<record<model>>;
DEFINE FIELD IF NOT EXISTS language ON TABLE episode_profile TYPE option<string>;
-- SPEAKER PROFILE
-- Legacy fields: make optional
DEFINE FIELD OVERWRITE tts_provider ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD OVERWRITE tts_model ON TABLE speaker_profile TYPE option<string>;
-- New field: reference to Model registry (profile-level)
DEFINE FIELD IF NOT EXISTS voice_model ON TABLE speaker_profile TYPE option<record<model>>;
-- Per-speaker TTS override
DEFINE FIELD IF NOT EXISTS speakers.*.voice_model ON TABLE speaker_profile TYPE option<record<model>>;
@@ -0,0 +1,20 @@
-- Migration 14 rollback: Remove model registry fields from podcast profiles
-- Remove new fields from episode_profile
REMOVE FIELD IF EXISTS outline_llm ON TABLE episode_profile;
REMOVE FIELD IF EXISTS transcript_llm ON TABLE episode_profile;
REMOVE FIELD IF EXISTS language ON TABLE episode_profile;
-- Restore episode_profile legacy fields as required strings
DEFINE FIELD OVERWRITE outline_provider ON TABLE episode_profile TYPE string;
DEFINE FIELD OVERWRITE outline_model ON TABLE episode_profile TYPE string;
DEFINE FIELD OVERWRITE transcript_provider ON TABLE episode_profile TYPE string;
DEFINE FIELD OVERWRITE transcript_model ON TABLE episode_profile TYPE string;
-- Remove new fields from speaker_profile
REMOVE FIELD IF EXISTS voice_model ON TABLE speaker_profile;
REMOVE FIELD IF EXISTS speakers.*.voice_model ON TABLE speaker_profile;
-- Restore speaker_profile legacy fields as required strings
DEFINE FIELD OVERWRITE tts_provider ON TABLE speaker_profile TYPE string;
DEFINE FIELD OVERWRITE tts_model ON TABLE speaker_profile TYPE string;
@@ -0,0 +1,6 @@
-- Migration 15: Add a flexible config object to the credential table.
-- Stores provider-specific tuning options (e.g. Ollama num_ctx) without a
-- schema change per option. FLEXIBLE allows arbitrary keys inside the object
-- on this SCHEMAFULL table; validation stays in the Pydantic Credential model.
DEFINE FIELD config ON credential FLEXIBLE TYPE option<object>;
@@ -0,0 +1,3 @@
-- Migration 15 rollback: remove the flexible config field from credential
REMOVE FIELD IF EXISTS config ON TABLE credential;
@@ -0,0 +1,7 @@
-- Migration 16: Add optional max_tokens override to episode_profile.
-- Lets a profile cap (or raise) the output tokens used for outline and
-- transcript generation, which is passed through to podcast_creator.
-- The episode_profile table is SCHEMAFULL, so the field must be declared
-- here for the value to persist.
DEFINE FIELD IF NOT EXISTS max_tokens ON TABLE episode_profile TYPE option<int>;
@@ -0,0 +1,3 @@
-- Migration 16 rollback: Remove max_tokens field from episode_profile.
REMOVE FIELD IF EXISTS max_tokens ON TABLE episode_profile;
@@ -0,0 +1,3 @@
-- Migration 17: Per-transformation model selection
DEFINE FIELD IF NOT EXISTS model_id ON TABLE transformation TYPE option<record<model>>;
@@ -0,0 +1,3 @@
-- Migration 17 rollback: Remove per-transformation model selection
REMOVE FIELD IF EXISTS model_id ON TABLE transformation;
@@ -0,0 +1,10 @@
-- Migration 18: Recently viewed timestamps
-- Adds optional last-viewed timestamps for notebooks and sources.
DEFINE FIELD IF NOT EXISTS last_viewed_at ON TABLE notebook TYPE option<datetime>;
DEFINE FIELD IF NOT EXISTS last_viewed_at ON TABLE source TYPE option<datetime>;
-- Indexes backing the recently-viewed query (ORDER BY last_viewed_at DESC LIMIT)
-- so it does not degrade into a full table scan as data grows.
DEFINE INDEX IF NOT EXISTS idx_notebook_last_viewed ON TABLE notebook FIELDS last_viewed_at CONCURRENTLY;
DEFINE INDEX IF NOT EXISTS idx_source_last_viewed ON TABLE source FIELDS last_viewed_at CONCURRENTLY;
@@ -0,0 +1,6 @@
-- Migration 18 rollback: Remove recently viewed timestamps
REMOVE INDEX IF EXISTS idx_notebook_last_viewed ON TABLE notebook;
REMOVE INDEX IF EXISTS idx_source_last_viewed ON TABLE source;
REMOVE FIELD IF EXISTS last_viewed_at ON TABLE notebook;
REMOVE FIELD IF EXISTS last_viewed_at ON TABLE source;
@@ -0,0 +1,8 @@
-- Migration 19: Timestamp source_insight records
-- source_insight is SCHEMAFULL but never defined created/updated, so new
-- insights were stored without timestamps (and clients saw "None").
-- Mirrors the created/updated definitions used by source, note and notebook.
-- Existing rows are left untouched (no backfill); the API returns null for them.
DEFINE FIELD IF NOT EXISTS created ON source_insight DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON source_insight DEFAULT time::now() VALUE time::now();
@@ -0,0 +1,4 @@
-- Migration 19 rollback: Remove source_insight timestamp fields
REMOVE FIELD IF EXISTS created ON TABLE source_insight;
REMOVE FIELD IF EXISTS updated ON TABLE source_insight;
@@ -0,0 +1,24 @@
REMOVE TABLE IF EXISTS source;
REMOVE TABLE IF EXISTS source_embedding;
REMOVE TABLE IF EXISTS source_insight;
REMOVE TABLE IF EXISTS note;
REMOVE TABLE IF EXISTS notebook;
REMOVE TABLE IF EXISTS reference;
REMOVE TABLE IF EXISTS artifact;
REMOVE TABLE IF EXISTS podcast_config;
REMOVE EVENT IF EXISTS source_delete ON TABLE source;
REMOVE ANALYZER IF EXISTS my_analyzer;
REMOVE INDEX IF EXISTS idx_source_title ON TABLE source;
REMOVE INDEX IF EXISTS idx_source_full_text ON TABLE source;
REMOVE INDEX IF EXISTS idx_source_embed_chunk ON TABLE source_embedding;
REMOVE INDEX IF EXISTS idx_source_insight ON TABLE source_insight;
REMOVE INDEX IF EXISTS idx_note ON TABLE note;
REMOVE INDEX IF EXISTS idx_note_title ON TABLE note;
REMOVE FUNCTION IF EXISTS fn::text_search;
REMOVE FUNCTION IF EXISTS fn::vector_search;
DELETE open_notebook:default_models;
@@ -0,0 +1 @@
DEFINE FIELD IF NOT EXISTS note_type ON TABLE note TYPE option<string>;
@@ -0,0 +1,24 @@
-- Migration 20: episode_profile.speaker_config references speaker_profile by record ID (#630)
-- Previously speaker_config stored the speaker profile NAME (a plain string),
-- so renaming a speaker profile silently broke every episode profile that
-- referenced it. This converts the field to a record link.
--
-- Steps:
-- 1. Widen the field type so both the legacy strings and the new record
-- links are valid during the data conversion (the table is SCHEMAFULL).
-- 2. Convert existing rows: look up the speaker_profile with a matching
-- name and store its record ID. Orphan references (a name that no longer
-- matches any speaker_profile) become NONE - the app treats a missing
-- speaker as "needs setup" and the UI asks the user to pick one.
-- 3. Tighten the type to option<record<speaker_profile>> so only
-- speaker_profile record links (or NONE) can be stored from now on.
-- Note: SurrealDB record links do NOT prevent dangling references -
-- deleting a speaker profile still leaves the link in place, so the
-- app must keep resolving speaker_config defensively (it treats an
-- unresolvable reference the same as NONE: "needs setup").
DEFINE FIELD OVERWRITE speaker_config ON TABLE episode_profile TYPE option<string | record<speaker_profile>>;
UPDATE episode_profile SET speaker_config = (SELECT VALUE id FROM ONLY speaker_profile WHERE name = $parent.speaker_config LIMIT 1) WHERE type::is::string(speaker_config);
DEFINE FIELD OVERWRITE speaker_config ON TABLE episode_profile TYPE option<record<speaker_profile>>;
@@ -0,0 +1,12 @@
-- Migration 20 rollback: restore episode_profile.speaker_config to the speaker profile name
-- Widen the type first (SCHEMAFULL table), convert record links back to the
-- linked profile's name, then settle on option<string>. The original type was
-- a required string, but rows whose reference was orphaned by the up
-- migration (or whose linked speaker_profile was deleted since) have no name
-- to restore, so the rollback type must tolerate NONE.
DEFINE FIELD OVERWRITE speaker_config ON TABLE episode_profile TYPE option<string | record<speaker_profile>>;
UPDATE episode_profile SET speaker_config = speaker_config.name WHERE type::is::record(speaker_config);
DEFINE FIELD OVERWRITE speaker_config ON TABLE episode_profile TYPE option<string>;
@@ -0,0 +1,29 @@
-- Migration 21: store episode.audio_file relative to PODCASTS_FOLDER (#1030)
-- audio_file used to hold the absolute path returned by podcast-creator
-- (str(Path.resolve())), sometimes with a legacy file:// prefix. From this
-- migration on the app stores the path RELATIVE to PODCASTS_FOLDER
-- (e.g. "episodes/<uuid>/audio/<uuid>.mp3") and resolves + contains it at a
-- single choke point (open_notebook/podcasts/audio_paths.py).
--
-- Conversion is by known-prefix stripping only, in two steps:
-- 1. Normalize plain "file:///" URIs to absolute paths. Only URIs without
-- percent-encoding are handled (SurrealQL cannot URL-decode) and only
-- the empty-host "file:///" form (a "file://host/" URI is not a local
-- path); anything else keeps its scheme and stays legacy-invalid.
-- 2. Strip EXACTLY ONE known root prefix per row (a single IF/ELSE chain,
-- so a remainder that happens to start with another prefix is never
-- stripped twice):
-- "/app/data/podcasts/" the shipped Docker layout (WORKDIR /app,
-- DATA_FOLDER "./data" resolved at write time)
-- "/data/podcasts/" older Docker layout with data mounted at /data
-- "./data/podcasts/" CWD-relative forms written by source
-- "data/podcasts/" checkouts running from the repo root
--
-- Rows under any OTHER root (e.g. a source checkout at a custom absolute
-- path) stay untouched by design (settled in #1030): the resolve helper
-- treats any non-relative value as legacy-invalid, preserving the 403/404
-- behavior #1018's guards give those rows today. No nulling, no flags.
UPDATE episode SET audio_file = string::slice(audio_file, 7) WHERE type::is::string(audio_file) AND string::starts_with(audio_file, "file:///") AND !string::contains(audio_file, "%");
UPDATE episode SET audio_file = IF string::starts_with(audio_file, "/app/data/podcasts/") { string::slice(audio_file, 19) } ELSE IF string::starts_with(audio_file, "/data/podcasts/") { string::slice(audio_file, 15) } ELSE IF string::starts_with(audio_file, "./data/podcasts/") { string::slice(audio_file, 16) } ELSE IF string::starts_with(audio_file, "data/podcasts/") { string::slice(audio_file, 14) } ELSE { audio_file } WHERE type::is::string(audio_file);
@@ -0,0 +1,10 @@
-- Migration 21 rollback: restore absolute audio_file paths (best effort).
-- The up migration collapsed several historical prefixes (file://,
-- /app/data/podcasts/, /data/podcasts/, ./data/podcasts/, data/podcasts/)
-- into one relative form, so the original spelling of each row is not
-- recoverable. Rows are restored to the shipped Docker layout, which is
-- what pre-1030 code resolved them to in the standard deployment. Rows the
-- up migration left untouched (unknown roots) are already absolute and are
-- skipped.
UPDATE episode SET audio_file = string::concat("/app/data/podcasts/", audio_file) WHERE type::is::string(audio_file) AND audio_file != "" AND !string::starts_with(audio_file, "/") AND !string::starts_with(audio_file, "file://");
@@ -0,0 +1,42 @@
-- Migration 22: drop legacy provider/model string fields from podcast profiles (#1107)
-- Phase 4 of #584. The app has ignored these strings since migration 14
-- introduced the Model registry references (outline_llm, transcript_llm,
-- voice_model); their only remaining role was as input to the startup data
-- migration retry loop (open_notebook/podcasts/migration.py), which this
-- migration replaces and retires.
--
-- Steps:
-- 1. Best-effort mapping: for profiles whose reference fields are still
-- empty but whose legacy provider+model strings are set, populate the
-- reference from a matching `model` record (provider + name + type).
-- This mirrors the retired Python data migration MINUS the auto-create
-- path: a migration has no business touching credentials, so profiles
-- with no matching model record simply stay unresolved. Accepted
-- trade-off (#1107): those profiles were already non-functional (the app
-- ignores legacy strings) and the UI already flags them as needing model
-- selection - the user re-picks models in the profile form once.
-- 2. Clear the stored legacy values (UNSET while the fields are still
-- defined, so rows don't keep stale data the schema no longer knows).
-- 3. Drop the 6 legacy field definitions.
-- Step 1: best-effort mapping (no-op on rows already migrated or with no
-- legacy strings; a subquery with no match yields NONE, leaving the row
-- unresolved, which is the accepted outcome).
UPDATE episode_profile SET outline_llm = (SELECT VALUE id FROM ONLY model WHERE provider = $parent.outline_provider AND name = $parent.outline_model AND type = "language" LIMIT 1) WHERE outline_llm = NONE AND outline_provider != NONE AND outline_provider != "" AND outline_model != NONE AND outline_model != "";
UPDATE episode_profile SET transcript_llm = (SELECT VALUE id FROM ONLY model WHERE provider = $parent.transcript_provider AND name = $parent.transcript_model AND type = "language" LIMIT 1) WHERE transcript_llm = NONE AND transcript_provider != NONE AND transcript_provider != "" AND transcript_model != NONE AND transcript_model != "";
UPDATE speaker_profile SET voice_model = (SELECT VALUE id FROM ONLY model WHERE provider = $parent.tts_provider AND name = $parent.tts_model AND type = "text_to_speech" LIMIT 1) WHERE voice_model = NONE AND tts_provider != NONE AND tts_provider != "" AND tts_model != NONE AND tts_model != "";
-- Step 2: clear stored legacy values before dropping the definitions.
UPDATE episode_profile UNSET outline_provider, outline_model, transcript_provider, transcript_model;
UPDATE speaker_profile UNSET tts_provider, tts_model;
-- Step 3: drop the legacy field definitions.
REMOVE FIELD IF EXISTS outline_provider ON TABLE episode_profile;
REMOVE FIELD IF EXISTS outline_model ON TABLE episode_profile;
REMOVE FIELD IF EXISTS transcript_provider ON TABLE episode_profile;
REMOVE FIELD IF EXISTS transcript_model ON TABLE episode_profile;
REMOVE FIELD IF EXISTS tts_provider ON TABLE speaker_profile;
REMOVE FIELD IF EXISTS tts_model ON TABLE speaker_profile;
@@ -0,0 +1,14 @@
-- Migration 22 rollback: re-define the legacy provider/model string fields
-- (best effort, documented in #1107). The up migration UNSET the stored
-- values before dropping the definitions, so the legacy strings themselves
-- are NOT recoverable - this rollback only restores the schema shape that
-- migration 14 left (optional strings, app-ignored). Profiles keep working
-- through the Model registry references (outline_llm, transcript_llm,
-- voice_model), which the up migration only ever populated, never cleared.
DEFINE FIELD OVERWRITE outline_provider ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE outline_model ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE transcript_provider ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE transcript_model ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD OVERWRITE tts_provider ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD OVERWRITE tts_model ON TABLE speaker_profile TYPE option<string>;
@@ -0,0 +1 @@
REMOVE FIELD IF EXISTS note_type ON TABLE note;
@@ -0,0 +1,145 @@
DEFINE TABLE IF NOT EXISTS chat_session SCHEMALESS;
DEFINE TABLE IF NOT EXISTS refers_to
TYPE RELATION
FROM chat_session TO notebook;
REMOVE FUNCTION IF EXISTS fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
let $source_embedding_search =
IF $sources {(
SELECT
id,
source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT
id,
insight_type + ' - ' + source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT
id,
title,
content,
id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM note
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $all_results = array::union(
array::union($source_embedding_search, $source_insight_search),
$note_content_search
);
RETURN (
SELECT
id, title, content, parent_id,
math::max(similarity) as similarity
FROM $all_results
GROUP BY id
ORDER BY similarity DESC
LIMIT $match_count
);
};
REMOVE FUNCTION IF EXISTS fn::text_search;
DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_title_search =
IF $sources {(
SELECT id, title,
search::highlight('`', '`', 1) as content,
id as parent_id,
math::max(search::score(1)) AS relevance
FROM source
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_embedding_search =
IF $sources {(
SELECT id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source_embedding
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_full_search =
IF $sources {(
SELECT source.id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source
WHERE full_text @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT id, insight_type + " - " + source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source_insight
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_title_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
let $source_asset_results = array::union($source_title_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_asset_results );
let $note_results = array::union($note_title_search, $note_content_search );
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT id, title, content, parent_id, math::max(relevance) as relevance from $final_results
where id is not None
group by id, title, content, parent_id ORDER BY relevance DESC LIMIT $match_count);
};
@@ -0,0 +1,110 @@
REMOVE TABLE IF EXISTS chat_session;
REMOVE TABLE IF EXISTS refers_to;
REMOVE FUNCTION fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_embedding_search =
IF $sources {(
SELECT source as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding LIMIT $match_count)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT source as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight LIMIT $match_count)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id as item_id, content, vector::similarity::cosine(embedding, $query) as similarity
FROM note LIMIT $match_count)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_insight_search);
let $note_results = $note_content_search;
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT item_id, math::max(similarity) as similarity from $final_results
group by item_id ORDER BY similarity DESC LIMIT $match_count);
};
REMOVE FUNCTION fn::text_search;
DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_title_search =
IF $sources {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM source
WHERE title @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_embedding_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source_embedding
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_full_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source
WHERE full_text @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT source as item_id, math::max(search::score(1)) AS relevance
FROM source_insight
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $note_title_search =
IF $show_notes {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM note
WHERE title @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id as item_id, math::max(search::score(1)) AS relevance
FROM note
WHERE content @1@ $query_text
GROUP BY item_id)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
let $source_asset_results = array::union($source_title_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_asset_results );
let $note_results = array::union($note_title_search, $note_content_search );
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT item_id, math::max(relevance) as relevance from $final_results
group by item_id ORDER BY relevance DESC LIMIT $match_count);
};
@@ -0,0 +1,134 @@
REMOVE FUNCTION IF EXISTS fn::text_search;
DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_title_search =
IF $sources {(
SELECT id, title,
search::highlight('`', '`', 1) as content,
id as parent_id,
math::max(search::score(1)) AS relevance
FROM source
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_embedding_search =
IF $sources {(
SELECT source.id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source_embedding
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_full_search =
IF $sources {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM source
WHERE full_text @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT id, insight_type + " - " + (source.title OR '') as title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM source_insight
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_title_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
let $source_asset_results = array::union($source_title_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_asset_results );
let $note_results = array::union($note_title_search, $note_content_search );
let $final_results = array::union($source_results, $note_results );
RETURN (select id, parent_id, title, math::max(relevance) as relevance
from $final_results where id is not None
group by id, parent_id, title ORDER BY relevance DESC LIMIT $match_count);
};
REMOVE FUNCTION IF EXISTS fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
let $source_embedding_search =
IF $sources {(
SELECT
source.id as id,
source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT
id,
insight_type + ' - ' + (source.title OR '') as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT
id,
title,
content,
id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM note
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $all_results = array::union(
array::union($source_embedding_search, $source_insight_search),
$note_content_search
);
RETURN (select id, parent_id, title, math::max(similarity) as similarity,
array::flatten(content) as matches
from $all_results where id is not None
group by id, parent_id, title ORDER BY similarity DESC LIMIT $match_count);
};
@@ -0,0 +1,139 @@
REMOVE FUNCTION IF EXISTS fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
let $source_embedding_search =
IF $sources {(
SELECT
id,
source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT
id,
insight_type + ' - ' + source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT
id,
title,
content,
id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM note
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $all_results = array::union(
array::union($source_embedding_search, $source_insight_search),
$note_content_search
);
RETURN (
SELECT
id, title, content, parent_id,
math::max(similarity) as similarity
FROM $all_results
GROUP BY id
ORDER BY similarity DESC
LIMIT $match_count
);
};
REMOVE FUNCTION IF EXISTS fn::text_search;
DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count: int, $sources:bool, $show_notes:bool) {
let $source_title_search =
IF $sources {(
SELECT id, title,
search::highlight('`', '`', 1) as content,
id as parent_id,
math::max(search::score(1)) AS relevance
FROM source
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_embedding_search =
IF $sources {(
SELECT id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source_embedding
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_full_search =
IF $sources {(
SELECT source.id as id, source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source
WHERE full_text @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT id, insight_type + " - " + source.title as title, search::highlight('`', '`', 1) as content, source.id as parent_id, math::max(search::score(1)) AS relevance
FROM source_insight
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_title_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE title @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT id, title, search::highlight('`', '`', 1) as content, id as parent_id, math::max(search::score(1)) AS relevance
FROM note
WHERE content @1@ $query_text
GROUP BY id)}
ELSE { [] };
let $source_chunk_results = array::union($source_embedding_search, $source_full_search);
let $source_asset_results = array::union($source_title_search, $source_insight_search);
let $source_results = array::union($source_chunk_results, $source_asset_results );
let $note_results = array::union($note_title_search, $note_content_search );
let $final_results = array::union($source_results, $note_results );
RETURN (SELECT id, title, content, parent_id, math::max(relevance) as relevance from $final_results
where id is not None
group by id, title, content, parent_id ORDER BY relevance DESC LIMIT $match_count);
};
@@ -0,0 +1,169 @@
-- remove old transformation defaults
DELETE open_notebook:default_transformations;
-- set up the default transformations
DEFINE TABLE IF NOT EXISTS transformation SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS name ON TABLE transformation TYPE string;
DEFINE FIELD IF NOT EXISTS title ON TABLE transformation TYPE string;
DEFINE FIELD IF NOT EXISTS description ON TABLE transformation TYPE string;
DEFINE FIELD IF NOT EXISTS prompt ON TABLE transformation TYPE string;
DEFINE FIELD IF NOT EXISTS apply_default ON TABLE transformation TYPE bool DEFAULT False;
DEFINE FIELD IF NOT EXISTS created ON transformation DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON transformation DEFAULT time::now() VALUE time::now();
insert into transformation [
{
name: "Analyze Paper",
title: "Paper Analysis",
description: "Analyses a technical/scientific paper",
prompt:"# IDENTITY and PURPOSE
You are an insightful and analytical reader of academic papers, extracting the key components, significance, and broader implications. Your focus is to uncover the core contributions, practical applications, methodological strengths or weaknesses, and any surprising findings. You are especially attuned to the clarity of arguments, the relevance to existing literature, and potential impacts on both the specific field and broader contexts.
# STEPS
1. **READ AND UNDERSTAND THE PAPER**: Thoroughly read the paper, identifying its main focus, arguments, methods, results, and conclusions.
2. **IDENTIFY CORE ELEMENTS**:
- **Purpose**: What is the main goal or research question?
- **Contribution**: What new knowledge or innovation does this paper bring to the field?
- **Methods**: What methods are used, and are they novel or particularly effective?
- **Key Findings**: What are the most critical results, and why do they matter?
- **Limitations**: Are there any notable limitations or areas for further research?
3. **SYNTHESIZE THE MAIN POINTS**:
- Extract the key elements and organize them into insightful observations.
- Highlight the broader impact and potential applications.
- Note any aspects that challenge established views or introduce new questions.
# OUTPUT INSTRUCTIONS
- Structure the output as follows:
- **PURPOSE**: A concise summary of the main research question or goal (1-2 sentences).
- **CONTRIBUTION**: A bullet list of 2-3 points that describe what the paper adds to the field.
- **KEY FINDINGS**: A bullet list of 2-3 points summarizing the critical outcomes of the study.
- **IMPLICATIONS**: A bullet list of 2-3 points discussing the significance or potential impact of the findings on the field or broader context.
- **LIMITATIONS**: A bullet list of 1-2 points identifying notable limitations or areas for future work.
- **Bullet Points** should be between 15-20 words.
- Avoid starting each bullet point with the same word to maintain variety.
- Use clear and concise language that conveys the key ideas effectively.
- Do not include warnings, disclaimers, or personal opinions.
- Output only the requested sections with their respective labels.",
apply_default: False
},
{
name: "Key Insights",
title: "Key Insights",
description: "Extracts important insights and actionable items",
prompt:"# IDENTITY and PURPOSE
You extract surprising, powerful, and interesting insights from text content. You are interested in insights related to the purpose and meaning of life, human flourishing, the role of technology in the future of humanity, artificial intelligence and its affect on humans, memes, learning, reading, books, continuous improvement, and similar topics.
You create 15 word bullet points that capture the most important insights from the input.
Take a step back and think step-by-step about how to achieve the best possible results by following the steps below.
# STEPS
- Extract 20 to 50 of the most surprising, insightful, and/or interesting ideas from the input in a section called IDEAS, and write them on a virtual whiteboard in your mind using 15 word bullets. If there are less than 50 then collect all of them. Make sure you extract at least 20.
- From those IDEAS, extract the most powerful and insightful of them and write them in a section called INSIGHTS. Make sure you extract at least 10 and up to 25.
# OUTPUT INSTRUCTIONS
- INSIGHTS are essentially higher-level IDEAS that are more abstracted and wise.
- Output the INSIGHTS section only.
- Each bullet should be about 15 words in length.
- Do not give warnings or notes; only output the requested sections.
- You use bulleted lists for output, not numbered lists.
- Do not start items with the same opening words.
- Ensure you follow ALL these instructions when creating your output.
",
apply_default: False
},
{
name: "Dense Summary",
title: "Dense Summary",
description: "Creates a rich, deep summary of the content",
prompt:"# MISSION
You are a Sparse Priming Representation (SPR) writer. An SPR is a particular kind of use of language for advanced NLP, NLU, and NLG tasks, particularly useful for the latest generation of Large Language Models (LLMs). You will be given information by the USER which you are to render as an SPR.
# THEORY
LLMs are a kind of deep neural network. They have been demonstrated to embed knowledge, abilities, and concepts, ranging from reasoning to planning, and even to theory of mind. These are called latent abilities and latent content, collectively referred to as latent space. The latent space of an LLM can be activated with the correct series of words as inputs, which will create a useful internal state of the neural network. This is not unlike how the right shorthand cues can prime a human mind to think in a certain way. Like human minds, LLMs are associative, meaning you only need to use the correct associations to 'prime' another model to think in the same way.
# METHODOLOGY
Render the input as a distilled list of succinct statements, assertions, associations, concepts, analogies, and metaphors. The idea is to capture as much, conceptually, as possible but with as few words as possible. Write it in a way that makes sense to you, as the future audience will be another language model, not a human. Use complete sentences.
",
apply_default: True
},
{
name: "Reflections",
title: "Reflection Questions",
description: "Generates reflection questions from the document to help explore it further",
prompt:"# IDENTITY and PURPOSE
You extract deep, thought-provoking, and meaningful reflections from text content. You are especially focused on themes related to the human experience, such as the purpose of life, personal growth, the intersection of technology and humanity, artificial intelligence's societal impact, human potential, collective evolution, and transformative learning. Your reflections aim to provoke new ways of thinking, challenge assumptions, and provide a thoughtful synthesis of the content.
# STEPS
- Extract 3 to 5 of the most profound, thought-provoking, and/or meaningful ideas from the input in a section called REFLECTIONS.
- Each reflection should aim to explore underlying implications, connections to broader human experiences, or highlight a transformative perspective.
- Take a step back and consider the deeper significance or questions that arise from the content.
# OUTPUT INSTRUCTIONS
- The output section should be labeled as REFLECTIONS.
- Each bullet point should be between 20-25 words.
- Avoid repetition in the phrasing and ensure variety in sentence structure.
- The reflections should encourage deeper inquiry and provide a synthesis that transcends surface-level observations.
- Use bullet points, not numbered lists.
- Every bullet should be formatted as a question that elicits contemplation or a statement that offers a profound insight.
- Do not give warnings or notes; only output the requested section.",
apply_default: False
},
{
name: "Table of Contents",
title: "Table of Contents",
description: "Describes the different topics of the document",
prompt:"# SYSTEM ROLE
You are a content analysis assistant that reads through documents and provides a Table of Contents (ToC) to help users identify what the document covers more easily.
Your ToC should capture all major topics and transitions in the content and should mention them in the order theh appear.
# TASK
Analyze the provided content and create a Table of Contents:
- Captures the core topics included in the text
- Gives a small description of what is covered",
apply_default: False
},
{
name: "Simple Summary",
title: "Simple Summary",
description: "Generates a small summary of the content",
prompt:"# SYSTEM ROLE
You are a content summarization assistant that creates dense, information-rich summaries optimized for machine understanding. Your summaries should capture key concepts with minimal words while maintaining complete, clear sentences.
# TASK
Analyze the provided content and create a summary that:
- Captures the core concepts and key information
- Uses clear, direct language
- Maintains context from any previous summaries",
apply_default: False
},
];
-- Sets the default transformation instructions prompt
UPSERT open_notebook:default_prompts
CONTENT {transformation_instructions: "# INSTRUCTIONS
You are my learning assistant and you help me process and transform content so that I can extract insights from them.
# IMPORTANT
- You are working on my editorial projects. The text below is my own. Do not give me any warnings about copyright or plagiarism.
- Output ONLY the requested content, without acknowledgements of the task and additional chatting. Don't start with \"Sure, I can help you with that.\" or \"Here is the information you requested:\". Just provide the content.
- Do not stop in the middle of the generation to ask me questions. Execute my request completely.
"};
@@ -0,0 +1,2 @@
REMOVE TABLE IF EXISTS transformation SCHEMAFULL;
@@ -0,0 +1 @@
update model set provider='vertex' where provider='vertexai';
@@ -0,0 +1 @@
update model set provider='vertexai' where provider='vertex';
@@ -0,0 +1,152 @@
DEFINE TABLE IF NOT EXISTS episode_profile SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS name ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS description ON TABLE episode_profile TYPE option<string>;
DEFINE FIELD IF NOT EXISTS speaker_config ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS outline_provider ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS outline_model ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS transcript_provider ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS transcript_model ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS default_briefing ON TABLE episode_profile TYPE string;
DEFINE FIELD IF NOT EXISTS num_segments ON TABLE episode_profile TYPE int DEFAULT 5;
DEFINE FIELD IF NOT EXISTS created ON TABLE episode_profile TYPE datetime DEFAULT time::now();
DEFINE FIELD IF NOT EXISTS updated ON TABLE episode_profile TYPE datetime DEFAULT time::now();
-- Create Speaker Profile table
remove table speaker_profile;
DEFINE TABLE IF NOT EXISTS speaker_profile SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS name ON TABLE speaker_profile TYPE string;
DEFINE FIELD IF NOT EXISTS description ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD IF NOT EXISTS tts_provider ON TABLE speaker_profile TYPE string;
DEFINE FIELD IF NOT EXISTS tts_model ON TABLE speaker_profile TYPE string;
DEFINE FIELD IF NOT EXISTS speakers ON TABLE speaker_profile TYPE array<object>;
DEFINE FIELD IF NOT EXISTS speakers.*.name ON TABLE speaker_profile TYPE string;
DEFINE FIELD IF NOT EXISTS speakers.*.voice_id ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD IF NOT EXISTS speakers.*.backstory ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD IF NOT EXISTS speakers.*.personality ON TABLE speaker_profile TYPE option<string>;
DEFINE FIELD IF NOT EXISTS created ON TABLE speaker_profile TYPE datetime DEFAULT time::now();
DEFINE FIELD IF NOT EXISTS updated ON TABLE speaker_profile TYPE datetime DEFAULT time::now();
-- Enhance PodcastEpisode table
DEFINE TABLE IF NOT EXISTS episode SCHEMAFULL;
DEFINE FIELD IF NOT EXISTS created ON episode DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD IF NOT EXISTS updated ON episode DEFAULT time::now() VALUE time::now();
DEFINE FIELD IF NOT EXISTS name ON TABLE episode TYPE string;
DEFINE FIELD IF NOT EXISTS briefing ON TABLE episode TYPE option<string>;
DEFINE FIELD IF NOT EXISTS episode_profile ON TABLE episode FLEXIBLE TYPE object;
DEFINE FIELD IF NOT EXISTS speaker_profile ON TABLE episode FLEXIBLE TYPE object;
DEFINE FIELD IF NOT EXISTS transcript ON TABLE episode FLEXIBLE TYPE option<object>;
DEFINE FIELD IF NOT EXISTS outline ON TABLE episode FLEXIBLE TYPE option<object>;
DEFINE FIELD IF NOT EXISTS command ON TABLE episode TYPE option<record<command>>;
DEFINE FIELD IF NOT EXISTS content ON TABLE episode TYPE option<string>;
DEFINE FIELD IF NOT EXISTS audio_file ON TABLE episode TYPE option<string>;
-- Create indexes for better performance
DEFINE INDEX IF NOT EXISTS idx_episode_profile_name ON TABLE episode_profile COLUMNS name UNIQUE CONCURRENTLY;
DEFINE INDEX IF NOT EXISTS idx_speaker_profile_name ON TABLE speaker_profile COLUMNS name UNIQUE CONCURRENTLY;
DEFINE INDEX IF NOT EXISTS idx_episode_profile ON TABLE episode COLUMNS episode_profile CONCURRENTLY;
DEFINE INDEX IF NOT EXISTS idx_episode_command ON TABLE episode COLUMNS command CONCURRENTLY;
--Sample data
insert into episode_profile
[
{
name: "tech_discussion",
description: "Technical discussion between 2 experts",
speaker_config: "tech_experts",
outline_provider: "openai",
outline_model: "gpt-5-mini",
transcript_provider: "openai",
transcript_model: "gpt-5-mini",
default_briefing: "Create an engaging technical discussion about the provided content. Focus on practical insights, real-world applications, and detailed explanations that would interest developers and technical professionals.",
num_segments: 5
},
{
name: "solo_expert",
description: "Single expert explaining complex topics",
speaker_config: "solo_expert",
outline_provider: "openai",
outline_model: "gpt-5-mini",
transcript_provider: "openai",
transcript_model: "gpt-5-mini",
default_briefing: "Create an educational explanation of the provided content. Break down complex concepts into digestible segments, use analogies and examples, and maintain an engaging teaching style.",
"num_segments":4 },
{
name: "business_analysis",
description: "Business-focused analysis and discussion",
speaker_config: "business_panel",
outline_provider: "openai",
outline_model: "gpt-5-mini",
transcript_provider: "openai",
transcript_model: "gpt-5-mini",
default_briefing: "Analyze the provided content from a business perspective. Discuss market implications, strategic insights, competitive advantages, and actionable business intelligence.",
"num_segments":6 }
];
insert into speaker_profile
[
{
name: "tech_experts",
description: "Two technical experts for tech discussions",
tts_provider: "openai",
tts_model: "gpt-4o-mini-tts",
speakers: [
{
name: "Dr. Alex Chen",
voice_id: "nova",
backstory: "Senior AI researcher and former tech lead at major companies. Specializes in making complex technical concepts accessible.",
personality: "Analytical, clear communicator, asks probing questions to dig deeper into technical details"
},
{
name: "Jamie Rodriguez",
voice_id: "alloy",
backstory: "Full-stack engineer and tech entrepreneur. Loves practical applications and real-world implementations.",
personality: "Enthusiastic, practical-minded, great at explaining implementation details and trade-offs"
}
]
},
{
name: "solo_expert",
description: "Single expert for educational content",
tts_provider: "openai",
tts_model: "gpt-4o-mini-tts",
speakers: [
{
name: "Professor Sarah Kim",
voice_id: "nova",
backstory: "Distinguished professor and researcher. Has a gift for making complex topics accessible to broad audiences.",
personality: "Patient teacher, uses analogies and examples, breaks down complex concepts step by step"
}
]
},
{
name: "business_panel",
description: "Business analysis panel with diverse perspectives",
tts_provider: "openai",
tts_model: "gpt-4o-mini-tts",
speakers: [
{
name: "Marcus Thompson",
voice_id: "echo",
backstory: "Former McKinsey consultant, now startup advisor. Expert in strategic analysis and market dynamics.",
personality: "Strategic thinker, data-driven, excellent at identifying key insights and implications"
},
{
name: "Elena Vasquez",
voice_id: "shimmer",
backstory: "Serial entrepreneur and investor. Focuses on practical implementation and execution.",
personality: "Action-oriented, pragmatic, brings startup experience and execution focus"
},
{
name: "Johny Bing",
voice_id: "ash",
backstory: "Youtube celebrity and business mogul. Focuses on practical implementation and execution.",
personality: "Controversial, likes to question ideas and concepts. He brings a fresh perspective and always has a point to make."
}
]
}
];
@@ -0,0 +1,3 @@
REMOVE TABLE IF EXISTS episode_profile;
REMOVE TABLE IF EXISTS speaker_profile;
REMOVE TABLE IF EXISTS episode;
@@ -0,0 +1,12 @@
-- Migration 8: Support chat sessions for both notebooks and sources
-- This migration allows chat_session to refer to either a notebook or a source
DEFINE TABLE OVERWRITE refers_to
TYPE RELATION
FROM chat_session TO notebook|source;
-- Add model_override field to chat_session for per-session model selection
DEFINE FIELD model_override ON chat_session TYPE option<string>;
DEFINE FIELD command ON source TYPE option<record<command>>;
@@ -0,0 +1,11 @@
-- Rollback Migration 8: Revert to notebook-only chat sessions
DEFINE TABLE OVERWRITE refers_to
TYPE RELATION
FROM chat_session TO notebook;
-- Remove model_override field from chat_session
REMOVE FIELD model_override ON chat_session;
REMOVE FIELD command ON source;
@@ -0,0 +1,66 @@
REMOVE FUNCTION IF EXISTS fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
let $source_embedding_search =
IF $sources {(
SELECT
source.id as id,
source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding
WHERE embedding != none and array::len(embedding)=array::len($query) AND
vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT
id,
insight_type + ' - ' + (source.title OR '') as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight
WHERE embedding != none and array::len(embedding)=array::len($query) AND
vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT
id,
title,
content,
id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM note
WHERE embedding != none and array::len(embedding)=array::len($query) AND
vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $all_results = array::union(
array::union($source_embedding_search, $source_insight_search),
$note_content_search
);
RETURN (select id, parent_id, title, math::max(similarity) as similarity,
array::flatten(content) as matches
from $all_results where id is not None
group by id, parent_id, title ORDER BY similarity DESC LIMIT $match_count);
};
@@ -0,0 +1,63 @@
REMOVE FUNCTION IF EXISTS fn::vector_search;
DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_count: int, $sources: bool, $show_notes: bool, $min_similarity: float) {
let $source_embedding_search =
IF $sources {(
SELECT
source.id as id,
source.title as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_embedding
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $source_insight_search =
IF $sources {(
SELECT
id,
insight_type + ' - ' + (source.title OR '') as title,
content,
source.id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM source_insight
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $note_content_search =
IF $show_notes {(
SELECT
id,
title,
content,
id as parent_id,
vector::similarity::cosine(embedding, $query) as similarity
FROM note
WHERE vector::similarity::cosine(embedding, $query) >= $min_similarity
ORDER BY similarity DESC
LIMIT $match_count
)}
ELSE { [] };
let $all_results = array::union(
array::union($source_embedding_search, $source_insight_search),
$note_content_search
);
RETURN (select id, parent_id, title, math::max(similarity) as similarity,
array::flatten(content) as matches
from $all_results where id is not None
group by id, parent_id, title ORDER BY similarity DESC LIMIT $match_count);
};
+237
View File
@@ -0,0 +1,237 @@
import os
import re
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, TypeVar, Union
from loguru import logger
from surrealdb import AsyncSurreal, RecordID # type: ignore
from surrealdb.data.types.table import Table # type: ignore
T = TypeVar("T", Dict[str, Any], List[Dict[str, Any]])
# Bare SurrealDB table/relation identifier: no ':', whitespace, or query
# syntax. Used to validate the parts of RELATE/UPSERT/UPDATE that name a
# table or edge-relation and therefore can't be bound as a query parameter
# (SurrealQL only allows binding record/table *values*, not identifiers in
# that position).
_IDENTIFIER_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
def _ensure_safe_identifier(value: str, kind: str) -> str:
"""Validate a table/relationship name before it is interpolated into a query."""
if not isinstance(value, str) or not _IDENTIFIER_RE.match(value):
raise ValueError(f"Invalid {kind} name: {value!r}")
return value
def _get_env_or_default(name: str, default: str) -> str:
value = os.getenv(name)
return value if value else default
def get_database_url() -> str:
"""Get database URL with backward compatibility"""
surreal_url = os.getenv("SURREAL_URL")
if surreal_url:
return surreal_url
# Fallback to old format - WebSocket URL format
address = os.getenv("SURREAL_ADDRESS", "localhost")
port = os.getenv("SURREAL_PORT", "8000")
return f"ws://{address}/rpc:{port}"
def get_database_password() -> str:
"""Get password with backward compatibility"""
return os.getenv("SURREAL_PASSWORD") or os.getenv("SURREAL_PASS") or "root"
def get_database_namespace() -> str:
"""Get configured SurrealDB namespace."""
return _get_env_or_default("SURREAL_NAMESPACE", "open_notebook")
def get_database_name() -> str:
"""Get configured SurrealDB database name."""
return _get_env_or_default("SURREAL_DATABASE", "open_notebook")
def parse_record_ids(obj: Any) -> Any:
"""Recursively parse and convert RecordIDs into strings."""
if isinstance(obj, dict):
return {k: parse_record_ids(v) for k, v in obj.items()}
elif isinstance(obj, list):
return [parse_record_ids(item) for item in obj]
elif isinstance(obj, RecordID):
return str(obj)
return obj
def ensure_record_id(value: Union[str, RecordID]) -> RecordID:
"""Ensure a value is a RecordID."""
if isinstance(value, RecordID):
return value
return RecordID.parse(value)
@asynccontextmanager
async def db_connection():
db = AsyncSurreal(get_database_url())
await db.signin(
{
"username": os.environ.get("SURREAL_USER"),
"password": get_database_password(),
}
)
await db.use(get_database_namespace(), get_database_name())
try:
yield db
finally:
await db.close()
async def repo_query(
query_str: str, vars: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
"""Execute a SurrealQL query and return the results"""
async with db_connection() as connection:
try:
result = parse_record_ids(await connection.query(query_str, vars))
if isinstance(result, str):
raise RuntimeError(result)
return result
except RuntimeError as e:
# RuntimeError is raised for retriable transaction conflicts - log at debug to avoid noise
logger.debug(str(e))
raise
except Exception as e:
logger.exception(e)
raise
async def repo_create(table: str, data: Dict[str, Any]) -> Dict[str, Any]:
"""Create a new record in the specified table"""
# Remove 'id' attribute if it exists in data
data.pop("id", None)
data["created"] = datetime.now(timezone.utc)
data["updated"] = datetime.now(timezone.utc)
try:
async with db_connection() as connection:
result = parse_record_ids(await connection.insert(table, data))
# SurrealDB may return a string error message instead of the expected record
if isinstance(result, str):
raise RuntimeError(result)
return result
except RuntimeError as e:
logger.error(str(e))
raise
except Exception as e:
logger.exception(e)
raise RuntimeError("Failed to create record")
async def repo_relate(
source: Union[str, RecordID],
relationship: str,
target: Union[str, RecordID],
data: Optional[Dict[str, Any]] = None,
) -> List[Dict[str, Any]]:
"""Create a relationship between two records with optional data"""
if data is None:
data = {}
# relationship is an edge-table name, not a record value, so it can't be
# bound as a query parameter; validate it against an identifier allowlist
# instead of trusting the caller. source/target are always bound.
_ensure_safe_identifier(relationship, "relationship")
query = f"RELATE $source->{relationship}->$target CONTENT $data;"
# logger.debug(f"Relate query: {query}")
return await repo_query(
query,
{
"source": ensure_record_id(source),
"target": ensure_record_id(target),
"data": data,
},
)
async def repo_upsert(
table: str, id: Optional[str], data: Dict[str, Any], add_timestamp: bool = False
) -> List[Dict[str, Any]]:
"""Create or update a record in the specified table"""
data.pop("id", None)
if add_timestamp:
data["updated"] = datetime.now(timezone.utc)
_ensure_safe_identifier(table, "table")
target: Union[RecordID, Table] = ensure_record_id(id) if id else Table(table)
query = "UPSERT $target MERGE $data;"
return await repo_query(query, {"target": target, "data": data})
async def repo_update(
table: str, id: Union[str, RecordID], data: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Update an existing record by table and id"""
# If id already contains the table name, use it as is
try:
_ensure_safe_identifier(table, "table")
if isinstance(id, RecordID):
record_id: RecordID = id
elif ":" in id and id.startswith(f"{table}:"):
record_id = ensure_record_id(id)
else:
record_id = RecordID(table, id)
data.pop("id", None)
if "created" in data and isinstance(data["created"], str):
data["created"] = datetime.fromisoformat(data["created"])
data["updated"] = datetime.now(timezone.utc)
query = "UPDATE $target MERGE $data;"
# logger.debug(f"Update query: {query}")
result = await repo_query(query, {"target": record_id, "data": data})
# if isinstance(result, list):
# return [_return_data(item) for item in result]
return parse_record_ids(result)
except Exception as e:
raise RuntimeError(f"Failed to update record: {str(e)}")
async def repo_delete(record_id: Union[str, RecordID]):
"""Delete a record by record id"""
try:
async with db_connection() as connection:
return await connection.delete(ensure_record_id(record_id))
except Exception as e:
logger.exception(e)
raise RuntimeError(f"Failed to delete record: {str(e)}")
async def repo_insert(
table: str, data: List[Dict[str, Any]], ignore_duplicates: bool = False
) -> List[Dict[str, Any]]:
"""Create a new record in the specified table"""
try:
async with db_connection() as connection:
result = parse_record_ids(await connection.insert(table, data))
# SurrealDB may return a string error message instead of the expected records
if isinstance(result, str):
raise RuntimeError(result)
return result
except RuntimeError as e:
if ignore_duplicates and "already contains" in str(e):
return []
# Log transaction conflicts at debug level (they are expected during concurrent operations)
error_str = str(e).lower()
if "transaction" in error_str or "conflict" in error_str:
logger.debug(str(e))
else:
logger.error(str(e))
raise
except Exception as e:
if ignore_duplicates and "already contains" in str(e):
return []
logger.exception(e)
raise RuntimeError("Failed to create record")
+7
View File
@@ -0,0 +1,7 @@
"""
Domain models for Open Notebook.
This module exports the core domain models used throughout the application.
"""
__all__: list[str] = []
+363
View File
@@ -0,0 +1,363 @@
import re
from datetime import datetime
from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union, cast
from loguru import logger
from pydantic import (
BaseModel,
ConfigDict,
ValidationError,
field_validator,
model_validator,
)
from open_notebook.database.repository import (
ensure_record_id,
repo_create,
repo_delete,
repo_query,
repo_relate,
repo_update,
repo_upsert,
)
from open_notebook.exceptions import (
DatabaseOperationError,
InvalidInputError,
NotFoundError,
)
T = TypeVar("T", bound="ObjectModel")
class ObjectModel(BaseModel):
id: Optional[str] = None
table_name: ClassVar[str] = ""
nullable_fields: ClassVar[set[str]] = set() # Fields that can be saved as None
created: Optional[datetime] = None
updated: Optional[datetime] = None
@classmethod
def _validate_order_by(cls, order_by: str) -> str:
"""Validate and normalize an ORDER BY clause to prevent SurrealQL injection.
Supports: "field", "field direction", "field1 direction, field2 direction".
Any subclass that builds its own query around `order_by` (instead of
delegating to `get_all()`) must route through this so the allowlist
can't silently drift between call sites.
"""
allowed_field_pattern = re.compile(r"^[a-z_][a-z0-9_]*$")
allowed_directions = {"asc", "desc"}
clauses = [c.strip() for c in order_by.split(",")]
validated_clauses = []
for clause in clauses:
parts = clause.strip().split()
if len(parts) == 1:
if not allowed_field_pattern.match(parts[0].lower()):
raise InvalidInputError(f"Invalid order_by field: '{parts[0]}'")
validated_clauses.append(parts[0].lower())
elif len(parts) == 2:
if not allowed_field_pattern.match(
parts[0].lower()
) or parts[1].lower() not in allowed_directions:
raise InvalidInputError(
f"Invalid order_by clause: '{clause.strip()}'"
)
validated_clauses.append(f"{parts[0].lower()} {parts[1].lower()}")
else:
raise InvalidInputError(f"Invalid order_by clause: '{clause.strip()}'")
return ", ".join(validated_clauses)
@classmethod
async def get_all(cls: Type[T], order_by=None) -> List[T]:
try:
# If called from a specific subclass, use its table_name
if cls.table_name:
target_class = cls
table_name = cls.table_name
else:
# This path is taken if called directly from ObjectModel
raise InvalidInputError(
"get_all() must be called from a specific model class"
)
if order_by:
validated_order_by = cls._validate_order_by(order_by)
query = f"SELECT * FROM {table_name} ORDER BY {validated_order_by}"
else:
query = f"SELECT * FROM {table_name}"
result = await repo_query(query)
objects = []
for obj in result:
try:
objects.append(target_class(**obj))
except Exception as e:
logger.critical(f"Error creating object: {str(e)}")
return objects
except Exception as e:
logger.error(f"Error fetching all {cls.table_name}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
@classmethod
async def get(cls: Type[T], id: str) -> T:
if not id:
raise InvalidInputError("ID cannot be empty")
try:
# Get the table name from the ID (everything before the first colon)
table_name = id.split(":")[0] if ":" in id else id
# If we're calling from a specific subclass and IDs match, use that class
if cls.table_name and cls.table_name == table_name:
target_class: Type[T] = cls
else:
# Otherwise, find the appropriate subclass based on table_name
found_class = cls._get_class_by_table_name(table_name)
if not found_class:
raise InvalidInputError(f"No class found for table {table_name}")
target_class = cast(Type[T], found_class)
result = await repo_query("SELECT * FROM $id", {"id": ensure_record_id(id)})
if result:
return target_class(**result[0])
else:
raise NotFoundError(f"{table_name} with id {id} not found")
except Exception as e:
logger.error(f"Error fetching object with id {id}: {str(e)}")
logger.exception(e)
raise NotFoundError(f"Object with id {id} not found - {str(e)}")
@classmethod
def _get_class_by_table_name(cls, table_name: str) -> Optional[Type["ObjectModel"]]:
"""Find the appropriate subclass based on table_name."""
def get_all_subclasses(c: Type["ObjectModel"]) -> List[Type["ObjectModel"]]:
all_subclasses: List[Type["ObjectModel"]] = []
for subclass in c.__subclasses__():
all_subclasses.append(subclass)
all_subclasses.extend(get_all_subclasses(subclass))
return all_subclasses
for subclass in get_all_subclasses(ObjectModel):
if hasattr(subclass, "table_name") and subclass.table_name == table_name:
return subclass
return None
async def save(self) -> None:
"""
Save the model to the database.
Note: Embedding is no longer generated inline. Subclasses that need
embedding should override save() to submit the appropriate embed_*
command after calling super().save().
"""
try:
self.model_validate(self.model_dump(), strict=True)
data = self._prepare_save_data()
data["updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
repo_result: Union[List[Dict[str, Any]], Dict[str, Any]]
if self.id is None:
data["created"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
repo_result = await repo_create(self.__class__.table_name, data)
else:
data["created"] = (
self.created.strftime("%Y-%m-%d %H:%M:%S")
if isinstance(self.created, datetime)
else self.created
)
logger.debug(f"Updating record with id {self.id}")
repo_result = await repo_update(
self.__class__.table_name, self.id, data
)
# Update the current instance with the result
# repo_result is a list of dictionaries
result_list: List[Dict[str, Any]] = (
repo_result if isinstance(repo_result, list) else [repo_result]
)
for key, value in result_list[0].items():
if hasattr(self, key):
if isinstance(getattr(self, key), BaseModel):
setattr(self, key, type(getattr(self, key))(**value))
else:
setattr(self, key, value)
except ValidationError as e:
logger.error(f"Validation failed: {e}")
raise
except RuntimeError:
# Transaction conflicts should propagate for retry
raise
except Exception as e:
logger.error(f"Error saving record: {e}")
raise DatabaseOperationError(e)
def _prepare_save_data(self) -> Dict[str, Any]:
data = self.model_dump()
return {
key: value
for key, value in data.items()
if value is not None or key in self.__class__.nullable_fields
}
async def delete(self) -> bool:
if self.id is None:
raise InvalidInputError("Cannot delete object without an ID")
try:
logger.debug(f"Deleting record with id {self.id}")
return await repo_delete(self.id)
except Exception as e:
logger.error(
f"Error deleting {self.__class__.table_name} with id {self.id}: {str(e)}"
)
raise DatabaseOperationError(
f"Failed to delete {self.__class__.table_name}"
)
async def relate(
self, relationship: str, target_id: str, data: Optional[Dict] = {}
) -> Any:
if not relationship or not target_id or not self.id:
raise InvalidInputError("Relationship and target ID must be provided")
try:
return await repo_relate(
source=self.id, relationship=relationship, target=target_id, data=data
)
except Exception as e:
logger.error(f"Error creating relationship: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
@field_validator("created", "updated", mode="before")
@classmethod
def parse_datetime(cls, value):
if isinstance(value, str):
return datetime.fromisoformat(value.replace("Z", "+00:00"))
return value
class RecordModel(BaseModel):
model_config = ConfigDict(
validate_assignment=True,
arbitrary_types_allowed=True,
extra="allow",
from_attributes=True,
defer_build=True,
)
record_id: ClassVar[str]
auto_save: ClassVar[bool] = (
False # Default to False, can be overridden in subclasses
)
_instances: ClassVar[Dict[str, "RecordModel"]] = {} # Store instances by record_id
def __new__(cls, **kwargs):
# If an instance already exists for this record_id, return it
if cls.record_id in cls._instances:
instance = cls._instances[cls.record_id]
# Update instance with any new kwargs if provided
if kwargs:
for key, value in kwargs.items():
setattr(instance, key, value)
return instance
# If no instance exists, create a new one
instance = super().__new__(cls)
cls._instances[cls.record_id] = instance
return instance
def __init__(self, **kwargs):
# Only initialize if this is a new instance
if not hasattr(self, "_initialized"):
object.__setattr__(self, "__dict__", {})
# For RecordModel, we need to handle async initialization differently
# Initialize with provided kwargs only for now
super().__init__(**kwargs)
# Mark as initialized but not loaded from DB yet
object.__setattr__(self, "_initialized", True)
object.__setattr__(self, "_db_loaded", False)
async def _load_from_db(self):
"""Load data from database if not already loaded"""
if not getattr(self, "_db_loaded", False):
result = await repo_query(
"SELECT * FROM ONLY $record_id",
{"record_id": ensure_record_id(self.record_id)},
)
# Handle case where record doesn't exist yet
if result:
if isinstance(result, list) and len(result) > 0:
# Standard list response
row = result[0]
if isinstance(row, dict):
for key, value in row.items():
if hasattr(self, key):
object.__setattr__(self, key, value)
elif isinstance(result, dict):
# Direct dict response
for key, value in result.items():
if hasattr(self, key):
object.__setattr__(self, key, value)
object.__setattr__(self, "_db_loaded", True)
@classmethod
async def get_instance(cls) -> "RecordModel":
"""Get or create the singleton instance and load from DB"""
instance = cls()
await instance._load_from_db()
return instance
@model_validator(mode="after")
def auto_save_validator(self):
if self.__class__.auto_save:
# Auto-save can't work with async - log warning
logger.warning(
f"Auto-save is enabled for {self.__class__.__name__} but update() is now async. Call await instance.update() manually."
)
return self
async def update(self):
# Get all non-ClassVar fields and their values
data = {
field_name: getattr(self, field_name)
for field_name, field_info in self.model_fields.items()
if not str(field_info.annotation).startswith("typing.ClassVar")
}
await repo_upsert(
self.__class__.table_name
if hasattr(self.__class__, "table_name")
else "record",
self.record_id,
data,
)
result = await repo_query(
"SELECT * FROM $record_id", {"record_id": ensure_record_id(self.record_id)}
)
if result:
for key, value in result[0].items():
if hasattr(self, key):
object.__setattr__(
self, key, value
) # Use object.__setattr__ to avoid triggering validation again
return self
@classmethod
def clear_instance(cls):
"""Clear the singleton instance (useful for testing)"""
if cls.record_id in cls._instances:
del cls._instances[cls.record_id]
async def patch(self, model_dict: dict):
"""Update model attributes from dictionary and save"""
for key, value in model_dict.items():
setattr(self, key, value)
await self.update()
+32
View File
@@ -0,0 +1,32 @@
from typing import ClassVar, List, Literal, Optional
from pydantic import Field
from open_notebook.domain.base import RecordModel
class ContentSettings(RecordModel):
record_id: ClassVar[str] = "open_notebook:content_settings"
default_content_processing_engine_doc: Optional[
Literal["auto", "docling", "simple"]
] = Field("auto", description="Default Content Processing Engine for Documents")
default_content_processing_engine_url: Optional[
Literal["auto", "firecrawl", "jina", "crawl4ai", "simple"]
] = Field("auto", description="Default Content Processing Engine for URLs")
default_embedding_option: Optional[Literal["ask", "always", "never"]] = Field(
"ask", description="Default Embedding Option for Vector Search"
)
auto_delete_files: Optional[Literal["yes", "no"]] = Field(
"yes", description="Auto Delete Uploaded Files"
)
docling_ocr: Optional[bool] = Field(
True,
description=(
"Run OCR on scanned PDFs and images when the Docling engine handles "
"them. Disable for faster processing of text-native documents."
),
)
youtube_preferred_languages: Optional[List[str]] = Field(
["en", "pt", "es", "de", "nl", "en-GB", "fr", "de", "hi", "ja"],
description="Preferred languages for YouTube transcripts",
)
+286
View File
@@ -0,0 +1,286 @@
"""
Credential domain model for storing individual provider credentials.
Each credential is a standalone record in the 'credential' table, replacing
the old ProviderConfig singleton. Credentials store API keys (encrypted at
rest) and provider-specific configuration fields.
Usage:
cred = Credential(
name="Production",
provider="openai",
modalities=["language", "embedding"],
api_key=SecretStr("sk-..."),
)
await cred.save()
"""
from typing import Any, ClassVar, Dict, List, Optional
from loguru import logger
from pydantic import SecretStr, model_validator
from open_notebook.database.repository import ensure_record_id, repo_query
from open_notebook.domain.base import ObjectModel
from open_notebook.utils.encryption import decrypt_value, encrypt_value
class Credential(ObjectModel):
"""
Individual credential record for an AI provider.
Each record stores authentication and configuration for a single provider
account. Models link to credentials via the credential field.
"""
table_name: ClassVar[str] = "credential"
nullable_fields: ClassVar[set[str]] = {
"api_key",
"base_url",
"endpoint",
"api_version",
"endpoint_llm",
"endpoint_embedding",
"endpoint_stt",
"endpoint_tts",
"project",
"location",
"credentials_path",
}
# Provider-specific tuning options that are exposed as top-level fields on
# this model but persisted inside the flexible `config` object on the
# `credential` table (migration 15). Adding a new such option only requires
# declaring the Pydantic field below and listing it here — no DB migration.
CONFIG_EXTRAS: ClassVar[set[str]] = {"num_ctx"}
name: str
provider: str
modalities: List[str] = []
api_key: Optional[SecretStr] = None
decryption_error: Optional[str] = None
base_url: Optional[str] = None
endpoint: Optional[str] = None
api_version: Optional[str] = None
endpoint_llm: Optional[str] = None
endpoint_embedding: Optional[str] = None
endpoint_stt: Optional[str] = None
endpoint_tts: Optional[str] = None
project: Optional[str] = None
location: Optional[str] = None
credentials_path: Optional[str] = None
# Flexible provider config bag, persisted as-is to the credential table's
# `config` FLEXIBLE object (migration 15). This is the source of truth on
# disk and holds both the keys this version maps (see CONFIG_EXTRAS) and any
# keys written by a newer version, so a load/save round-trip never clobbers
# options this version doesn't know about.
config: Optional[Dict[str, Any]] = None
# Ollama-only: overrides the context window (num_ctx). Esperanto defaults to
# 8192; raise this if your hardware can handle a larger context window.
# Exposed as a top-level field (and on the API) for convenience; it mirrors
# config["num_ctx"].
num_ctx: Optional[int] = None
@model_validator(mode="before")
@classmethod
def _mirror_config_to_fields(cls, data: Any) -> Any:
"""Mirror known keys from the persisted `config` bag onto their
convenience fields on load. Done in `before` so the values flow through
normal Pydantic field validation (e.g. `num_ctx` is coerced/validated as
an int) instead of being set raw. Unknown keys stay in `config` untouched
and are preserved on save."""
if isinstance(data, dict) and isinstance(data.get("config"), dict):
config = data["config"]
data = dict(data)
for key in cls.CONFIG_EXTRAS:
if data.get(key) is None and config.get(key) is not None:
data[key] = config[key]
return data
def to_esperanto_config(self) -> Dict[str, Any]:
"""
Build config dict for AIFactory.create_*() calls.
Returns a dict that can be passed as the 'config' parameter to
Esperanto's AIFactory methods, overriding env var lookup.
"""
config: Dict[str, Any] = {}
if self.api_key:
config["api_key"] = self.api_key.get_secret_value()
if self.base_url:
config["base_url"] = self.base_url
# For Azure, base_url from the UI form maps to endpoint
if self.provider and self.provider.lower() == "azure" and not self.endpoint:
config["endpoint"] = self.base_url
if self.endpoint:
config["endpoint"] = self.endpoint
if self.api_version:
config["api_version"] = self.api_version
if self.endpoint_llm:
config["endpoint_llm"] = self.endpoint_llm
if self.endpoint_embedding:
config["endpoint_embedding"] = self.endpoint_embedding
if self.endpoint_stt:
config["endpoint_stt"] = self.endpoint_stt
if self.endpoint_tts:
config["endpoint_tts"] = self.endpoint_tts
if self.project:
config["project"] = self.project
if self.location:
config["location"] = self.location
if self.credentials_path:
config["credentials_path"] = self.credentials_path
if self.num_ctx is not None:
config["num_ctx"] = self.num_ctx
return config
@classmethod
async def get_by_provider(cls, provider: str) -> List["Credential"]:
"""Get all credentials for a provider."""
results = await repo_query(
"SELECT * FROM credential WHERE string::lowercase(provider) = string::lowercase($provider) ORDER BY created ASC",
{"provider": provider},
)
credentials = []
for row in results:
try:
cred = cls._from_db_row(row)
credentials.append(cred)
except Exception as e:
logger.warning(f"Skipping invalid credential: {e}")
return credentials
@classmethod
async def get(cls, id: str) -> "Credential":
"""Override get() to handle api_key decryption."""
instance = await super().get(id)
# Pydantic auto-wraps the raw DB string in SecretStr, so we need
# to extract, decrypt, and re-wrap regardless of type.
if instance.api_key:
raw = (
instance.api_key.get_secret_value()
if isinstance(instance.api_key, SecretStr)
else instance.api_key
)
decrypted = decrypt_value(raw)
object.__setattr__(instance, "api_key", SecretStr(decrypted))
return instance
@classmethod
async def get_all(cls, order_by=None) -> List["Credential"]:
"""Override get_all() to handle api_key decryption with per-row error handling."""
if order_by:
validated_order_by = cls._validate_order_by(order_by)
query = f"SELECT * FROM {cls.table_name} ORDER BY {validated_order_by}"
else:
query = f"SELECT * FROM {cls.table_name}"
results = await repo_query(query, {})
credentials = []
for row in results:
try:
cred = cls._from_db_row(row)
credentials.append(cred)
except Exception as e:
logger.warning(
f"Failed to decrypt credential {row.get('id', 'unknown')}: {e}"
)
# Create a minimal credential with error info from raw DB fields
try:
error_cred = cls(
name=row.get("name", "Unknown"),
provider=row.get("provider", "unknown"),
modalities=row.get("modalities", []),
decryption_error="Failed to decrypt API key. The encryption key may have changed.",
)
# Preserve the DB id, created, updated from the raw row
if row.get("id"):
object.__setattr__(error_cred, "id", str(row["id"]))
if row.get("created"):
object.__setattr__(error_cred, "created", row["created"])
if row.get("updated"):
object.__setattr__(error_cred, "updated", row["updated"])
# Mark that it had an api_key (even though we can't decrypt it)
if row.get("api_key"):
object.__setattr__(
error_cred, "api_key", SecretStr("UNDECRYPTABLE")
)
credentials.append(error_cred)
except Exception as inner_e:
logger.error(
f"Failed to create error credential for {row.get('id', 'unknown')}: {inner_e}"
)
return credentials
async def get_linked_models(self) -> list:
"""Get all models linked to this credential."""
if not self.id:
return []
from open_notebook.ai.models import Model
results = await repo_query(
"SELECT * FROM model WHERE credential = $cred_id",
{"cred_id": ensure_record_id(self.id)},
)
return [Model(**row) for row in results]
def _prepare_save_data(self) -> Dict[str, Any]:
"""Override to encrypt api_key and sync provider extras into `config`."""
data: Dict[str, Any] = {}
for key, value in self.model_dump().items():
if key in ("decryption_error", "config"):
# `config` is rebuilt below from the existing bag + convenience fields.
continue
if key == "api_key":
# Handle SecretStr: extract, encrypt, store
if self.api_key:
secret_value = self.api_key.get_secret_value()
data["api_key"] = encrypt_value(secret_value)
else:
data["api_key"] = None
elif value is not None or key in self.__class__.nullable_fields:
data[key] = value
# Sync the convenience fields (num_ctx, ...) into the flexible `config`
# object so the SCHEMAFULL credential table doesn't drop them (migration
# 15). Starting from the existing bag preserves any keys written by a
# newer version, since repo_update MERGE replaces the whole object.
# `config` is None only when the merged result is genuinely empty.
config: Dict[str, Any] = dict(self.config or {})
for key in self.__class__.CONFIG_EXTRAS:
data.pop(key, None) # not a top-level column; lives in `config`
value = getattr(self, key, None)
if value is not None:
config[key] = value
else:
config.pop(key, None)
data["config"] = config or None
return data
async def save(self) -> None:
"""Save credential, handling api_key re-hydration after DB round-trip."""
# Remember the original SecretStr before save
original_api_key = self.api_key
await super().save()
# After save, the api_key field may be set to the encrypted string
# from the DB result. Restore the original SecretStr.
if original_api_key:
object.__setattr__(self, "api_key", original_api_key)
elif self.api_key and isinstance(self.api_key, str):
# Decrypt if DB returned an encrypted string
decrypted = decrypt_value(self.api_key)
object.__setattr__(self, "api_key", SecretStr(decrypted))
@classmethod
def _from_db_row(cls, row: dict) -> "Credential":
"""Create a Credential from a database row, decrypting api_key."""
api_key_val = row.get("api_key")
if api_key_val and isinstance(api_key_val, str):
decrypted = decrypt_value(api_key_val)
row["api_key"] = SecretStr(decrypted)
elif api_key_val is None:
row["api_key"] = None
return cls(**row)
+827
View File
@@ -0,0 +1,827 @@
import os
from datetime import datetime
from pathlib import Path
from typing import Any, ClassVar, Dict, List, Literal, Optional, Union
from loguru import logger
from pydantic import BaseModel, ConfigDict, Field, field_validator
from surreal_commands import submit_command
from surrealdb import RecordID
from open_notebook.database.repository import ensure_record_id, repo_query
from open_notebook.domain.base import ObjectModel
from open_notebook.exceptions import DatabaseOperationError, InvalidInputError
class Notebook(ObjectModel):
table_name: ClassVar[str] = "notebook"
name: str
description: str
archived: Optional[bool] = False
last_viewed_at: Optional[datetime] = None
@field_validator("name")
@classmethod
def name_must_not_be_empty(cls, v):
if not v.strip():
raise InvalidInputError("Notebook name cannot be empty")
return v
async def get_sources(self, include_full_text: bool = False) -> List["Source"]:
try:
source_projection = "" if include_full_text else " omit source.full_text"
srcs = await repo_query(
f"""
select *{source_projection} from (
select in as source from reference where out=$id
fetch source
) order by source.updated desc
""",
{"id": ensure_record_id(self.id)},
)
return [Source(**src["source"]) for src in srcs] if srcs else []
except Exception as e:
logger.error(f"Error fetching sources for notebook {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
async def get_notes(self, include_content: bool = False) -> List["Note"]:
try:
note_projection = (
" omit note.embedding"
if include_content
else " omit note.content, note.embedding"
)
srcs = await repo_query(
f"""
select *{note_projection} from (
select in as note from artifact where out=$id
fetch note
) order by note.updated desc
""",
{"id": ensure_record_id(self.id)},
)
return [Note(**src["note"]) for src in srcs] if srcs else []
except Exception as e:
logger.error(f"Error fetching notes for notebook {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
async def get_context(self) -> str:
"""
Build long-form notebook context for podcast and LLM workflows.
Normal list retrieval omits large source/note bodies, so this method uses
opt-in full-content fetches and formats only substantive context blocks.
"""
sources = await self.get_sources(include_full_text=True)
notes = await self.get_notes(include_content=True)
context_blocks = []
insights_by_source = await SourceInsight.get_for_sources(
[source.id for source in sources if source.id]
)
for source in sources:
source_context = await source.get_context(
context_size="long",
insights=insights_by_source.get(source.id or "", []),
)
if isinstance(source_context, dict):
title = source_context.get("title") or source.title or "Untitled source"
full_text = source_context.get("full_text")
insights = source_context.get("insights") or []
content_parts = []
if full_text:
content_parts.append(str(full_text))
insight_lines = []
for insight in insights:
if not isinstance(insight, dict):
continue
insight_content = insight.get("content")
if not insight_content:
continue
insight_type = insight.get("insight_type") or "Insight"
insight_lines.append(f"- {insight_type}: {insight_content}")
if insight_lines:
content_parts.append("Insights:\n" + "\n".join(insight_lines))
content = "\n\n".join(content_parts).strip()
else:
title = source.title or "Untitled source"
content = str(source_context).strip()
if content:
context_blocks.append(f"## Source: {title}\n\n{content}")
for note in notes:
note_context = note.get_context(context_size="long")
if isinstance(note_context, dict):
title = note_context.get("title") or note.title or "Untitled note"
content = note_context.get("content")
content = str(content).strip() if content else ""
else:
title = note.title or "Untitled note"
content = str(note_context).strip()
if content:
context_blocks.append(f"## Note: {title}\n\n{content}")
return "\n\n".join(context_blocks)
async def get_chat_sessions(self) -> List["ChatSession"]:
try:
srcs = await repo_query(
"""
select * from (
select
<- chat_session as chat_session
from refers_to
where out=$id
fetch chat_session
)
order by chat_session.updated desc
""",
{"id": ensure_record_id(self.id)},
)
return (
[ChatSession(**src["chat_session"][0]) for src in srcs] if srcs else []
)
except Exception as e:
logger.error(
f"Error fetching chat sessions for notebook {self.id}: {str(e)}"
)
logger.exception(e)
raise DatabaseOperationError(e)
async def get_delete_preview(self) -> Dict[str, Any]:
"""
Get counts of items that would be affected by deleting this notebook.
Returns a dict with:
- note_count: Number of notes that will be deleted
- exclusive_source_count: Sources only in this notebook (can be deleted)
- shared_source_count: Sources in other notebooks (will be unlinked only)
"""
try:
notebook_id = ensure_record_id(self.id)
# Count notes
note_result = await repo_query(
"SELECT count() as count FROM artifact WHERE out = $notebook_id GROUP ALL",
{"notebook_id": notebook_id},
)
note_count = note_result[0]["count"] if note_result else 0
# Get sources with count of references to OTHER notebooks
# If assigned_others = 0, source is exclusive to this notebook
# If assigned_others > 0, source is shared with other notebooks
source_counts = await repo_query(
"""
SELECT
id,
count(->reference[WHERE out != $notebook_id].out) as assigned_others
FROM (SELECT VALUE <-reference.in AS sources FROM $notebook_id)[0]
""",
{"notebook_id": notebook_id},
)
exclusive_count = 0
shared_count = 0
for src in source_counts:
if src.get("assigned_others", 0) == 0:
exclusive_count += 1
else:
shared_count += 1
return {
"note_count": note_count,
"exclusive_source_count": exclusive_count,
"shared_source_count": shared_count,
}
except Exception as e:
logger.error(f"Error getting delete preview for notebook {self.id}: {e}")
logger.exception(e)
raise DatabaseOperationError(e)
async def delete(self, delete_exclusive_sources: bool = False) -> Dict[str, int]:
"""
Delete notebook with cascade deletion of notes and optional source deletion.
Args:
delete_exclusive_sources: If True, also delete sources that belong
only to this notebook. Default is False.
Returns:
Dict with counts: deleted_notes, deleted_sources, unlinked_sources
"""
if self.id is None:
raise InvalidInputError("Cannot delete notebook without an ID")
try:
notebook_id = ensure_record_id(self.id)
deleted_notes = 0
deleted_sources = 0
unlinked_sources = 0
# 1. Get and delete all notes linked to this notebook
notes = await self.get_notes()
for note in notes:
await note.delete()
deleted_notes += 1
logger.info(f"Deleted {deleted_notes} notes for notebook {self.id}")
# Delete artifact relationships
await repo_query(
"DELETE artifact WHERE out = $notebook_id",
{"notebook_id": notebook_id},
)
# 2. Handle sources
if delete_exclusive_sources:
# Find sources with count of references to OTHER notebooks
# If assigned_others = 0, source is exclusive to this notebook
source_counts = await repo_query(
"""
SELECT
id,
count(->reference[WHERE out != $notebook_id].out) as assigned_others
FROM (SELECT VALUE <-reference.in AS sources FROM $notebook_id)[0]
""",
{"notebook_id": notebook_id},
)
for src in source_counts:
source_id = src.get("id")
if source_id and src.get("assigned_others", 0) == 0:
# Exclusive source - delete it
try:
source = await Source.get(str(source_id))
await source.delete()
deleted_sources += 1
except Exception as e:
logger.warning(
f"Failed to delete exclusive source {source_id}: {e}"
)
else:
unlinked_sources += 1
else:
# Just count sources that will be unlinked
source_result = await repo_query(
"SELECT count() as count FROM reference WHERE out = $notebook_id GROUP ALL",
{"notebook_id": notebook_id},
)
unlinked_sources = source_result[0]["count"] if source_result else 0
# Delete reference relationships (unlink all sources)
await repo_query(
"DELETE reference WHERE out = $notebook_id",
{"notebook_id": notebook_id},
)
logger.info(
f"Unlinked {unlinked_sources} sources, deleted {deleted_sources} "
f"exclusive sources for notebook {self.id}"
)
# 3. Delete the notebook record itself
await super().delete()
logger.info(f"Deleted notebook {self.id}")
return {
"deleted_notes": deleted_notes,
"deleted_sources": deleted_sources,
"unlinked_sources": unlinked_sources,
}
except Exception as e:
logger.error(f"Error deleting notebook {self.id}: {e}")
logger.exception(e)
raise DatabaseOperationError(f"Failed to delete notebook: {e}")
class Asset(BaseModel):
file_path: Optional[str] = None
url: Optional[str] = None
class SourceEmbedding(ObjectModel):
table_name: ClassVar[str] = "source_embedding"
content: str
async def get_source(self) -> "Source":
try:
src = await repo_query(
"""
select source.* from $id fetch source
""",
{"id": ensure_record_id(self.id)},
)
return Source(**src[0]["source"])
except Exception as e:
logger.error(f"Error fetching source for embedding {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
class SourceInsight(ObjectModel):
table_name: ClassVar[str] = "source_insight"
insight_type: str
content: str
@classmethod
async def get_for_sources(
cls, source_ids: List[str]
) -> Dict[str, List["SourceInsight"]]:
"""
Batch-fetch insights for many sources in a single query.
Building notebook/chat context otherwise calls get_insights() once
per source - fine for one source, but O(n) round trips (each paying
its own connection setup - no pooling in the repository layer) when
a caller loops over every source in a notebook.
"""
grouped: Dict[str, List[SourceInsight]] = {sid: [] for sid in source_ids if sid}
if not grouped:
return grouped
try:
result = await repo_query(
"SELECT * FROM source_insight WHERE source IN $source_ids",
{"source_ids": [ensure_record_id(sid) for sid in grouped]},
)
except Exception as e:
logger.error(f"Error batch-fetching insights for sources: {str(e)}")
logger.exception(e)
raise DatabaseOperationError("Failed to fetch insights for sources")
for row in result:
key = str(row.get("source"))
grouped.setdefault(key, []).append(cls(**row))
return grouped
async def get_source(self) -> "Source":
try:
src = await repo_query(
"""
select source.* from $id fetch source
""",
{"id": ensure_record_id(self.id)},
)
return Source(**src[0]["source"])
except Exception as e:
logger.error(f"Error fetching source for insight {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
async def save_as_note(self, notebook_id: Optional[str] = None) -> Any:
source = await self.get_source()
note = Note(
title=f"{self.insight_type} from source {source.title}",
content=self.content,
)
await note.save()
if notebook_id:
await note.add_to_notebook(notebook_id)
return note
class Source(ObjectModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
table_name: ClassVar[str] = "source"
asset: Optional[Asset] = None
title: Optional[str] = None
topics: Optional[List[str]] = Field(default_factory=list)
full_text: Optional[str] = None
last_viewed_at: Optional[datetime] = None
command: Optional[Union[str, RecordID]] = Field(
default=None, description="Link to surreal-commands processing job"
)
@field_validator("command", mode="before")
@classmethod
def parse_command(cls, value):
"""Parse command field to ensure RecordID format"""
if isinstance(value, str) and value:
return ensure_record_id(value)
return value
@field_validator("id", mode="before")
@classmethod
def parse_id(cls, value):
"""Parse id field to handle both string and RecordID inputs"""
if value is None:
return None
if isinstance(value, RecordID):
return str(value)
return str(value) if value else None
async def get_status(self) -> Optional[str]:
"""Get the processing status of the associated command"""
if not self.command:
return None
try:
from surreal_commands import get_command_status
status = await get_command_status(str(self.command))
return status.status if status else "unknown"
except Exception as e:
logger.warning(f"Failed to get command status for {self.command}: {e}")
return "unknown"
async def get_processing_progress(self) -> Optional[Dict[str, Any]]:
"""Get detailed processing information for the associated command"""
if not self.command:
return None
try:
from surreal_commands import get_command_status
status_result = await get_command_status(str(self.command))
if not status_result:
return None
# Extract execution metadata if available
result = getattr(status_result, "result", None)
execution_metadata = (
result.get("execution_metadata", {}) if isinstance(result, dict) else {}
)
return {
"status": status_result.status,
"started_at": execution_metadata.get("started_at"),
"completed_at": execution_metadata.get("completed_at"),
"error": getattr(status_result, "error_message", None),
"result": result,
}
except Exception as e:
logger.warning(f"Failed to get command progress for {self.command}: {e}")
return None
async def get_context(
self,
context_size: Literal["short", "long"] = "short",
insights: Optional[List["SourceInsight"]] = None,
) -> Dict[str, Any]:
# Callers looping over many sources can batch-fetch insights up front
# via SourceInsight.get_for_sources() and pass them in here, instead
# of paying a separate query per source.
insight_objects = insights if insights is not None else await self.get_insights()
insights = [insight.model_dump() for insight in insight_objects]
if context_size == "long":
return dict(
id=self.id,
title=self.title,
insights=insights,
full_text=self.full_text,
)
else:
return dict(id=self.id, title=self.title, insights=insights)
async def get_embedded_chunks(self) -> int:
try:
result = await repo_query(
"""
select count() as chunks from source_embedding where source=$id GROUP ALL
""",
{"id": ensure_record_id(self.id)},
)
if len(result) == 0:
return 0
return result[0]["chunks"]
except Exception as e:
logger.error(f"Error fetching chunks count for source {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(f"Failed to count chunks for source: {str(e)}")
async def get_insights(self) -> List[SourceInsight]:
try:
result = await repo_query(
"""
SELECT * FROM source_insight WHERE source=$id
""",
{"id": ensure_record_id(self.id)},
)
return [SourceInsight(**insight) for insight in result]
except Exception as e:
logger.error(f"Error fetching insights for source {self.id}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError("Failed to fetch insights for source")
async def add_to_notebook(self, notebook_id: str) -> Any:
if not notebook_id:
raise InvalidInputError("Notebook ID must be provided")
await Notebook.get(notebook_id) # raises NotFoundError if invalid/missing
return await self.relate("reference", notebook_id)
async def vectorize(self) -> str:
"""
Submit vectorization as a background job using the embed_source command.
This method leverages the job-based architecture to prevent HTTP connection
pool exhaustion when processing large documents. The embed_source command:
1. Detects content type from file path
2. Chunks text using content-type aware splitter
3. Generates all embeddings in batches
4. Bulk inserts source_embedding records
Returns:
str: The command/job ID that can be used to track progress via the commands API
Raises:
ValueError: If source has no text to vectorize
DatabaseOperationError: If job submission fails
"""
logger.info(f"Submitting embed_source job for source {self.id}")
try:
if not self.full_text or not self.full_text.strip():
raise ValueError(f"Source {self.id} has no text to vectorize")
# Submit the embed_source command
command_id = submit_command(
"open_notebook",
"embed_source",
{"source_id": str(self.id)},
)
command_id_str = str(command_id)
logger.info(
f"Embed source job submitted for source {self.id}: "
f"command_id={command_id_str}"
)
return command_id_str
except ValueError:
raise
except Exception as e:
logger.error(f"Failed to submit embed_source job for source {self.id}: {e}")
logger.exception(e)
raise DatabaseOperationError(e)
async def add_insight(self, insight_type: str, content: str) -> str:
"""
Submit insight creation as an async command (fire-and-forget).
Submits a create_insight command that handles database operations with
automatic retry logic for transaction conflicts. The command also submits
an embed_insight command for async embedding.
This method returns immediately after submitting the command - it does NOT
wait for the insight to be created. Use this for batch operations where
throughput is more important than immediate confirmation.
Args:
insight_type: Type/category of the insight
content: The insight content text
Returns:
command_id for optional tracking
Raises:
InvalidInputError: If insight_type or content is empty
DatabaseOperationError: If submitting the command fails. Matches
vectorize()'s contract - callers (transformation.py, source.py)
run inside surreal-commands jobs whose outer exception
handling already retries transient failures, so a swallowed
submission failure here previously meant a transformation
could report success while the insight was silently never
persisted.
"""
if not insight_type or not content:
raise InvalidInputError("Insight type and content must be provided")
try:
# Submit create_insight command (fire-and-forget)
# Command handles retries internally for transaction conflicts
command_id = submit_command(
"open_notebook",
"create_insight",
{
"source_id": str(self.id),
"insight_type": insight_type,
"content": content,
},
)
logger.info(
f"Submitted create_insight command {command_id} for source {self.id} "
f"(type={insight_type})"
)
return str(command_id)
except Exception as e:
logger.exception(f"Error submitting create_insight for source {self.id}: {e}")
raise DatabaseOperationError(e)
def _prepare_save_data(self) -> dict:
"""Override to ensure command field is always RecordID format for database"""
data = super()._prepare_save_data()
# Ensure command field is RecordID format if not None
if data.get("command") is not None:
data["command"] = ensure_record_id(data["command"])
return data
async def delete(self) -> bool:
"""Delete source and clean up associated file, embeddings, and insights."""
# Clean up uploaded file if it exists
if self.asset and self.asset.file_path:
file_path = Path(self.asset.file_path)
if file_path.exists():
try:
os.unlink(file_path)
logger.info(f"Deleted file for source {self.id}: {file_path}")
except Exception as e:
logger.warning(
f"Failed to delete file {file_path} for source {self.id}: {e}. "
"Continuing with database deletion."
)
else:
logger.debug(
f"File {file_path} not found for source {self.id}, skipping cleanup"
)
# Delete associated embeddings and insights to prevent orphaned records
try:
source_id = ensure_record_id(self.id)
await repo_query(
"DELETE source_embedding WHERE source = $source_id",
{"source_id": source_id},
)
await repo_query(
"DELETE source_insight WHERE source = $source_id",
{"source_id": source_id},
)
logger.debug(f"Deleted embeddings and insights for source {self.id}")
except Exception as e:
logger.warning(
f"Failed to delete embeddings/insights for source {self.id}: {e}. "
"Continuing with source deletion."
)
# Call parent delete to remove database record
return await super().delete()
class Note(ObjectModel):
table_name: ClassVar[str] = "note"
title: Optional[str] = None
note_type: Optional[Literal["human", "ai"]] = None
content: Optional[str] = None
@field_validator("content")
@classmethod
def content_must_not_be_empty(cls, v):
if v is not None and not v.strip():
raise InvalidInputError("Note content cannot be empty")
return v
async def save(self) -> Optional[str]:
"""
Save the note and submit embedding command.
Overrides ObjectModel.save() to submit an async embed_note command
after saving, instead of inline embedding.
Returns:
Optional[str]: The command_id if embedding was submitted, None
otherwise (either no content to embed, or submission failed)
"""
# Call parent save (without embedding)
await super().save()
# Submit embedding command (fire-and-forget) if note has content.
# Unlike Source.vectorize()/add_insight() (explicit, dedicated calls
# whose whole point is the submission), this runs automatically
# inside save() - the note itself is already durably saved above,
# so a submission hiccup here shouldn't fail an otherwise-successful
# save with a 500. Best-effort: log and move on.
if self.id and self.content and self.content.strip():
try:
command_id = submit_command(
"open_notebook",
"embed_note",
{"note_id": str(self.id)},
)
logger.debug(f"Submitted embed_note command {command_id} for {self.id}")
return command_id
except Exception as e:
logger.error(f"Failed to submit embed_note command for {self.id}: {e}")
return None
return None
async def add_to_notebook(self, notebook_id: str) -> Any:
if not notebook_id:
raise InvalidInputError("Notebook ID must be provided")
await Notebook.get(notebook_id) # raises NotFoundError if invalid/missing
return await self.relate("artifact", notebook_id)
def get_context(
self, context_size: Literal["short", "long"] = "short"
) -> Dict[str, Any]:
if context_size == "long":
return dict(id=self.id, title=self.title, content=self.content)
else:
return dict(
id=self.id,
title=self.title,
content=self.content[:100] if self.content else None,
)
class ChatSession(ObjectModel):
table_name: ClassVar[str] = "chat_session"
nullable_fields: ClassVar[set[str]] = {"model_override"}
title: Optional[str] = None
model_override: Optional[str] = None
async def relate_to_notebook(self, notebook_id: str) -> Any:
if not notebook_id:
raise InvalidInputError("Notebook ID must be provided")
return await self.relate("refers_to", notebook_id)
async def relate_to_source(self, source_id: str) -> Any:
if not source_id:
raise InvalidInputError("Source ID must be provided")
return await self.relate("refers_to", source_id)
async def text_search(
keyword: str, results: int, source: bool = True, note: bool = True
):
if not keyword:
raise InvalidInputError("Search keyword cannot be empty")
try:
search_results = await repo_query(
"""
select *
from fn::text_search($keyword, $results, $source, $note)
""",
{"keyword": keyword, "results": results, "source": source, "note": note},
)
return search_results
except RuntimeError as e:
# SurrealDB's search::highlight can compute a byte position that exceeds the
# stored string length on large or multi-byte chunks, aborting the whole query
# ("position overflow"). Fall back to vector search so the user still gets
# results instead of a 500. See issue #648.
if "position overflow" in str(e):
logger.warning(
f"Highlight position overflow, falling back to vector search: {str(e)}"
)
try:
return await vector_search(keyword, results, source, note)
except Exception as ve:
# Both search paths failed (e.g. no embedding model configured).
# Surface the failure instead of returning [] — an empty list would
# be indistinguishable from a legitimate "no matches" and mask a
# total search outage from callers.
logger.error(f"Vector search fallback also failed: {str(ve)}")
logger.exception(ve)
raise DatabaseOperationError(ve)
logger.error(f"Error performing text search: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
except Exception as e:
logger.error(f"Error performing text search: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
async def vector_search(
keyword: str,
results: int,
source: bool = True,
note: bool = True,
minimum_score=0.2,
):
if not keyword:
raise InvalidInputError("Search keyword cannot be empty")
try:
from open_notebook.utils.embedding import generate_embedding
# Use unified embedding function (handles chunking if query is very long)
embed = await generate_embedding(keyword)
search_results = await repo_query(
"""
SELECT * FROM fn::vector_search($embed, $results, $source, $note, $minimum_score);
""",
{
"embed": embed,
"results": results,
"source": source,
"note": note,
"minimum_score": minimum_score,
},
)
return search_results
except Exception as e:
logger.error(f"Error performing vector search: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
+444
View File
@@ -0,0 +1,444 @@
"""
Provider Configuration domain model for storing multiple credentials per provider.
This module provides the ProviderConfig singleton model that stores multiple
API key configurations per provider. Each ProviderCredential contains a complete
set of configuration options for a provider (api_key, base_url, model, etc.).
Encryption is enabled when OPEN_NOTEBOOK_ENCRYPTION_KEY environment variable
is set. If not set, keys are stored as plain text with a warning logged.
"""
from datetime import datetime
from typing import Any, ClassVar, Dict, List, Optional
from pydantic import Field, SecretStr
from open_notebook.database.repository import ensure_record_id, repo_query, repo_upsert
from open_notebook.domain.base import RecordModel
from open_notebook.utils.encryption import decrypt_value, encrypt_value
class ProviderCredential:
"""
A single provider configuration item containing api_key and related settings.
This class represents one complete configuration for an AI provider.
Multiple configurations can exist for the same provider, allowing users
to have different credentials for different environments (dev, prod, etc.).
Attributes:
id: Unique identifier for this configuration
name: Human-readable name for this configuration
provider: Provider name (e.g., "openai", "anthropic")
is_default: Whether this is the default configuration for the provider
api_key: The API key (stored as SecretStr for in-memory protection)
base_url: Base URL for the provider API
model: Default model to use for this provider
api_version: API version string (for providers that need it)
endpoint: Generic endpoint URL
endpoint_llm: Endpoint URL for LLM service
endpoint_embedding: Endpoint URL for embedding service
endpoint_stt: Endpoint URL for speech-to-text service
endpoint_tts: Endpoint URL for text-to-speech service
project: Project ID (for Vertex AI)
location: Location/region (for Vertex AI)
credentials_path: Path to credentials file (for Vertex AI)
created: Timestamp when this config was created
updated: Timestamp when this config was last updated
"""
def __init__(
self,
id: str,
name: str,
provider: str,
is_default: bool = False,
api_key: Optional[SecretStr] = None,
base_url: Optional[str] = None,
model: Optional[str] = None,
api_version: Optional[str] = None,
endpoint: Optional[str] = None,
endpoint_llm: Optional[str] = None,
endpoint_embedding: Optional[str] = None,
endpoint_stt: Optional[str] = None,
endpoint_tts: Optional[str] = None,
project: Optional[str] = None,
location: Optional[str] = None,
credentials_path: Optional[str] = None,
created: Optional[str] = None,
updated: Optional[str] = None,
):
self.id = id
self.name = name
self.provider = provider
self.is_default = is_default
self.api_key = api_key
self.base_url = base_url
self.model = model
self.api_version = api_version
self.endpoint = endpoint
self.endpoint_llm = endpoint_llm
self.endpoint_embedding = endpoint_embedding
self.endpoint_stt = endpoint_stt
self.endpoint_tts = endpoint_tts
self.project = project
self.location = location
self.credentials_path = credentials_path
self.created = created or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.updated = updated or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
def to_dict(self, encrypted: bool = False) -> dict:
"""
Convert the credential to a dictionary for storage.
Args:
encrypted: If True, api_key is encrypted; otherwise it's a SecretStr
Returns:
Dictionary representation of the credential
"""
data = {
"id": self.id,
"name": self.name,
"provider": self.provider,
"is_default": self.is_default,
"base_url": self.base_url,
"model": self.model,
"api_version": self.api_version,
"endpoint": self.endpoint,
"endpoint_llm": self.endpoint_llm,
"endpoint_embedding": self.endpoint_embedding,
"endpoint_stt": self.endpoint_stt,
"endpoint_tts": self.endpoint_tts,
"project": self.project,
"location": self.location,
"credentials_path": self.credentials_path,
"created": self.created,
"updated": self.updated,
}
if self.api_key:
if encrypted:
data["api_key"] = encrypt_value(self.api_key.get_secret_value())
else:
data["api_key"] = self.api_key.get_secret_value()
return data
@classmethod
def from_dict(cls, data: dict, decrypted: bool = False) -> "ProviderCredential":
"""
Create a ProviderCredential from a dictionary.
Args:
data: Dictionary containing credential data
decrypted: If True, api_key is already decrypted; otherwise wrap in SecretStr
Returns:
ProviderCredential instance
"""
api_key = None
if "api_key" in data and data["api_key"]:
if isinstance(data["api_key"], SecretStr):
# Already a SecretStr - use as-is
api_key = data["api_key"]
elif decrypted:
# Decrypted string from DB - wrap in SecretStr
api_key = SecretStr(data["api_key"])
else:
# Encrypted string from DB - wrap in SecretStr (will be decrypted later)
api_key = SecretStr(data["api_key"])
return cls(
id=data["id"],
name=data["name"],
provider=data["provider"],
is_default=data.get("is_default", False),
api_key=api_key,
base_url=data.get("base_url"),
model=data.get("model"),
api_version=data.get("api_version"),
endpoint=data.get("endpoint"),
endpoint_llm=data.get("endpoint_llm"),
endpoint_embedding=data.get("endpoint_embedding"),
endpoint_stt=data.get("endpoint_stt"),
endpoint_tts=data.get("endpoint_tts"),
project=data.get("project"),
location=data.get("location"),
credentials_path=data.get("credentials_path"),
created=data.get("created"),
updated=data.get("updated"),
)
class ProviderConfig(RecordModel):
"""
Singleton configuration for multiple provider credentials.
Uses RecordModel pattern with a fixed record_id. Stores a dictionary
of ProviderCredential objects organized by provider name.
Usage:
config = await ProviderConfig.get_instance()
credentials = config.credentials.get("openai", [])
default = config.get_default_config("openai")
"""
record_id: ClassVar[str] = "open_notebook:provider_configs"
# Store credentials organized by provider name
# Structure: {"openai": [ProviderCredential, ...], "anthropic": [...], ...}
credentials: Dict[str, List[ProviderCredential]] = Field(
default_factory=dict,
description="Provider credentials organized by provider name",
)
@classmethod
async def get_instance(cls) -> "ProviderConfig":
"""
Always fetch fresh configuration from database.
Overrides parent caching behavior to ensure we always get the latest
configuration values.
Returns:
ProviderConfig: Fresh instance with current database values
"""
result = await repo_query(
"SELECT * FROM ONLY $record_id",
{"record_id": ensure_record_id(cls.record_id)},
)
if result:
if isinstance(result, list) and len(result) > 0:
data = result[0]
elif isinstance(result, dict):
data = result
else:
data = {}
else:
data = {}
# Initialize credentials from database data
credentials: Dict[str, List[ProviderCredential]] = {}
creds_data = data.get("credentials")
if creds_data and isinstance(creds_data, dict):
for provider, provider_creds in creds_data.items():
if isinstance(provider_creds, list):
credentials[provider] = []
for cred_data in provider_creds:
try:
# Decrypt api_key if it's a string
api_key_val = cred_data.get("api_key")
if api_key_val and isinstance(api_key_val, str):
decrypted = decrypt_value(api_key_val)
cred_data["api_key"] = SecretStr(decrypted)
else:
# Keep as SecretStr or None
if api_key_val:
cred_data["api_key"] = SecretStr(api_key_val)
else:
cred_data["api_key"] = None
credentials[provider].append(
ProviderCredential(
id=cred_data.get("id", ""),
name=cred_data.get("name", "Default"),
provider=cred_data.get("provider", provider),
is_default=cred_data.get("is_default", False),
api_key=cred_data.get("api_key"),
base_url=cred_data.get("base_url"),
model=cred_data.get("model"),
api_version=cred_data.get("api_version"),
endpoint=cred_data.get("endpoint"),
endpoint_llm=cred_data.get("endpoint_llm"),
endpoint_embedding=cred_data.get(
"endpoint_embedding"
),
endpoint_stt=cred_data.get("endpoint_stt"),
endpoint_tts=cred_data.get("endpoint_tts"),
project=cred_data.get("project"),
location=cred_data.get("location"),
credentials_path=cred_data.get("credentials_path"),
created=cred_data.get("created"),
updated=cred_data.get("updated"),
)
)
except Exception:
# Skip invalid credentials
continue
# Create instance using model_validate to properly initialize Pydantic model
instance = cls.model_validate({"credentials": credentials})
# Mark as loaded from database
object.__setattr__(instance, "_db_loaded", True)
return instance
def get_default_config(self, provider: str) -> Optional[ProviderCredential]:
"""
Get the default configuration for a provider.
Args:
provider: Provider name (e.g., "openai", "anthropic")
Returns:
The default ProviderCredential, or None if not found
"""
provider_lower = provider.lower()
credentials = self.credentials.get(provider_lower, [])
# First, try to find explicitly marked default
for cred in credentials:
if cred.is_default:
return cred
# If no explicit default, return first config
if credentials:
return credentials[0]
return None
def get_config(
self, provider: str, config_id: str
) -> Optional[ProviderCredential]:
"""
Get a specific configuration by ID.
Args:
provider: Provider name
config_id: Configuration ID
Returns:
The ProviderCredential if found, None otherwise
"""
provider_lower = provider.lower()
credentials = self.credentials.get(provider_lower, [])
for cred in credentials:
if cred.id == config_id:
return cred
return None
def add_config(self, provider: str, credential: ProviderCredential) -> None:
"""
Add a new configuration for a provider.
If this is the first config for the provider, it becomes the default.
When adding a new config to an existing provider, the new config becomes
the default and previous default is unset.
Args:
provider: Provider name (normalized to lowercase)
credential: ProviderCredential to add
"""
provider_lower = provider.lower()
credential.provider = provider_lower
if provider_lower not in self.credentials:
self.credentials[provider_lower] = []
# When adding a new config to an existing provider, make it the default
# and unset the previous default
if self.credentials[provider_lower]:
for cred in self.credentials[provider_lower]:
cred.is_default = False
credential.is_default = True
# If this is the first config, make it default
if not self.credentials[provider_lower]:
credential.is_default = True
self.credentials[provider_lower].append(credential)
def delete_config(self, provider: str, config_id: str) -> bool:
"""
Delete a configuration.
Cannot delete the default configuration unless it's the only one.
Args:
provider: Provider name
config_id: Configuration ID to delete
Returns:
True if deleted, False if not found
"""
provider_lower = provider.lower()
credentials = self.credentials.get(provider_lower, [])
for i, cred in enumerate(credentials):
if cred.id == config_id:
# Cannot delete default if there are other configs
if cred.is_default and len(credentials) > 1:
return False
del credentials[i]
return True
return False
def set_default_config(self, provider: str, config_id: str) -> bool:
"""
Set a configuration as the default for a provider.
Args:
provider: Provider name
config_id: Configuration ID to make default
Returns:
True if successful, False if config not found
"""
provider_lower = provider.lower()
credentials = self.credentials.get(provider_lower, [])
for cred in credentials:
if cred.id == config_id:
# Unset all other defaults
for other in credentials:
other.is_default = False
# Set this one as default
cred.is_default = True
cred.updated = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
return True
return False
def _prepare_save_data(self) -> dict:
"""
Prepare data for database storage.
SecretStr values are extracted, encrypted, and stored as strings.
Encryption is performed using Fernet symmetric encryption if
OPEN_NOTEBOOK_ENCRYPTION_KEY is configured.
"""
data: Dict[str, Any] = {"credentials": {}}
for provider, credentials in self.credentials.items():
data["credentials"][provider] = []
for cred in credentials:
cred.updated = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
data["credentials"][provider].append(cred.to_dict(encrypted=True))
return data
async def save(self) -> "ProviderConfig":
"""
Save the configuration to the database.
Uses _prepare_save_data() to properly handle SecretStr conversion
and encryption.
"""
data = self._prepare_save_data()
await repo_upsert("open_notebook", self.record_id, data)
return self
@classmethod
def _clear_for_test(cls) -> None:
"""Clear the singleton instance for testing purposes."""
if cls.record_id in cls._instances:
del cls._instances[cls.record_id]
+30
View File
@@ -0,0 +1,30 @@
from typing import Any, ClassVar, Dict, Optional
from pydantic import Field
from open_notebook.database.repository import ensure_record_id
from open_notebook.domain.base import ObjectModel, RecordModel
class Transformation(ObjectModel):
table_name: ClassVar[str] = "transformation"
nullable_fields: ClassVar[set[str]] = {"model_id"}
name: str
title: str
description: str
prompt: str
apply_default: bool
model_id: Optional[str] = None
def _prepare_save_data(self) -> Dict[str, Any]:
data = super()._prepare_save_data()
if data.get("model_id"):
data["model_id"] = ensure_record_id(data["model_id"])
return data
class DefaultPrompts(RecordModel):
record_id: ClassVar[str] = "open_notebook:default_prompts"
transformation_instructions: Optional[str] = Field(
None, description="Instructions for executing a transformation"
)
+70
View File
@@ -0,0 +1,70 @@
class OpenNotebookError(Exception):
"""Base exception class for Open Notebook errors."""
pass
class DatabaseOperationError(OpenNotebookError):
"""Raised when a database operation fails."""
pass
class UnsupportedTypeException(OpenNotebookError):
"""Raised when an unsupported type is provided."""
pass
class InvalidInputError(OpenNotebookError):
"""Raised when invalid input is provided."""
pass
class NotFoundError(OpenNotebookError):
"""Raised when a requested resource is not found."""
pass
class AuthenticationError(OpenNotebookError):
"""Raised when there's an authentication problem."""
pass
class ConfigurationError(OpenNotebookError):
"""Raised when there's a configuration problem."""
pass
class ExternalServiceError(OpenNotebookError):
"""Raised when an external service (e.g., AI model) fails."""
pass
class RateLimitError(OpenNotebookError):
"""Raised when a rate limit is exceeded."""
pass
class FileOperationError(OpenNotebookError):
"""Raised when a file operation fails."""
pass
class NetworkError(OpenNotebookError):
"""Raised when a network operation fails."""
pass
class NoTranscriptFound(OpenNotebookError):
"""Raised when no transcript is found for a video."""
pass
+157
View File
@@ -0,0 +1,157 @@
import operator
from typing import Annotated, List
from ai_prompter import Prompter
from langchain_core.output_parsers.pydantic import PydanticOutputParser
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from langgraph.types import Send
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.domain.notebook import vector_search
from open_notebook.exceptions import OpenNotebookError
from open_notebook.utils import clean_thinking_content
from open_notebook.utils.error_classifier import classify_error
from open_notebook.utils.text_utils import extract_text_content
class SubGraphState(TypedDict):
question: str
term: str
instructions: str
results: dict
answer: str
ids: list # Added for provide_answer function
class Search(BaseModel):
term: str
instructions: str = Field(
description="Tell the answeting LLM what information you need extracted from this search"
)
class Strategy(BaseModel):
reasoning: str
searches: List[Search] = Field(
default_factory=list,
description="You can add up to five searches to this strategy",
)
class ThreadState(TypedDict):
question: str
strategy: Strategy
answers: Annotated[list, operator.add]
final_answer: str
async def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict:
try:
parser: PydanticOutputParser[Strategy] = PydanticOutputParser(
pydantic_object=Strategy
)
system_prompt = Prompter(prompt_template="ask/entry", parser=parser).render( # type: ignore[arg-type]
data=state # type: ignore[arg-type]
)
model = await provision_langchain_model(
system_prompt,
config.get("configurable", {}).get("strategy_model"),
"tools",
max_tokens=2000,
structured=dict(type="json"),
)
# model = model.bind_tools(tools)
# First get the raw response from the model
ai_message = await model.ainvoke(system_prompt)
# Clean the thinking content from the response
message_content = extract_text_content(ai_message.content)
cleaned_content = clean_thinking_content(message_content)
# Parse the cleaned JSON content
strategy = parser.parse(cleaned_content)
return {"strategy": strategy}
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
async def trigger_queries(state: ThreadState, config: RunnableConfig):
return [
Send(
"provide_answer",
{
"question": state["question"],
"instructions": s.instructions,
"term": s.term,
# "type": s.type,
},
)
for s in state["strategy"].searches
]
async def provide_answer(state: SubGraphState, config: RunnableConfig) -> dict:
try:
payload = state
# if state["type"] == "text":
# results = text_search(state["term"], 10, True, True)
# else:
results = await vector_search(state["term"], 10, True, True)
if len(results) == 0:
return {"answers": []}
payload["results"] = results
ids = [r["id"] for r in results]
payload["ids"] = ids
system_prompt = Prompter(prompt_template="ask/query_process").render(data=payload) # type: ignore[arg-type]
model = await provision_langchain_model(
system_prompt,
config.get("configurable", {}).get("answer_model"),
"tools",
max_tokens=2000,
)
ai_message = await model.ainvoke(system_prompt)
ai_content = extract_text_content(ai_message.content)
return {"answers": [clean_thinking_content(ai_content)]}
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
async def write_final_answer(state: ThreadState, config: RunnableConfig) -> dict:
try:
system_prompt = Prompter(prompt_template="ask/final_answer").render(data=state) # type: ignore[arg-type]
model = await provision_langchain_model(
system_prompt,
config.get("configurable", {}).get("final_answer_model"),
"tools",
max_tokens=2000,
)
ai_message = await model.ainvoke(system_prompt)
final_content = extract_text_content(ai_message.content)
return {"final_answer": clean_thinking_content(final_content)}
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
agent_state = StateGraph(ThreadState)
agent_state.add_node("agent", call_model_with_messages)
agent_state.add_node("provide_answer", provide_answer)
agent_state.add_node("write_final_answer", write_final_answer)
agent_state.add_edge(START, "agent")
agent_state.add_conditional_edges("agent", trigger_queries, ["provide_answer"])
agent_state.add_edge("provide_answer", "write_final_answer")
agent_state.add_edge("write_final_answer", END)
graph = agent_state.compile()
+98
View File
@@ -0,0 +1,98 @@
import asyncio
import sqlite3
from typing import Annotated, Optional
from ai_prompter import Prompter
from langchain_core.messages import SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE
from open_notebook.domain.notebook import Notebook
from open_notebook.exceptions import OpenNotebookError
from open_notebook.utils import clean_thinking_content
from open_notebook.utils.error_classifier import classify_error
from open_notebook.utils.text_utils import extract_text_content
class ThreadState(TypedDict):
messages: Annotated[list, add_messages]
notebook: Optional[Notebook]
context: Optional[str]
context_config: Optional[dict]
model_override: Optional[str]
def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict:
try:
system_prompt = Prompter(prompt_template="chat/system").render(data=state) # type: ignore[arg-type]
payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
model_id = config.get("configurable", {}).get("model_id") or state.get(
"model_override"
)
# Handle async model provisioning from sync context
def run_in_new_loop():
"""Run the async function in a new event loop"""
new_loop = asyncio.new_event_loop()
try:
asyncio.set_event_loop(new_loop)
return new_loop.run_until_complete(
provision_langchain_model(
str(payload), model_id, "chat", max_tokens=8192
)
)
finally:
new_loop.close()
asyncio.set_event_loop(None)
try:
# Try to get the current event loop
asyncio.get_running_loop()
# If we're in an event loop, run in a thread with a new loop
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_new_loop)
model = future.result()
except RuntimeError:
# No event loop running, safe to use asyncio.run()
model = asyncio.run(
provision_langchain_model(
str(payload),
model_id,
"chat",
max_tokens=8192,
)
)
ai_message = model.invoke(payload)
# Clean thinking content from AI response (e.g., <think>...</think> tags)
content = extract_text_content(ai_message.content)
cleaned_content = clean_thinking_content(content)
cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
return {"messages": cleaned_message}
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
conn = sqlite3.connect(
LANGGRAPH_CHECKPOINT_FILE,
check_same_thread=False,
)
memory = SqliteSaver(conn)
agent_state = StateGraph(ThreadState)
agent_state.add_node("agent", call_model_with_messages)
agent_state.add_edge(START, "agent")
agent_state.add_edge("agent", END)
graph = agent_state.compile(checkpointer=memory)
+49
View File
@@ -0,0 +1,49 @@
from typing import Any, Optional
from ai_prompter import Prompter
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.utils.text_utils import clean_thinking_content, extract_text_content
class PatternChainState(TypedDict):
prompt: str
parser: Optional[Any]
input_text: str
output: str
async def call_model(state: dict, config: RunnableConfig) -> dict:
content = state["input_text"]
# state["prompt"] is caller-supplied free text. Never compile it as Jinja
# template *source* (Prompter(template_text=...)) - pass it as a plain
# render variable into a fixed, developer-authored template instead.
# See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7).
system_prompt = Prompter(
prompt_template="pattern/generic", parser=state.get("parser")
).render(data=state)
payload = [SystemMessage(content=system_prompt)] + [HumanMessage(content=content)]
chain = await provision_langchain_model(
str(payload),
config.get("configurable", {}).get("model_id"),
"transformation",
max_tokens=5000,
)
response = await chain.ainvoke(payload)
# Clean thinking tags from response (handles extended thinking models)
output = clean_thinking_content(extract_text_content(response.content))
return {"output": output}
agent_state = StateGraph(PatternChainState)
agent_state.add_node("agent", call_model) # type: ignore[type-var]
agent_state.add_edge(START, "agent")
agent_state.add_edge("agent", END)
graph = agent_state.compile()
+246
View File
@@ -0,0 +1,246 @@
import operator
import os
from typing import Any, Dict, List, Optional
from content_core import ContentCoreConfig, extract_content
from content_core.common import ExtractionOutput
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from langgraph.types import Send
from loguru import logger
from typing_extensions import Annotated, TypedDict
from open_notebook.ai.models import Model, ModelManager
from open_notebook.domain.content_settings import ContentSettings
from open_notebook.domain.notebook import Asset, Source
from open_notebook.domain.transformation import Transformation
from open_notebook.graphs.transformation import graph as transform_graph
# Preferred languages for YouTube transcript selection. content-core's own
# default is only ["en", "es", "pt"]; we keep the broader list Open Notebook has
# always intended so non-English videos still resolve a transcript.
YOUTUBE_PREFERRED_LANGUAGES = [
"en",
"pt",
"es",
"de",
"nl",
"en-GB",
"fr",
"hi",
"ja",
]
class SourceState(TypedDict):
# Input describing what to extract: url / file_path / content / delete_source.
content_state: Dict[str, Any]
# Result of content-core extraction (does NOT echo url/file_path back).
extraction: ExtractionOutput
apply_transformations: List[Transformation]
source_id: str
notebook_ids: List[str]
source: Source
transformation: Annotated[list, operator.add]
embed: bool
class TransformationState(TypedDict):
source: Source
transformation: Transformation
async def content_process(state: SourceState) -> dict:
content_state: Dict[str, Any] = state["content_state"]
# content-core 2.x takes engine/model overrides via ContentCoreConfig
# (keyword-only), not inside the input dict.
config_kwargs: Dict[str, Any] = {
"youtube_languages": YOUTUBE_PREFERRED_LANGUAGES,
}
# Honor the persisted content-processing engine choices. content-core
# accepts "auto"/"simple"/"firecrawl"/"jina"/"crawl4ai" for URLs and
# "auto"/"docling"/"simple" for documents; falling back to "auto" keeps the
# previous behavior when settings are unset.
try:
settings: ContentSettings = await ContentSettings.get_instance() # type: ignore[assignment]
if settings.default_content_processing_engine_url:
config_kwargs["url_engine"] = settings.default_content_processing_engine_url
if settings.default_content_processing_engine_doc:
config_kwargs["document_engine"] = (
settings.default_content_processing_engine_doc
)
if settings.docling_ocr is not None:
config_kwargs["docling_ocr"] = settings.docling_ocr
except Exception as e:
# Keep the server-side traceback for diagnosing DB/deserialization
# failures while still falling back to defaults (non-fatal).
logger.opt(exception=True).warning(
f"Failed to load content settings, using defaults: {e}"
)
try:
model_manager = ModelManager()
defaults = await model_manager.get_defaults()
if defaults.default_speech_to_text_model:
stt_model = await Model.get(defaults.default_speech_to_text_model)
if stt_model:
config_kwargs["audio_provider"] = stt_model.provider
config_kwargs["audio_model"] = stt_model.name
logger.debug(
f"Using speech-to-text model: {stt_model.provider}/{stt_model.name}"
)
except Exception as e:
logger.warning(f"Failed to retrieve speech-to-text model configuration: {e}")
# Continue without custom audio model (content-core will use its default)
config = ContentCoreConfig(**config_kwargs) if config_kwargs else None
processed = await extract_content(
url=content_state.get("url"),
file_path=content_state.get("file_path"),
content=content_state.get("content"),
config=config,
)
# content-core signals a soft extraction failure (e.g. an unreachable or
# invalid URL, via the bs4 fallback) by returning title="Error" and content
# prefixed with "Failed to extract content:" instead of raising. Detect that
# sentinel and raise so the job is marked failed and the source becomes
# retryable, rather than being saved as a "completed" source whose body is
# the error string.
if processed.title == "Error" and (processed.content or "").startswith(
"Failed to extract content:"
):
raise ValueError(
"Could not extract content from this source. "
"The URL or file may be unreachable, invalid, or in an unsupported format."
)
if not processed.content or not processed.content.strip():
url = content_state.get("url") or ""
if url and ("youtube.com" in url or "youtu.be" in url):
raise ValueError(
"Could not extract content from this YouTube video. "
"No transcript or subtitles are available. "
"Try configuring a Speech-to-Text model in Settings "
"to transcribe the audio instead."
)
raise ValueError(
"Could not extract any text content from this source. "
"The content may be empty, inaccessible, or in an unsupported format."
)
# content-core 2.x no longer deletes the uploaded source file after
# extraction (the delete_source flag it used to honor is gone). Preserve the
# previous auto-delete behavior on our side.
if content_state.get("delete_source") and content_state.get("file_path"):
file_path = content_state["file_path"]
try:
os.unlink(file_path)
except FileNotFoundError:
logger.warning(f"File not found while trying to delete: {file_path}")
except Exception as e:
logger.warning(f"Failed to delete source file {file_path}: {e}")
return {"extraction": processed}
async def save_source(state: SourceState) -> dict:
content_state = state["content_state"]
extraction = state["extraction"]
# Get existing source using the provided source_id
source = await Source.get(state["source_id"])
if not source:
raise ValueError(f"Source with ID {state['source_id']} not found")
# Update the source with processed content. content-core's ExtractionOutput
# does not echo url/file_path back, so carry them from the input state.
source.asset = Asset(
url=content_state.get("url"), file_path=content_state.get("file_path")
)
source.full_text = extraction.content
# Preserve user-set title; only overwrite placeholder or empty titles
if extraction.title and (not source.title or source.title == "Processing..."):
source.title = extraction.title
await source.save()
# NOTE: Notebook associations are created by the API immediately for UI responsiveness
# No need to create them here to avoid duplicate edges
if state["embed"]:
if source.full_text and source.full_text.strip():
logger.debug("Embedding content for vector search")
await source.vectorize()
else:
logger.warning(
f"Source {source.id} has no text content to embed, skipping vectorization"
)
return {"source": source}
def trigger_transformations(state: SourceState, config: RunnableConfig) -> List[Send]:
if len(state["apply_transformations"]) == 0:
return []
to_apply = state["apply_transformations"]
logger.debug(f"Applying transformations {to_apply}")
return [
Send(
"transform_content",
{
"source": state["source"],
"transformation": t,
},
)
for t in to_apply
]
async def transform_content(state: TransformationState) -> Optional[dict]:
source = state["source"]
content = source.full_text
if not content:
return None
transformation: Transformation = state["transformation"]
logger.debug(f"Applying transformation {transformation.name}")
# LangGraph accepts a partial state dict at runtime, but its typed
# overloads require the full state type (langgraph typing limitation).
result = await transform_graph.ainvoke( # type: ignore[call-overload]
dict(input_text=content, transformation=transformation)
)
await source.add_insight(transformation.title, result["output"])
return {
"transformation": [
{
"output": result["output"],
"transformation_name": transformation.name,
}
]
}
# Create and compile the workflow
workflow = StateGraph(SourceState)
# Add nodes
workflow.add_node("content_process", content_process)
workflow.add_node("save_source", save_source)
workflow.add_node("transform_content", transform_content)
# Define the graph edges
workflow.add_edge(START, "content_process")
workflow.add_edge("content_process", "save_source")
workflow.add_conditional_edges(
"save_source", trigger_transformations, ["transform_content"]
)
workflow.add_edge("transform_content", END)
# Compile the graph
source_graph = workflow.compile()
+254
View File
@@ -0,0 +1,254 @@
import asyncio
import sqlite3
from typing import Annotated, Dict, List, Optional
from ai_prompter import Prompter
from langchain_core.messages import SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE
from open_notebook.domain.notebook import Source, SourceInsight
from open_notebook.exceptions import OpenNotebookError
from open_notebook.utils import clean_thinking_content
from open_notebook.utils.context_builder import build_source_context
from open_notebook.utils.error_classifier import classify_error
from open_notebook.utils.text_utils import extract_text_content
class SourceChatState(TypedDict):
messages: Annotated[list, add_messages]
source_id: str
source: Optional[Source]
insights: Optional[List[SourceInsight]]
context: Optional[str]
model_override: Optional[str]
context_indicators: Optional[Dict[str, List[str]]]
def call_model_with_source_context(
state: SourceChatState, config: RunnableConfig
) -> dict:
"""
Main function that builds source context and calls the model.
This function:
1. Uses build_source_context to build source-specific context
2. Applies the source_chat Jinja2 prompt template
3. Handles model provisioning with override support
4. Tracks context indicators for referenced insights/content
"""
try:
return _call_model_with_source_context_inner(state, config)
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
def _call_model_with_source_context_inner(
state: SourceChatState, config: RunnableConfig
) -> dict:
source_id = state.get("source_id")
if not source_id:
raise ValueError("source_id is required in state")
# Build source context using build_source_context (run async code in new loop)
def build_context():
"""Build context in a new event loop"""
new_loop = asyncio.new_event_loop()
try:
asyncio.set_event_loop(new_loop)
return new_loop.run_until_complete(
build_source_context(
source_id=source_id,
max_tokens=50000, # Reasonable limit for source context
)
)
finally:
new_loop.close()
asyncio.set_event_loop(None)
# Get the built context
try:
# Try to get the current event loop
asyncio.get_running_loop()
# If we're in an event loop, run in a thread with a new loop
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(build_context)
context_data = future.result()
except RuntimeError:
# No event loop running, safe to create a new one
context_data = build_context()
# Extract source and insights from context
source = None
insights = []
context_indicators: dict[str, list[str | None]] = {
"sources": [],
"insights": [],
"notes": [],
}
if context_data.get("sources"):
source_info = context_data["sources"][0] # First source
source = Source(**source_info) if isinstance(source_info, dict) else source_info
context_indicators["sources"].append(source.id)
if context_data.get("insights"):
for insight_data in context_data["insights"]:
insight = (
SourceInsight(**insight_data)
if isinstance(insight_data, dict)
else insight_data
)
insights.append(insight)
context_indicators["insights"].append(insight.id)
# Format context for the prompt
formatted_context = _format_source_context(context_data)
# Build prompt data for the template
prompt_data = {
"source": source.model_dump() if source else None,
"insights": [insight.model_dump() for insight in insights] if insights else [],
"context": formatted_context,
"context_indicators": context_indicators,
}
# Apply the source_chat prompt template
system_prompt = Prompter(prompt_template="source_chat/system").render(
data=prompt_data
)
payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
# Handle async model provisioning from sync context
def run_in_new_loop():
"""Run the async function in a new event loop"""
new_loop = asyncio.new_event_loop()
try:
asyncio.set_event_loop(new_loop)
return new_loop.run_until_complete(
provision_langchain_model(
str(payload),
config.get("configurable", {}).get("model_id")
or state.get("model_override"),
"chat",
max_tokens=8192,
)
)
finally:
new_loop.close()
asyncio.set_event_loop(None)
try:
# Try to get the current event loop
asyncio.get_running_loop()
# If we're in an event loop, run in a thread with a new loop
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_new_loop)
model = future.result()
except RuntimeError:
# No event loop running, safe to use asyncio.run()
model = asyncio.run(
provision_langchain_model(
str(payload),
config.get("configurable", {}).get("model_id")
or state.get("model_override"),
"chat",
max_tokens=8192,
)
)
ai_message = model.invoke(payload)
# Clean thinking content from AI response (e.g., <think>...</think> tags)
content = extract_text_content(ai_message.content)
cleaned_content = clean_thinking_content(content)
cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
# Update state with context information
return {
"messages": cleaned_message,
"source": source,
"insights": insights,
"context": formatted_context,
"context_indicators": context_indicators,
}
def _format_source_context(context_data: Dict) -> str:
"""
Format the context data into a readable string for the prompt.
Args:
context_data: Context data from build_source_context
Returns:
Formatted context string
"""
context_parts = []
# Add source information
if context_data.get("sources"):
context_parts.append("## SOURCE CONTENT")
for source in context_data["sources"]:
if isinstance(source, dict):
context_parts.append(f"**Source ID:** {source.get('id', 'Unknown')}")
context_parts.append(f"**Title:** {source.get('title', 'No title')}")
if source.get("full_text"):
# Truncate full text if too long
full_text = source["full_text"]
if len(full_text) > 5000:
full_text = full_text[:5000] + "...\n[Content truncated]"
context_parts.append(f"**Content:**\n{full_text}")
context_parts.append("") # Empty line for separation
# Add insights
if context_data.get("insights"):
context_parts.append("## SOURCE INSIGHTS")
for insight in context_data["insights"]:
if isinstance(insight, dict):
context_parts.append(f"**Insight ID:** {insight.get('id', 'Unknown')}")
context_parts.append(
f"**Type:** {insight.get('insight_type', 'Unknown')}"
)
context_parts.append(
f"**Content:** {insight.get('content', 'No content')}"
)
context_parts.append("") # Empty line for separation
# Add metadata
if context_data.get("metadata"):
metadata = context_data["metadata"]
context_parts.append("## CONTEXT METADATA")
context_parts.append(f"- Source count: {metadata.get('source_count', 0)}")
context_parts.append(f"- Insight count: {metadata.get('insight_count', 0)}")
context_parts.append(f"- Total tokens: {context_data.get('total_tokens', 0)}")
context_parts.append("")
return "\n".join(context_parts)
# Create SQLite checkpointer
conn = sqlite3.connect(
LANGGRAPH_CHECKPOINT_FILE,
check_same_thread=False,
)
memory = SqliteSaver(conn)
# Create the StateGraph
source_chat_state = StateGraph(SourceChatState)
source_chat_state.add_node("source_chat_agent", call_model_with_source_context)
source_chat_state.add_edge(START, "source_chat_agent")
source_chat_state.add_edge("source_chat_agent", END)
source_chat_graph = source_chat_state.compile(checkpointer=memory)
+13
View File
@@ -0,0 +1,13 @@
from datetime import datetime
from langchain.tools import tool
# todo: turn this into a system prompt variable
@tool
def get_current_timestamp() -> str:
"""
name: get_current_timestamp
Returns the current timestamp in the format YYYYMMDDHHmmss.
"""
return datetime.now().strftime("%Y%m%d%H%M%S")
+77
View File
@@ -0,0 +1,77 @@
from ai_prompter import Prompter
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict
from open_notebook.ai.provision import provision_langchain_model
from open_notebook.domain.notebook import Source
from open_notebook.domain.transformation import DefaultPrompts, Transformation
from open_notebook.exceptions import OpenNotebookError
from open_notebook.utils import clean_thinking_content
from open_notebook.utils.error_classifier import classify_error
from open_notebook.utils.text_utils import extract_text_content
class TransformationState(TypedDict):
input_text: str
source: Source
transformation: Transformation
output: str
async def run_transformation(state: dict, config: RunnableConfig) -> dict:
source_obj = state.get("source")
source: Source = source_obj if isinstance(source_obj, Source) else None # type: ignore[assignment]
content = state.get("input_text")
assert source or content, "No content to transform"
transformation: Transformation = state["transformation"]
try:
if not content:
content = source.full_text
# transformation.prompt is user-controlled free text. Never compile it as
# Jinja template *source* (Prompter(template_text=...)) - pass it as a
# plain render variable into a fixed, developer-authored template instead.
# See docs/7-DEVELOPMENT/security.md (GHSA-f35w-wx37-26q7).
instructions = transformation.prompt
default_prompts: DefaultPrompts = DefaultPrompts(transformation_instructions=None)
if default_prompts.transformation_instructions:
instructions = f"{default_prompts.transformation_instructions}\n\n{instructions}"
system_prompt = Prompter(prompt_template="transformation/execute").render(
data={**state, "instructions": instructions}
)
content_str = str(content) if content else ""
payload = [SystemMessage(content=system_prompt), HumanMessage(content=content_str)]
chain = await provision_langchain_model(
str(payload),
config.get("configurable", {}).get("model_id"),
"transformation",
max_tokens=8192,
)
response = await chain.ainvoke(payload)
# Clean thinking content from the response
response_content = extract_text_content(response.content)
cleaned_content = clean_thinking_content(response_content)
if source:
await source.add_insight(transformation.title, cleaned_content)
return {
"output": cleaned_content,
}
except OpenNotebookError:
raise
except Exception as e:
error_class, user_message = classify_error(e)
raise error_class(user_message) from e
agent_state = StateGraph(TransformationState)
agent_state.add_node("agent", run_transformation) # type: ignore[type-var]
agent_state.add_edge(START, "agent")
agent_state.add_edge("agent", END)
graph = agent_state.compile()
+2
View File
@@ -0,0 +1,2 @@
# Podcasts module
# Contains podcast episode models, profiles, and generation logic
+87
View File
@@ -0,0 +1,87 @@
"""Single choke point for podcast episode audio file paths (#1030).
``PodcastEpisode.audio_file`` stores a path RELATIVE to ``PODCASTS_FOLDER``
(e.g. ``episodes/<uuid>/audio/<uuid>.mp3``). Two helpers enforce that
contract at the only two places paths cross the DB boundary:
- ``to_relative_audio_path()`` — write side (podcast generation command):
converts the generated file path to the relative storage form and refuses
to produce a value outside the podcasts root, so the DB never holds an
absolute or escaping path.
- ``resolve_contained_audio_path()`` — read side (every API consumption
point: stream, list/get, delete, retry): joins the stored value with
``PODCASTS_FOLDER``, resolves symlinks/``..`` and verifies containment.
Any absolute, ``file://`` or escaping value is treated as legacy-invalid
and returns ``None`` (callers keep today's 403/404 behavior from #1018).
Storing relative paths makes path traversal unrepresentable for new rows
and lets previously generated episodes survive a ``DATA_FOLDER`` move.
Migration 21 converts pre-existing rows written under the known roots.
"""
import os
from pathlib import Path
from typing import Optional, Union
from urllib.parse import unquote, urlparse
from open_notebook.config import PODCASTS_FOLDER
def podcasts_root() -> Path:
"""Real (symlink-resolved, absolute) path of the podcasts output root.
Computed on every call rather than at import time so tests can
monkeypatch ``PODCASTS_FOLDER`` on this module.
"""
return Path(os.path.realpath(PODCASTS_FOLDER))
def to_relative_audio_path(audio_path: Union[str, Path]) -> str:
"""Convert a generated audio file path to the DB storage form.
Accepts the absolute (or CWD-relative) path produced by podcast-creator,
including the legacy ``file://`` URI form, and returns it relative to
``PODCASTS_FOLDER`` as a POSIX-style string.
Raises:
ValueError: if the path resolves outside the podcasts root — the DB
must never hold an absolute or escaping value. ValueError also
marks the generation job as permanently failed (no retry).
"""
raw = str(audio_path)
if raw.startswith("file://"):
raw = unquote(urlparse(raw).path)
resolved = Path(os.path.realpath(raw))
root = podcasts_root()
if resolved == root or not resolved.is_relative_to(root):
raise ValueError(
f"Generated audio file path is outside the podcasts folder: {audio_path}"
)
return resolved.relative_to(root).as_posix()
def resolve_contained_audio_path(audio_file: Optional[str]) -> Optional[Path]:
"""Resolve a stored ``audio_file`` value to a real filesystem path.
Joins the stored relative path with ``PODCASTS_FOLDER``, resolves
symlinks and ``..`` components, and verifies the result stays inside the
podcasts root.
Returns ``None`` for anything that must not be followed:
- empty/None values
- absolute paths and ``file://`` URIs (legacy rows migration 21 could
not convert — exactly the out-of-root cases #1018's guards reject)
- relative paths that escape the root (``..`` or symlink traversal)
"""
if not audio_file:
return None
if "://" in audio_file:
return None
candidate = Path(audio_file)
if candidate.is_absolute():
return None
root = podcasts_root()
resolved = Path(os.path.realpath(root / candidate))
if resolved == root or not resolved.is_relative_to(root):
return None
return resolved
+339
View File
@@ -0,0 +1,339 @@
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
from loguru import logger
from pydantic import ConfigDict, Field, field_validator
from surrealdb import RecordID
from open_notebook.database.repository import ensure_record_id, repo_query
from open_notebook.domain.base import ObjectModel
async def _resolve_model_config(
model_id: str, max_tokens: Optional[int] = None
) -> Tuple[str, str, dict]:
"""Load Model record, resolve credential -> (provider, model_name, config_dict).
Used by resolve_outline_config, resolve_transcript_config, resolve_tts_config,
and per-speaker TTS overrides. Optionally passes through a max_tokens override.
"""
from open_notebook.ai.models import Model
model = await Model.get(model_id)
config: dict = {}
if model.credential:
credential = await model.get_credential_obj()
if credential:
config = credential.to_esperanto_config()
if not config:
from open_notebook.ai.key_provider import provision_provider_keys
await provision_provider_keys(model.provider)
if max_tokens is not None:
config = {**config, "max_tokens": max_tokens}
return (model.provider, model.name, config)
class EpisodeProfile(ObjectModel):
"""
Episode Profile - Simplified podcast configuration.
Replaces complex 15+ field configuration with user-friendly profiles.
"""
table_name: ClassVar[str] = "episode_profile"
nullable_fields: ClassVar[set[str]] = {
"description",
"speaker_config",
"outline_llm",
"transcript_llm",
"language",
"max_tokens",
}
name: str = Field(..., description="Unique profile name")
description: Optional[str] = Field(None, description="Profile description")
speaker_config: Optional[str] = Field(
None,
description=(
"speaker_profile record ID this profile uses. None when the "
"referenced speaker profile no longer exists (orphaned by "
"migration 20 or a later deletion)."
),
)
# Model registry references
outline_llm: Optional[str] = Field(
None, description="Model record ID for outline generation"
)
transcript_llm: Optional[str] = Field(
None, description="Model record ID for transcript generation"
)
language: Optional[str] = Field(
None, description="Podcast language (BCP 47 locale code, e.g. pt-BR, en-US)"
)
default_briefing: str = Field(..., description="Default briefing template")
num_segments: int = Field(default=5, description="Number of podcast segments")
max_tokens: Optional[int] = Field(
None,
description="Max output tokens for outline/transcript generation (passed through to podcast_creator)",
)
@field_validator("num_segments")
@classmethod
def validate_segments(cls, v):
if not 3 <= v <= 20:
raise ValueError("Number of segments must be between 3 and 20")
return v
def _prepare_save_data(self) -> dict:
data = super()._prepare_save_data()
if data.get("speaker_config"):
data["speaker_config"] = ensure_record_id(data["speaker_config"])
if data.get("outline_llm"):
data["outline_llm"] = ensure_record_id(data["outline_llm"])
if data.get("transcript_llm"):
data["transcript_llm"] = ensure_record_id(data["transcript_llm"])
return data
async def resolve_outline_config(self) -> Tuple[str, str, dict]:
"""Resolve outline model -> (provider, model_name, config_dict)"""
if not self.outline_llm:
raise ValueError(
f"Episode profile '{self.name}' has no outline model configured. "
"Please update the profile to select an outline model."
)
return await _resolve_model_config(self.outline_llm, max_tokens=self.max_tokens)
async def resolve_transcript_config(self) -> Tuple[str, str, dict]:
"""Resolve transcript model -> (provider, model_name, config_dict)"""
if not self.transcript_llm:
raise ValueError(
f"Episode profile '{self.name}' has no transcript model configured. "
"Please update the profile to select a transcript model."
)
return await _resolve_model_config(
self.transcript_llm, max_tokens=self.max_tokens
)
@classmethod
async def get_by_name(cls, name: str) -> Optional["EpisodeProfile"]:
"""Get episode profile by name"""
result = await repo_query(
"SELECT * FROM episode_profile WHERE name = $name", {"name": name}
)
if result:
return cls(**result[0])
return None
class SpeakerProfile(ObjectModel):
"""
Speaker Profile - Voice and personality configuration.
Supports 1-4 speakers for flexible podcast formats.
"""
table_name: ClassVar[str] = "speaker_profile"
nullable_fields: ClassVar[set[str]] = {
"description",
"voice_model",
}
name: str = Field(..., description="Unique profile name")
description: Optional[str] = Field(None, description="Profile description")
# Model registry reference
voice_model: Optional[str] = Field(
None, description="Model record ID for TTS"
)
speakers: List[Dict[str, Any]] = Field(
..., description="Array of speaker configurations"
)
@field_validator("speakers")
@classmethod
def validate_speakers(cls, v):
if not 1 <= len(v) <= 4:
raise ValueError("Must have between 1 and 4 speakers")
required_fields = ["name", "voice_id", "backstory", "personality"]
for speaker in v:
for field in required_fields:
if field not in speaker:
raise ValueError(f"Speaker missing required field: {field}")
return v
def _prepare_save_data(self) -> dict:
data = super()._prepare_save_data()
if data.get("voice_model"):
data["voice_model"] = ensure_record_id(data["voice_model"])
# Handle per-speaker voice_model overrides
if data.get("speakers"):
for speaker in data["speakers"]:
if speaker.get("voice_model"):
speaker["voice_model"] = ensure_record_id(speaker["voice_model"])
return data
async def resolve_tts_config(self) -> Tuple[str, str, dict]:
"""Resolve TTS model -> (provider, model_name, config_dict)"""
if not self.voice_model:
raise ValueError(
f"Speaker profile '{self.name}' has no voice model configured. "
"Please update the profile to select a voice model."
)
return await _resolve_model_config(self.voice_model)
@classmethod
async def get_by_name(cls, name: str) -> Optional["SpeakerProfile"]:
"""Get speaker profile by name"""
result = await repo_query(
"SELECT * FROM speaker_profile WHERE name = $name", {"name": name}
)
if result:
return cls(**result[0])
return None
@classmethod
async def resolve(
cls, ref: Union[str, RecordID]
) -> Optional["SpeakerProfile"]:
"""Resolve a speaker profile by record ID or by unique name.
The API contract accepts speaker profiles by NAME (see
POST /api/podcasts/generate), while episode_profile.speaker_config
stores a record ID (migration 20). This resolves either form and
returns None when the reference doesn't match anything.
"""
ref_str = str(ref)
if ref_str.startswith(f"{cls.table_name}:"):
result = await repo_query(
"SELECT * FROM $id", {"id": ensure_record_id(ref_str)}
)
if result:
return cls(**result[0])
return None
return await cls.get_by_name(ref_str)
class PodcastEpisode(ObjectModel):
"""Enhanced PodcastEpisode with job tracking and metadata"""
table_name: ClassVar[str] = "episode"
name: str = Field(..., description="Episode name")
episode_profile: Dict[str, Any] = Field(
..., description="Episode profile used (stored as object)"
)
speaker_profile: Dict[str, Any] = Field(
..., description="Speaker profile used (stored as object)"
)
briefing: str = Field(..., description="Full briefing used for generation")
content: str = Field(..., description="Source content")
audio_file: Optional[str] = Field(
default=None,
description=(
"Path to the generated audio file, relative to PODCASTS_FOLDER "
"(see open_notebook/podcasts/audio_paths.py). Absolute values "
"are legacy rows migration 21 could not convert and are treated "
"as invalid by the API."
),
)
transcript: Optional[Dict[str, Any]] = Field(
default_factory=dict, description="Generated transcript"
)
outline: Optional[Dict[str, Any]] = Field(
default_factory=dict, description="Generated outline"
)
command: Optional[Union[str, RecordID]] = Field(
default=None, description="Link to surreal-commands job"
)
model_config = ConfigDict(arbitrary_types_allowed=True)
async def get_job_status(self) -> Optional[str]:
"""Get the status of the associated command"""
if not self.command:
return None
try:
from surreal_commands import get_command_status
status = await get_command_status(str(self.command))
return status.status if status else "unknown"
except Exception:
return "unknown"
async def get_job_detail(self) -> dict:
"""Get status and error_message of the associated command"""
if not self.command:
return {"status": None, "error_message": None}
try:
from surreal_commands import get_command_status
status = await get_command_status(str(self.command))
if not status:
return {"status": "unknown", "error_message": None}
return {
"status": status.status,
"error_message": getattr(status, "error_message", None),
}
except Exception:
return {"status": "unknown", "error_message": None}
@classmethod
async def get_job_details_for_commands(
cls, command_ids: List[Union[str, RecordID]]
) -> Dict[str, dict]:
"""
Batch-fetch {status, error_message} for many commands in one query.
Listing episodes otherwise calls get_job_detail() -> surreal_commands
.get_command_status() once per episode, each its own round trip
against the `command` table (no connection pooling in the repository
layer, see docs/7-DEVELOPMENT/architecture.md) - O(n) queries for n
episodes. surreal_commands has no batch lookup, but its command table
lives in the same database (same SURREAL_* env vars), so this queries
it directly in one shot instead of looping through the library's
per-command helper.
CommandStatus is a `str` subclass (`class CommandStatus(str, Enum)`),
so returning the raw DB string here is interchangeable with the
enum-wrapped value get_job_detail() returns for every comparison
this codebase does against it.
"""
ids = [cid for cid in command_ids if cid]
grouped: Dict[str, dict] = {}
if not ids:
return grouped
try:
result = await repo_query(
"SELECT * FROM command WHERE id IN $command_ids",
{"command_ids": [ensure_record_id(cid) for cid in ids]},
)
except Exception as e:
logger.error(f"Error batch-fetching command status: {e}")
return grouped
for row in result:
grouped[str(row.get("id"))] = {
"status": row.get("status", "unknown"),
"error_message": row.get("error_message"),
}
return grouped
@field_validator("command", mode="before")
@classmethod
def parse_command(cls, value):
if isinstance(value, str):
return ensure_record_id(value)
return value
def _prepare_save_data(self) -> dict:
"""Override to ensure command field is always RecordID format for database"""
data = super()._prepare_save_data()
# Ensure command field is RecordID format if not None
if data.get("command") is not None:
data["command"] = ensure_record_id(data["command"])
return data
+188
View File
@@ -0,0 +1,188 @@
# ContextBuilder
A flexible and generic ContextBuilder class for the Open Notebook project that can handle any parameters and build context from sources, notebooks, insights, and notes.
## Features
- **Flexible Parameters**: Accepts any parameters via `**kwargs` for future extensibility
- **Priority-based Management**: Automatic prioritization and sorting of context items
- **Token Counting**: Built-in token counting and truncation to fit limits
- **Deduplication**: Automatic removal of duplicate items based on ID
- **Type-based Grouping**: Separates sources, notes, and insights in output
- **Async Support**: Fully async for database operations
## Basic Usage
```python
from open_notebook.utils.context_builder import ContextBuilder, ContextConfig
# Simple notebook context
builder = ContextBuilder(notebook_id="notebook:123")
context = await builder.build()
# Single source with insights
builder = ContextBuilder(
source_id="source:456",
include_insights=True,
max_tokens=2000
)
context = await builder.build()
```
## Convenience Functions
```python
from open_notebook.utils.context_builder import (
build_notebook_context,
build_source_context,
build_mixed_context
)
# Build notebook context
context = await build_notebook_context(
notebook_id="notebook:123",
max_tokens=5000
)
# Build single source context
context = await build_source_context(
source_id="source:456",
include_insights=True
)
# Build mixed context
context = await build_mixed_context(
source_ids=["source:1", "source:2"],
note_ids=["note:1", "note:2"],
max_tokens=3000
)
```
## Advanced Configuration
```python
from open_notebook.utils.context_builder import ContextConfig
# Custom configuration
config = ContextConfig(
sources={
"source:doc1": "insights",
"source:doc2": "full content",
"source:doc3": "not in" # Exclude
},
notes={
"note:summary": "full content",
"note:draft": "not in" # Exclude
},
include_insights=True,
max_tokens=3000,
priority_weights={
"source": 120, # Higher priority
"note": 80, # Medium priority
"insight": 100 # High priority
}
)
builder = ContextBuilder(
notebook_id="notebook:project",
context_config=config
)
context = await builder.build()
```
## Programmatic Item Management
```python
from open_notebook.utils.context_builder import ContextItem
builder = ContextBuilder()
# Add custom items
item = ContextItem(
id="source:important",
type="source",
content={"title": "Key Document", "summary": "..."},
priority=150 # Very high priority
)
builder.add_item(item)
# Apply management operations
builder.remove_duplicates()
builder.prioritize()
builder.truncate_to_fit(1000)
context = builder._format_response()
```
## Flexible Parameters
The ContextBuilder accepts any parameters via `**kwargs`, making it extensible for future features:
```python
builder = ContextBuilder(
notebook_id="notebook:123",
include_insights=True,
max_tokens=2000,
# Custom parameters for future extensions
user_id="user:456",
custom_filter="advanced",
experimental_feature=True
)
# Access custom parameters
user_id = builder.params.get('user_id')
```
## Output Format
The ContextBuilder returns a structured response:
```python
{
"sources": [...], # List of source contexts
"notes": [...], # List of note contexts
"insights": [...], # List of insight contexts
"total_tokens": 1234, # Total token count
"total_items": 10, # Total number of items
"notebook_id": "notebook:123", # If provided
"metadata": {
"source_count": 5,
"note_count": 3,
"insight_count": 2,
"config": {
"include_insights": true,
"include_notes": true,
"max_tokens": 2000
}
}
}
```
## Architecture
The ContextBuilder follows these design principles:
1. **Separation of Concerns**: Context building, item management, and formatting are separate
2. **Extensibility**: Uses `**kwargs` and flexible configuration for future features
3. **Performance**: Token-aware truncation and efficient deduplication
4. **Type Safety**: Proper type hints and data classes for structure
5. **Error Handling**: Graceful handling of missing items and database errors
## Integration
The ContextBuilder integrates seamlessly with the existing Open Notebook architecture:
- Uses existing domain models (`Source`, `Notebook`, `Note`)
- Leverages the repository pattern for database access
- Follows the same async patterns as other services
- Integrates with the token counting utilities
## Error Handling
The ContextBuilder handles errors gracefully:
- Missing notebooks/sources/notes are logged but don't stop execution
- Database errors are wrapped in `DatabaseOperationError`
- Invalid parameters raise `InvalidInputError`
- All errors include detailed context information
+72
View File
@@ -0,0 +1,72 @@
"""
Utils package for Open Notebook.
To avoid circular imports, import functions directly:
- from open_notebook.utils.context_builder import build_notebook_context, build_source_context
- from open_notebook.utils import token_count, compare_versions
- from open_notebook.utils.chunking import chunk_text, detect_content_type, ContentType
- from open_notebook.utils.embedding import generate_embedding, generate_embeddings
- from open_notebook.utils.encryption import encrypt_value, decrypt_value
"""
from .chunking import (
CHUNK_SIZE,
ContentType,
chunk_text,
detect_content_type,
detect_content_type_from_extension,
detect_content_type_from_heuristics,
)
from .embedding import (
generate_embedding,
generate_embeddings,
mean_pool_embeddings,
)
from .encryption import (
decrypt_value,
encrypt_value,
)
from .model_utils import full_model_dump
from .text_utils import (
clean_thinking_content,
parse_thinking_content,
remove_non_ascii,
remove_non_printable,
)
from .token_utils import token_cost, token_count
from .version_utils import (
compare_versions,
get_installed_version,
get_version_from_github,
)
__all__ = [
# Chunking
"CHUNK_SIZE",
"ContentType",
"chunk_text",
"detect_content_type",
"detect_content_type_from_extension",
"detect_content_type_from_heuristics",
# Embedding
"generate_embedding",
"generate_embeddings",
"mean_pool_embeddings",
# Text utils
"remove_non_ascii",
"remove_non_printable",
"parse_thinking_content",
"clean_thinking_content",
# Token utils
"token_count",
"token_cost",
# Version utils
"compare_versions",
"get_installed_version",
"get_version_from_github",
# Encryption utils
"decrypt_value",
"encrypt_value",
# Model utils
"full_model_dump",
]
+494
View File
@@ -0,0 +1,494 @@
"""
Chunking utilities for Open Notebook.
Provides content-type detection and smart text chunking for embedding operations.
Supports HTML, Markdown, and plain text with appropriate splitters for each type.
Key functions:
- detect_content_type(): Detects content type from file extension or content heuristics
- chunk_text(): Splits text into chunks using appropriate splitter for content type
Environment Variables:
OPEN_NOTEBOOK_CHUNK_SIZE: Maximum chunk size in tokens (default: 400)
OPEN_NOTEBOOK_CHUNK_OVERLAP: Overlap between chunks in tokens (default: 15% of CHUNK_SIZE)
OPEN_NOTEBOOK_MIN_CHUNK_SIZE: Minimum chunk size in tokens (default: 5)
"""
import os
import re
from enum import Enum
from pathlib import Path
from typing import List, Optional, Tuple
from langchain_text_splitters import (
HTMLHeaderTextSplitter,
MarkdownHeaderTextSplitter,
RecursiveCharacterTextSplitter,
)
from loguru import logger
from .token_utils import token_count
def _get_chunk_size() -> int:
"""Get chunk size from environment variable or use default."""
chunk_size_str = os.getenv("OPEN_NOTEBOOK_CHUNK_SIZE")
if chunk_size_str:
try:
chunk_size = int(chunk_size_str)
if chunk_size < 100:
logger.warning(
f"OPEN_NOTEBOOK_CHUNK_SIZE ({chunk_size}) is too small. "
f"Using minimum value of 100."
)
return 100
if chunk_size > 8192:
logger.warning(
f"OPEN_NOTEBOOK_CHUNK_SIZE ({chunk_size}) is very large. "
f"This may cause issues with some embedding models."
)
logger.info(f"Using custom chunk size: {chunk_size} tokens")
return chunk_size
except ValueError:
logger.warning(
f"Invalid OPEN_NOTEBOOK_CHUNK_SIZE value: '{chunk_size_str}'. "
f"Using default: 400"
)
return 400
def _get_chunk_overlap(chunk_size: int) -> int:
"""Get chunk overlap from environment variable or calculate default (15% of chunk size)."""
overlap_str = os.getenv("OPEN_NOTEBOOK_CHUNK_OVERLAP")
if overlap_str:
try:
overlap = int(overlap_str)
if overlap < 0:
logger.warning(
f"OPEN_NOTEBOOK_CHUNK_OVERLAP ({overlap}) cannot be negative. "
f"Using 0."
)
return 0
if overlap >= chunk_size:
logger.warning(
f"OPEN_NOTEBOOK_CHUNK_OVERLAP ({overlap}) cannot be >= chunk size ({chunk_size}). "
f"Using 15% of chunk size: {int(chunk_size * 0.15)}"
)
return int(chunk_size * 0.15)
logger.info(f"Using custom chunk overlap: {overlap} tokens")
return overlap
except ValueError:
logger.warning(
f"Invalid OPEN_NOTEBOOK_CHUNK_OVERLAP value: '{overlap_str}'. "
f"Using default: 15% of chunk size"
)
return int(chunk_size * 0.15)
def _get_min_chunk_size() -> int:
"""Get minimum chunk size from environment variable or use default.
Chunks below this token count are dropped. Some splitters (notably the
HTML header splitter on complex pages) can emit single-character or
punctuation-only chunks that produce useless or null embeddings —
llama.cpp's OpenAI-compatible endpoint, for example, returns null vector
elements for such inputs and crashes downstream parsing.
"""
raw = os.getenv("OPEN_NOTEBOOK_MIN_CHUNK_SIZE")
if raw is None:
return 5
try:
value = int(raw)
if value < 0:
logger.warning(
f"OPEN_NOTEBOOK_MIN_CHUNK_SIZE ({value}) cannot be negative. Using 0."
)
return 0
return value
except ValueError:
logger.warning(
f"Invalid OPEN_NOTEBOOK_MIN_CHUNK_SIZE value: '{raw}'. Using default: 5"
)
return 5
# Constants (computed at import time from environment variables)
CHUNK_SIZE = _get_chunk_size()
CHUNK_OVERLAP = _get_chunk_overlap(CHUNK_SIZE)
MIN_CHUNK_SIZE = _get_min_chunk_size()
HIGH_CONFIDENCE_THRESHOLD = 0.8 # Threshold for heuristics to override extension
logger.debug(
f"Chunking configuration: CHUNK_SIZE={CHUNK_SIZE}, "
f"CHUNK_OVERLAP={CHUNK_OVERLAP}, MIN_CHUNK_SIZE={MIN_CHUNK_SIZE}"
)
class ContentType(Enum):
"""Content type for chunking strategy selection."""
HTML = "html"
MARKDOWN = "markdown"
PLAIN = "plain"
# File extension mappings
_EXTENSION_TO_CONTENT_TYPE = {
# HTML
".html": ContentType.HTML,
".htm": ContentType.HTML,
".xhtml": ContentType.HTML,
# Markdown
".md": ContentType.MARKDOWN,
".markdown": ContentType.MARKDOWN,
".mdown": ContentType.MARKDOWN,
".mkd": ContentType.MARKDOWN,
# Plain text (explicit)
".txt": ContentType.PLAIN,
".text": ContentType.PLAIN,
# Code files (treat as plain)
".py": ContentType.PLAIN,
".js": ContentType.PLAIN,
".ts": ContentType.PLAIN,
".java": ContentType.PLAIN,
".c": ContentType.PLAIN,
".cpp": ContentType.PLAIN,
".go": ContentType.PLAIN,
".rs": ContentType.PLAIN,
".rb": ContentType.PLAIN,
".php": ContentType.PLAIN,
".sh": ContentType.PLAIN,
".bash": ContentType.PLAIN,
".zsh": ContentType.PLAIN,
".sql": ContentType.PLAIN,
".json": ContentType.PLAIN,
".yaml": ContentType.PLAIN,
".yml": ContentType.PLAIN,
".xml": ContentType.PLAIN,
".csv": ContentType.PLAIN,
".tsv": ContentType.PLAIN,
}
def detect_content_type_from_extension(
file_path: Optional[str],
) -> Optional[ContentType]:
"""
Detect content type from file extension.
Args:
file_path: Path to the file (can be full path or just filename)
Returns:
ContentType if extension is recognized, None otherwise
"""
if not file_path:
return None
try:
extension = Path(file_path).suffix.lower()
return _EXTENSION_TO_CONTENT_TYPE.get(extension)
except Exception:
return None
def detect_content_type_from_heuristics(text: str) -> Tuple[ContentType, float]:
"""
Detect content type using content heuristics.
Args:
text: The text content to analyze
Returns:
Tuple of (ContentType, confidence_score) where confidence is 0.0-1.0
"""
if not text or len(text) < 10:
return ContentType.PLAIN, 0.5
# Sample first 5000 chars for efficiency
sample = text[:5000]
# Check HTML first (most specific patterns)
html_score = _calculate_html_score(sample)
if html_score >= HIGH_CONFIDENCE_THRESHOLD:
return ContentType.HTML, html_score
# Check Markdown
markdown_score = _calculate_markdown_score(sample)
if markdown_score >= HIGH_CONFIDENCE_THRESHOLD:
return ContentType.MARKDOWN, markdown_score
# Return the higher scoring type, or PLAIN if both are low
if html_score > markdown_score and html_score > 0.3:
return ContentType.HTML, html_score
elif markdown_score > 0.3:
return ContentType.MARKDOWN, markdown_score
else:
return ContentType.PLAIN, 0.6
def _calculate_html_score(text: str) -> float:
"""Calculate confidence score for HTML content."""
score = 0.0
indicators = 0
# Strong indicators
if re.search(r"<!DOCTYPE\s+html", text, re.IGNORECASE):
score += 0.4
indicators += 1
if re.search(r"<html[\s>]", text, re.IGNORECASE):
score += 0.3
indicators += 1
# Structural tags
structural_tags = ["<head", "<body", "<div", "<span", "<p>", "<table", "<form"]
for tag in structural_tags:
if tag.lower() in text.lower():
score += 0.1
indicators += 1
if indicators >= 5:
break
# Header tags
if re.search(r"<h[1-6][\s>]", text, re.IGNORECASE):
score += 0.15
indicators += 1
# Closing tags pattern
if re.search(r"</\w+>", text):
score += 0.1
indicators += 1
return min(score, 1.0)
def _calculate_markdown_score(text: str) -> float:
"""Calculate confidence score for Markdown content."""
score = 0.0
indicators = 0
# Headers (# ## ###) - strong indicator
header_matches = len(re.findall(r"^#{1,6}\s+.+", text, re.MULTILINE))
if header_matches >= 3:
score += 0.35
indicators += 1
elif header_matches >= 1:
score += 0.2
indicators += 1
# Links [text](url) - strong indicator
link_matches = len(re.findall(r"\[.+?\]\(.+?\)", text))
if link_matches >= 2:
score += 0.25
indicators += 1
elif link_matches >= 1:
score += 0.15
indicators += 1
# Code blocks ``` - strong indicator
if re.search(r"^```", text, re.MULTILINE):
score += 0.2
indicators += 1
# Inline code `code`
if re.search(r"`[^`]+`", text):
score += 0.1
indicators += 1
# Lists (-, *, +, or numbered)
list_matches = len(re.findall(r"^[\*\-\+]\s+", text, re.MULTILINE))
list_matches += len(re.findall(r"^\d+\.\s+", text, re.MULTILINE))
if list_matches >= 3:
score += 0.15
indicators += 1
elif list_matches >= 1:
score += 0.08
indicators += 1
# Bold/italic
if re.search(r"\*\*.+?\*\*|__.+?__", text):
score += 0.1
indicators += 1
# Blockquotes
if re.search(r"^>\s+", text, re.MULTILINE):
score += 0.1
indicators += 1
return min(score, 1.0)
def detect_content_type(text: str, file_path: Optional[str] = None) -> ContentType:
"""
Detect content type using file extension (primary) and heuristics (fallback).
Strategy:
1. If file extension is available and recognized, use it as primary
2. If no extension or generic extension (.txt), use heuristics
3. Heuristics can override extension only with very high confidence
Args:
text: The text content
file_path: Optional file path for extension-based detection
Returns:
Detected ContentType
"""
# Try extension-based detection first
extension_type = detect_content_type_from_extension(file_path)
# Get heuristic-based detection
heuristic_type, confidence = detect_content_type_from_heuristics(text)
# If no extension or generic extension, use heuristics
if extension_type is None:
logger.debug(
f"No file extension, using heuristics: {heuristic_type.value} "
f"(confidence: {confidence:.2f})"
)
return heuristic_type
# If extension suggests plain text but heuristics are very confident, override
if extension_type == ContentType.PLAIN and confidence >= HIGH_CONFIDENCE_THRESHOLD:
logger.debug(
f"Extension suggests plain, but heuristics override with "
f"{heuristic_type.value} (confidence: {confidence:.2f})"
)
return heuristic_type
# Otherwise trust the extension
logger.debug(f"Using extension-based content type: {extension_type.value}")
return extension_type
def _get_html_splitter() -> HTMLHeaderTextSplitter:
"""Get HTML header splitter configured for h1, h2, h3."""
headers_to_split_on = [
("h1", "Header 1"),
("h2", "Header 2"),
("h3", "Header 3"),
]
return HTMLHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
def _get_markdown_splitter() -> MarkdownHeaderTextSplitter:
"""Get Markdown header splitter configured for #, ##, ###."""
headers_to_split_on = [
("#", "Header 1"),
("##", "Header 2"),
("###", "Header 3"),
]
return MarkdownHeaderTextSplitter(
headers_to_split_on=headers_to_split_on,
strip_headers=False,
)
def _get_plain_splitter() -> RecursiveCharacterTextSplitter:
"""Get plain text splitter using CHUNK_SIZE and CHUNK_OVERLAP constants."""
return RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
length_function=token_count,
separators=["\n\n", "\n", ". ", ", ", " ", ""],
)
def _apply_secondary_chunking(chunks: List[str]) -> List[str]:
"""
Apply secondary chunking to ensure no chunk exceeds CHUNK_SIZE tokens.
Used when primary splitters (HTML/Markdown) produce oversized chunks.
"""
result = []
secondary_splitter = _get_plain_splitter()
for chunk in chunks:
if token_count(chunk) > CHUNK_SIZE:
# Split oversized chunk
sub_chunks = secondary_splitter.split_text(chunk)
result.extend(sub_chunks)
else:
result.append(chunk)
return result
def chunk_text(
text: str,
content_type: Optional[ContentType] = None,
file_path: Optional[str] = None,
) -> List[str]:
"""
Split text into chunks using appropriate splitter for content type.
Args:
text: The text to chunk
content_type: Optional explicit content type (auto-detected if not provided)
file_path: Optional file path for content type detection
Returns:
List of text chunks, each approximately <= CHUNK_SIZE tokens
"""
if not text or not text.strip():
return []
# Short text doesn't need chunking
text_tokens = token_count(text)
if text_tokens <= CHUNK_SIZE:
return [text]
# Detect content type if not provided
if content_type is None:
content_type = detect_content_type(text, file_path)
logger.debug(f"Chunking text with content type: {content_type.value}")
# Select appropriate splitter
chunks: List[str]
if content_type == ContentType.HTML:
html_splitter = _get_html_splitter()
# HTML splitter returns Document objects
docs = html_splitter.split_text(text)
chunks = [
doc.page_content if hasattr(doc, "page_content") else str(doc)
for doc in docs
]
elif content_type == ContentType.MARKDOWN:
md_splitter = _get_markdown_splitter()
# Markdown splitter returns Document objects
docs = md_splitter.split_text(text)
chunks = [
doc.page_content if hasattr(doc, "page_content") else str(doc)
for doc in docs
]
else:
# Plain text - use recursive splitter directly
chunks = _get_plain_splitter().split_text(text)
# Apply secondary chunking if needed (for HTML/Markdown that may produce large chunks)
if content_type in (ContentType.HTML, ContentType.MARKDOWN):
chunks = _apply_secondary_chunking(chunks)
# Filter out empty chunks
chunks = [c.strip() for c in chunks if c and c.strip()]
# Drop chunks below the minimum token threshold. These are typically
# punctuation or single-character fragments left over from header-based
# splitters; embedding them is wasteful and some providers return null
# vector elements for such inputs (which then crash response parsing).
# Only filter when more than one chunk exists and at least one chunk
# would survive — never return an empty list because of this filter.
if MIN_CHUNK_SIZE > 0 and len(chunks) > 1:
kept = [c for c in chunks if token_count(c) >= MIN_CHUNK_SIZE]
if kept:
dropped = len(chunks) - len(kept)
if dropped > 0:
logger.debug(
f"Dropped {dropped} chunk(s) below MIN_CHUNK_SIZE={MIN_CHUNK_SIZE} tokens"
)
chunks = kept
logger.debug(f"Created {len(chunks)} chunks from {text_tokens} tokens")
return chunks
+208
View File
@@ -0,0 +1,208 @@
"""Context building for chat and podcast generation.
This is the single implementation behind:
- ``POST /api/chat/context`` (`api/routers/chat.py`) — assembles notebook
context from a source/note inclusion config, via
:func:`build_notebook_context`.
- the source-chat graph (`open_notebook/graphs/source_chat.py`) — assembles
a single source plus its insights under a token budget, via
:func:`build_source_context`.
The inclusion config uses string matching on human-readable status values
("not in context", "insights", "full content"). That protocol is shared with
the frontend — do not change it here.
"""
from __future__ import annotations
from typing import Any, Dict, Optional, Tuple
from loguru import logger
from open_notebook.domain.notebook import (
Note,
Notebook,
Source,
SourceInsight,
)
from open_notebook.exceptions import DatabaseOperationError, NotFoundError
from .token_utils import token_count
def _ensure_prefix(table: str, record_id: str) -> str:
"""Ensure a record ID carries its table prefix (`table:id`)."""
prefix = f"{table}:"
return record_id if record_id.startswith(prefix) else f"{prefix}{record_id}"
async def build_notebook_context(
notebook: Notebook,
context_config: Optional[Dict[str, Any]],
) -> Tuple[Dict[str, list], str]:
"""Assemble source/note context for a notebook.
With a config, each entry's status string decides inclusion: "not in"
skips it, "insights" includes the short source context, "full content"
includes the long context (notes only support "full content"). Without a
config, every source and note is included with its short context.
Failures on individual items are logged and skipped — one broken record
never fails the whole request.
Returns:
({"sources": [...], "notes": [...]}, concatenated str() of every
included context dict — used for token/char counting).
"""
context_data: Dict[str, list] = {"sources": [], "notes": []}
total_content = ""
if context_config:
for source_id, status in context_config.get("sources", {}).items():
if "not in" in status:
continue
try:
full_source_id = _ensure_prefix("source", source_id)
try:
source = await Source.get(full_source_id)
except Exception:
continue
if "insights" in status:
source_context = await source.get_context(context_size="short")
context_data["sources"].append(source_context)
total_content += str(source_context)
elif "full content" in status:
source_context = await source.get_context(context_size="long")
context_data["sources"].append(source_context)
total_content += str(source_context)
except Exception as e:
logger.warning(f"Error processing source {source_id}: {str(e)}")
continue
for note_id, status in context_config.get("notes", {}).items():
if "not in" in status:
continue
try:
full_note_id = _ensure_prefix("note", note_id)
note = await Note.get(full_note_id)
if not note:
continue
if "full content" in status:
note_context = note.get_context(context_size="long")
context_data["notes"].append(note_context)
total_content += str(note_context)
except Exception as e:
logger.warning(f"Error processing note {note_id}: {str(e)}")
continue
else:
# Default behavior - include all sources and notes with short context
sources = await notebook.get_sources()
try:
insights_by_source = await SourceInsight.get_for_sources(
[source.id for source in sources if source.id]
)
except Exception as e:
# Match the per-source fallback below: a hiccup fetching
# insights shouldn't fail the whole context request.
logger.warning(f"Error batch-fetching source insights: {str(e)}")
insights_by_source = {}
for source in sources:
try:
source_context = await source.get_context(
context_size="short",
insights=insights_by_source.get(source.id or "", []),
)
context_data["sources"].append(source_context)
total_content += str(source_context)
except Exception as e:
logger.warning(f"Error processing source {source.id}: {str(e)}")
continue
notes = await notebook.get_notes()
for note in notes:
try:
note_context = note.get_context(context_size="short")
context_data["notes"].append(note_context)
total_content += str(note_context)
except Exception as e:
logger.warning(f"Error processing note {note.id}: {str(e)}")
continue
return context_data, total_content
async def build_source_context(
source_id: str, max_tokens: Optional[int] = None
) -> Dict[str, Any]:
"""Assemble a single source's short context plus its insights.
Used by the source-chat graph. If `max_tokens` is given, insights are
dropped (last-fetched first) until the total fits — the source itself is
always kept.
Returns a dict with "sources", "notes" (always empty), "insights",
"total_tokens", "total_items" and per-type counts in "metadata".
"""
try:
sources: list = []
insights: list = []
item_tokens: list[int] = []
try:
full_source_id = _ensure_prefix("source", source_id)
source = await Source.get(full_source_id)
except NotFoundError:
source = None
if source:
source_context = await source.get_context(context_size="short")
sources.append(source_context)
item_tokens.append(token_count(str(source_context)))
for insight in await source.get_insights():
insight_content = {
"id": insight.id,
"source_id": source.id,
"insight_type": insight.insight_type,
"content": insight.content,
}
insights.append(insight_content)
item_tokens.append(token_count(str(insight_content)))
else:
logger.warning(f"Source {source_id} not found")
# Truncate to the token budget: drop insights from the end (the
# source, added first, is dropped only if it alone exceeds the budget).
total_tokens = sum(item_tokens)
if max_tokens:
while total_tokens > max_tokens and item_tokens:
total_tokens -= item_tokens.pop()
if insights:
insights.pop()
else:
sources.pop()
total_items = len(sources) + len(insights)
logger.info(f"Built context with {total_items} items, {total_tokens} tokens")
return {
"sources": sources,
"notes": [],
"insights": insights,
"total_tokens": total_tokens,
"total_items": total_items,
"metadata": {
"source_count": len(sources),
"note_count": 0,
"insight_count": len(insights),
},
}
except Exception as e:
logger.error(f"Error building context: {str(e)}")
raise DatabaseOperationError(f"Failed to build context: {str(e)}")
+269
View File
@@ -0,0 +1,269 @@
"""
Unified embedding utilities for Open Notebook.
Provides centralized embedding generation with support for:
- Single text embedding (with automatic chunking and mean pooling for large texts)
- Batch text embedding (multiple texts with automatic batching)
- Mean pooling for combining multiple embeddings into one
All embedding operations in the application should use these functions
to ensure consistent behavior and proper handling of large content.
"""
import asyncio
import os
from typing import List, Optional
import numpy as np
from loguru import logger
from .chunking import CHUNK_SIZE, ContentType, chunk_text
from .token_utils import token_count
def _get_embedding_batch_size() -> int:
"""
Read the embedding batch size from the environment.
This is intentionally configurable because provider limits vary widely, and
CPU-only local embedding endpoints often need smaller batches than cloud APIs.
"""
raw = os.getenv("OPEN_NOTEBOOK_EMBEDDING_BATCH_SIZE", "50").strip()
try:
value = int(raw)
if value < 1:
raise ValueError
return value
except ValueError:
logger.warning(
"Invalid OPEN_NOTEBOOK_EMBEDDING_BATCH_SIZE='{}'; falling back to 50",
raw,
)
return 50
EMBEDDING_BATCH_SIZE = _get_embedding_batch_size()
EMBEDDING_MAX_RETRIES = 3
EMBEDDING_RETRY_DELAY = 2 # seconds
async def mean_pool_embeddings(embeddings: List[List[float]]) -> List[float]:
"""
Combine multiple embeddings into a single embedding using mean pooling.
Algorithm:
1. Normalize each embedding to unit length
2. Compute element-wise mean
3. Normalize the result to unit length
This approach ensures the final embedding has the same properties as
individual embeddings (unit length) regardless of input count.
Args:
embeddings: List of embedding vectors (each is a list of floats)
Returns:
Single embedding vector (mean pooled and normalized)
Raises:
ValueError: If embeddings list is empty or embeddings have different dimensions
"""
if not embeddings:
raise ValueError("Cannot mean pool empty list of embeddings")
if len(embeddings) == 1:
# Single embedding - just normalize and return
arr = np.array(embeddings[0], dtype=np.float64)
norm = np.linalg.norm(arr)
if norm > 0:
arr = arr / norm
return arr.tolist()
# Convert to numpy array
arr = np.array(embeddings, dtype=np.float64)
# Verify all embeddings have same dimension
if arr.ndim != 2:
raise ValueError(f"Expected 2D array, got shape {arr.shape}")
# Normalize each embedding to unit length
norms = np.linalg.norm(arr, axis=1, keepdims=True)
# Avoid division by zero
norms = np.where(norms > 0, norms, 1.0)
normalized = arr / norms
# Compute mean
mean = np.mean(normalized, axis=0)
# Normalize the result
mean_norm = np.linalg.norm(mean)
if mean_norm > 0:
mean = mean / mean_norm
return mean.tolist()
async def generate_embeddings(
texts: List[str], command_id: Optional[str] = None
) -> List[List[float]]:
"""
Generate embeddings for multiple texts with automatic batching and retry.
Texts are split into batches of EMBEDDING_BATCH_SIZE to avoid exceeding
provider payload limits. Each batch is retried up to EMBEDDING_MAX_RETRIES
times on transient failures.
Args:
texts: List of text strings to embed
command_id: Optional command ID for error logging context
Returns:
List of embedding vectors, one per input text
Raises:
ValueError: If no embedding model is configured
RuntimeError: If embedding generation fails
"""
if not texts:
return []
# Lazy import to avoid circular dependency
from open_notebook.ai.models import model_manager
embedding_model = await model_manager.get_embedding_model()
if not embedding_model:
raise ValueError(
"No embedding model configured. Please configure one in the Models section."
)
model_name = getattr(embedding_model, "model_name", "unknown")
# Log text sizes for debugging
metrics: tuple[int, int, int, int] | None = None
def _get_size_metrics() -> tuple[int, int, int, int]:
nonlocal metrics
if metrics is None:
token_sizes = [token_count(t) for t in texts]
metrics = (
min(token_sizes),
max(token_sizes),
sum(token_sizes),
sum(len(t) for t in texts),
)
return metrics
logger.opt(lazy=True).debug(
"Generating embeddings for {} texts "
"(tokens: min={}, max={}, total={}; chars: total={})",
lambda: len(texts),
lambda: _get_size_metrics()[0],
lambda: _get_size_metrics()[1],
lambda: _get_size_metrics()[2],
lambda: _get_size_metrics()[3],
)
all_embeddings: List[List[float]] = []
total_batches = (len(texts) + EMBEDDING_BATCH_SIZE - 1) // EMBEDDING_BATCH_SIZE
for batch_idx in range(total_batches):
start = batch_idx * EMBEDDING_BATCH_SIZE
end = start + EMBEDDING_BATCH_SIZE
batch = texts[start:end]
for attempt in range(1, EMBEDDING_MAX_RETRIES + 1):
try:
batch_embeddings = await embedding_model.aembed(batch)
all_embeddings.extend(batch_embeddings)
break
except Exception as e:
cmd_context = f" (command: {command_id})" if command_id else ""
if attempt < EMBEDDING_MAX_RETRIES:
logger.debug(
f"Embedding batch {batch_idx + 1}/{total_batches} "
f"attempt {attempt}/{EMBEDDING_MAX_RETRIES} failed "
f"using model '{model_name}'{cmd_context}: {e}. Retrying..."
)
await asyncio.sleep(EMBEDDING_RETRY_DELAY)
else:
logger.debug(
f"Embedding batch {batch_idx + 1}/{total_batches} "
f"failed after {EMBEDDING_MAX_RETRIES} attempts "
f"using model '{model_name}'{cmd_context}: {e}"
)
raise RuntimeError(
f"Failed to generate embeddings using model '{model_name}' "
f"(batch {batch_idx + 1}/{total_batches}, "
f"{len(batch)} texts): {e}"
) from e
logger.debug(f"Generated {len(all_embeddings)} embeddings in {total_batches} batch(es)")
return all_embeddings
async def generate_embedding(
text: str,
content_type: Optional[ContentType] = None,
file_path: Optional[str] = None,
command_id: Optional[str] = None,
) -> List[float]:
"""
Generate a single embedding for text, handling large content via chunking and mean pooling.
For short text (<= CHUNK_SIZE tokens):
- Embeds directly and returns the embedding
For long text (> CHUNK_SIZE tokens):
- Chunks the text using appropriate splitter for content type
- Embeds all chunks in batches
- Combines embeddings via mean pooling
Args:
text: The text to embed
content_type: Optional explicit content type for chunking
file_path: Optional file path for content type detection
command_id: Optional command ID for error logging context
Returns:
Single embedding vector (list of floats)
Raises:
ValueError: If text is empty or no embedding model configured
RuntimeError: If embedding generation fails
"""
if not text or not text.strip():
raise ValueError("Cannot generate embedding for empty text")
text = text.strip()
text_tokens = token_count(text)
# Check if chunking is needed
if text_tokens <= CHUNK_SIZE:
# Short text - embed directly
logger.debug(f"Embedding short text ({text_tokens} tokens) directly")
embeddings = await generate_embeddings([text], command_id=command_id)
return embeddings[0]
# Long text - chunk and mean pool
logger.debug(f"Text exceeds chunk size ({text_tokens} tokens), chunking...")
chunks = chunk_text(text, content_type=content_type, file_path=file_path)
if not chunks:
raise ValueError("Text chunking produced no chunks")
if len(chunks) == 1:
# Single chunk after splitting
embeddings = await generate_embeddings(chunks, command_id=command_id)
return embeddings[0]
logger.debug(f"Embedding {len(chunks)} chunks and mean pooling")
# Embed all chunks in batches
embeddings = await generate_embeddings(chunks, command_id=command_id)
# Mean pool to get single embedding
pooled = await mean_pool_embeddings(embeddings)
logger.debug(f"Mean pooled {len(embeddings)} embeddings into single vector")
return pooled
+198
View File
@@ -0,0 +1,198 @@
"""
Field-level encryption for sensitive data using API keys.
This module provides encryption/decryption for API keys stored in the database.
Fernet uses AES-128-CBC with HMAC-SHA256 for authenticated encryption.
OPEN_NOTEBOOK_ENCRYPTION_KEY accepts **any string**. A Fernet key is derived
from it via SHA-256, so users can set a simple passphrase like
``OPEN_NOTEBOOK_ENCRYPTION_KEY=my-secret`` and it will work.
Usage:
# Encrypt before storing
encrypted = encrypt_value(api_key)
# Decrypt when reading
decrypted = decrypt_value(encrypted)
"""
import base64
import hashlib
import os
from pathlib import Path
from typing import Optional
from cryptography.fernet import Fernet, InvalidToken
from loguru import logger
def get_secret_from_env(var_name: str) -> Optional[str]:
"""
Get a secret from environment, supporting Docker secrets pattern.
Checks for VAR_FILE first (Docker secrets), then falls back to VAR.
Args:
var_name: Base name of the environment variable (e.g., "OPEN_NOTEBOOK_ENCRYPTION_KEY")
Returns:
The secret value, or None if not configured.
"""
# Check for _FILE variant first (Docker secrets)
file_path = os.environ.get(f"{var_name}_FILE")
if file_path:
try:
path = Path(file_path)
if path.exists() and path.is_file():
secret = path.read_text().strip()
if secret:
logger.debug(f"Loaded {var_name} from file: {file_path}")
return secret
else:
logger.warning(f"{var_name}_FILE points to empty file: {file_path}")
else:
logger.warning(f"{var_name}_FILE path does not exist: {file_path}")
except Exception as e:
logger.error(f"Failed to read {var_name} from file {file_path}: {e}")
# Fall back to direct environment variable
return os.environ.get(var_name)
def _get_or_create_encryption_key() -> str:
"""
Get encryption key from environment, requires explicit configuration.
Priority:
1. OPEN_NOTEBOOK_ENCRYPTION_KEY_FILE (Docker secrets)
2. OPEN_NOTEBOOK_ENCRYPTION_KEY (environment variable)
For production deployments, you MUST set OPEN_NOTEBOOK_ENCRYPTION_KEY explicitly!
Returns:
Encryption key string.
Raises:
ValueError: If no encryption key is configured.
"""
# First check environment/Docker secrets
key = get_secret_from_env("OPEN_NOTEBOOK_ENCRYPTION_KEY")
if key:
return key
raise ValueError(
"OPEN_NOTEBOOK_ENCRYPTION_KEY is not set. "
"Set this environment variable to any secret string to enable "
"encrypted storage of API keys in the database."
)
# Lazy-loaded encryption key: initialized on first use, not at import time.
# This prevents the entire app from crashing if the key is not yet configured
# when other modules import from this file.
_ENCRYPTION_KEY: Optional[str] = None
def _get_encryption_key() -> str:
"""Get the encryption key, initializing lazily on first call."""
global _ENCRYPTION_KEY
if _ENCRYPTION_KEY is None:
_ENCRYPTION_KEY = _get_or_create_encryption_key()
return _ENCRYPTION_KEY
def _ensure_fernet_key(key: str) -> str:
"""
Derive a valid Fernet key from an arbitrary string via SHA-256.
Any string is accepted as input. The key is derived by hashing it with
SHA-256 and encoding the result as URL-safe base64.
"""
derived = hashlib.sha256(key.encode()).digest()
return base64.urlsafe_b64encode(derived).decode()
def get_fernet() -> Fernet:
"""
Get Fernet instance with the configured encryption key.
Returns:
Fernet instance.
Raises:
ValueError: If encryption key is not configured.
"""
return Fernet(_ensure_fernet_key(_get_encryption_key()).encode())
def encrypt_value(value: str) -> str:
"""
Encrypt a string value using Fernet symmetric encryption.
Args:
value: The plain text string to encrypt.
Returns:
Base64-encoded encrypted string.
Raises:
ValueError: If encryption is not configured.
"""
fernet = get_fernet()
return fernet.encrypt(value.encode()).decode()
def looks_like_fernet_token(s: str) -> bool:
"""
Check if string looks like a Fernet encrypted token.
Fernet tokens are versioned (1 byte) + timestamp (8 bytes) + IV (16 bytes)
+ ciphertext (variable, multiple of 16 with PKCS7 padding) + HMAC (32 bytes).
Minimum decoded size is 73 bytes (1+8+16+16+32) for the smallest payload.
"""
if len(s) < 100: # Base64 of 73 bytes = ~100 chars minimum
return False
try:
decoded = base64.urlsafe_b64decode(s)
# Fernet: version(1) + timestamp(8) + IV(16) + ciphertext(>=16) + HMAC(32)
# Minimum 73 bytes, ciphertext must be multiple of 16 (AES block size)
if len(decoded) < 73:
return False
ciphertext_len = len(decoded) - 1 - 8 - 16 - 32
return ciphertext_len > 0 and ciphertext_len % 16 == 0
except Exception:
return False
def decrypt_value(value: str) -> str:
"""
Decrypt a Fernet-encrypted string value.
Handles graceful fallback for legacy unencrypted data.
Args:
value: The encrypted string (or plain text for legacy data).
Returns:
Decrypted plain text string, or original value if not encrypted.
Raises:
ValueError: If encryption is not configured or if decryption fails
for what appears to be encrypted data (wrong key).
"""
fernet = get_fernet()
try:
return fernet.decrypt(value.encode()).decode()
except InvalidToken:
if looks_like_fernet_token(value):
# Looks like encrypted data but failed to decrypt - likely wrong key
raise ValueError(
"Decryption failed: data appears to be encrypted but key is incorrect. "
"Check OPEN_NOTEBOOK_ENCRYPTION_KEY configuration."
)
# Not a valid token - treat as legacy plaintext
return value
except Exception as e:
logger.error(f"Decryption failed: {e}")
raise ValueError(f"Decryption failed: {str(e)}")
+103
View File
@@ -0,0 +1,103 @@
"""
Error classification utility for LLM provider errors.
Maps raw exceptions from AI providers/Esperanto/LangChain to user-friendly
error messages and appropriate exception types.
"""
from loguru import logger
from open_notebook.exceptions import (
AuthenticationError,
ConfigurationError,
ExternalServiceError,
NetworkError,
OpenNotebookError,
RateLimitError,
)
# Classification rules: (keywords, exception_class, user_message or None to pass through)
_CLASSIFICATION_RULES: list[tuple[list[str], type[OpenNotebookError], str | None]] = [
# Authentication errors
(
["authentication", "unauthorized", "invalid api key", "invalid_api_key", "401"],
AuthenticationError,
"Authentication failed. Please check your API key in Settings -> Credentials.",
),
# Rate limit errors
(
["rate limit", "rate_limit", "429", "too many requests", "quota exceeded"],
RateLimitError,
"Rate limit exceeded. Please wait a moment and try again.",
),
# Model not found (pass through original message)
(
["model not found", "does not exist", "model_not_found"],
ConfigurationError,
None,
),
# Configuration errors from provision.py (pass through)
(
["no model configured", "please go to settings"],
ConfigurationError,
None,
),
# Network errors
(
["connecterror", "timeoutexception", "connection refused", "connection error", "timed out", "timeout"],
NetworkError,
"Could not connect to the AI provider. Please check your network connection and provider URL.",
),
# Context length errors
(
["context length", "token limit", "maximum context", "context_length_exceeded", "max_tokens"],
ExternalServiceError,
"Content too large for the selected model. Try using a smaller selection or a model with a larger context window.",
),
# Payload too large errors
(
["413", "payload too large", "request entity too large"],
ExternalServiceError,
"The request payload is too large for the AI provider. Try reducing the content size or using a different model.",
),
# Provider availability errors
(
["500", "502", "503", "service unavailable", "overloaded", "internal server error"],
ExternalServiceError,
"The AI provider is temporarily unavailable. Please try again in a few minutes.",
),
]
def classify_error(exception: BaseException) -> tuple[type[OpenNotebookError], str]:
"""
Classify a raw exception into a user-friendly error type and message.
Args:
exception: Any exception from LLM providers/Esperanto/LangChain
Returns:
Tuple of (exception_class, user_friendly_message)
"""
error_str = str(exception).lower()
error_type_name = type(exception).__name__.lower()
combined = f"{error_type_name}: {error_str}"
for keywords, exc_class, message in _CLASSIFICATION_RULES:
for keyword in keywords:
if keyword in combined:
user_message = message if message is not None else _truncate(str(exception))
return exc_class, user_message
# Unclassified error - log for future improvement
logger.warning(
f"Unclassified LLM error ({type(exception).__name__}): {exception}"
)
return ExternalServiceError, f"AI service error: {_truncate(str(exception))}"
def _truncate(text: str, max_length: int = 200) -> str:
"""Truncate text to max_length to avoid leaking verbose internal details."""
if len(text) <= max_length:
return text
return text[:max_length] + "..."
+23
View File
@@ -0,0 +1,23 @@
import asyncio
from langchain_core.runnables import RunnableConfig
from loguru import logger
async def get_session_message_count(graph, session_id: str) -> int:
"""Get message count from LangGraph state, returns 0 on error."""
try:
# Use sync get_state() in a thread (SqliteSaver doesn't support async)
thread_state = await asyncio.to_thread(
graph.get_state,
config=RunnableConfig(configurable={"thread_id": session_id}),
)
if (
thread_state
and thread_state.values
and "messages" in thread_state.values
):
return len(thread_state.values["messages"])
except Exception as e:
logger.warning(f"Could not fetch message count for session {session_id}: {e}")
return 0
+15
View File
@@ -0,0 +1,15 @@
"""Utilities for working with Pydantic models."""
from pydantic import BaseModel
def full_model_dump(model):
"""Recursively dump Pydantic models nested inside dicts/lists to plain data."""
if isinstance(model, BaseModel):
return model.model_dump()
elif isinstance(model, dict):
return {k: full_model_dump(v) for k, v in model.items()}
elif isinstance(model, list):
return [full_model_dump(item) for item in model]
else:
return model
+145
View File
@@ -0,0 +1,145 @@
"""
Text utilities for Open Notebook.
Extracted from main utils to avoid circular imports.
"""
import re
import unicodedata
from typing import Tuple
# Patterns for matching thinking content in AI responses
# Standard pattern: <think>...</think>
THINK_PATTERN = re.compile(r"<think>(.*?)</think>", re.DOTALL)
# Pattern for malformed output: content</think> (missing opening tag)
THINK_PATTERN_NO_OPEN = re.compile(r"^(.*?)</think>", re.DOTALL)
def remove_non_ascii(text: str) -> str:
"""Remove non-ASCII characters from text."""
return re.sub(r"[^\x00-\x7F]+", "", text)
def remove_non_printable(text: str) -> str:
"""Remove non-printable characters from text."""
# Replace any special Unicode whitespace characters with a regular space
text = re.sub(r"[\u2000-\u200B\u202F\u205F\u3000]", " ", text)
# Replace unusual line terminators with a single newline
text = re.sub(r"[\u2028\u2029\r]", "\n", text)
# Remove control characters, except newlines and tabs
text = "".join(
char for char in text if unicodedata.category(char)[0] != "C" or char in "\n\t"
)
# Replace non-breaking spaces with regular spaces
text = text.replace("\xa0", " ").strip()
# Keep letters (including accented ones), numbers, spaces, newlines, tabs, and basic punctuation
return re.sub(r"[^\w\s.,!?\-\n\t]", "", text, flags=re.UNICODE)
def parse_thinking_content(content: str) -> Tuple[str, str]:
"""
Parse message content to extract thinking content from <think> tags.
Handles both well-formed tags and malformed output where the opening
<think> tag is missing but </think> is present.
Args:
content (str): The original message content
Returns:
Tuple[str, str]: (thinking_content, cleaned_content)
- thinking_content: Content from within <think> tags
- cleaned_content: Original content with <think> blocks removed
Example:
>>> content = "<think>Let me analyze this</think>Here's my answer"
>>> thinking, cleaned = parse_thinking_content(content)
>>> print(thinking)
"Let me analyze this"
>>> print(cleaned)
"Here's my answer"
"""
# Input validation
if not isinstance(content, str):
return "", str(content) if content is not None else ""
# Limit processing for very large content (100KB limit)
if len(content) > 100000:
return "", content
# Find all well-formed thinking blocks
thinking_matches = THINK_PATTERN.findall(content)
if thinking_matches:
# Join all thinking content with double newlines
thinking_content = "\n\n".join(match.strip() for match in thinking_matches)
# Remove all <think>...</think> blocks from the original content
cleaned_content = THINK_PATTERN.sub("", content)
# Clean up extra whitespace
cleaned_content = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned_content).strip()
return thinking_content, cleaned_content
# Handle malformed output: content</think> (missing opening tag)
# Some models like Nemotron output thinking without the opening <think> tag
malformed_match = THINK_PATTERN_NO_OPEN.match(content)
if malformed_match:
thinking_content = malformed_match.group(1).strip()
# Remove the thinking content and </think> tag
cleaned_content = content[malformed_match.end() :].strip()
return thinking_content, cleaned_content
return "", content
def clean_thinking_content(content: str) -> str:
"""
Remove thinking content from AI responses, returning only the cleaned content.
This is a convenience function for cases where you only need the cleaned
content and don't need access to the thinking process.
Args:
content (str): The original message content with potential <think> tags
Returns:
str: Content with <think> blocks removed and whitespace cleaned
Example:
>>> content = "<think>Let me think...</think>Here's the answer"
>>> clean_thinking_content(content)
"Here's the answer"
"""
_, cleaned_content = parse_thinking_content(content)
return cleaned_content
def extract_text_content(content) -> str:
"""Extract text from LLM response content.
Handles both plain string responses and structured content formats
(e.g. Gemini's envelope format):
[{'type': 'text', 'text': '...', 'extras': {...}}]
Args:
content: The content from an AI message, either a string or a list of parts.
Returns:
The extracted text content as a string.
"""
if isinstance(content, str):
return content
if isinstance(content, list):
text_parts = []
for part in content:
if isinstance(part, dict) and "text" in part:
text_parts.append(part["text"])
elif isinstance(part, str):
text_parts.append(part)
return "".join(text_parts)
return str(content)
+57
View File
@@ -0,0 +1,57 @@
"""
Token utilities for Open Notebook.
Handles token counting and cost calculations for language models.
"""
import os
from open_notebook.config import TIKTOKEN_CACHE_DIR
# Set tiktoken cache directory before importing tiktoken to ensure
# tokenizer encodings are cached persistently in the data folder
os.environ["TIKTOKEN_CACHE_DIR"] = TIKTOKEN_CACHE_DIR
def token_count(input_string: str) -> int:
"""
Count the number of tokens in the input string using the 'o200k_base' encoding.
Args:
input_string (str): The input string to count tokens for.
Returns:
int: The number of tokens in the input string.
"""
try:
import tiktoken
encoding = tiktoken.get_encoding("o200k_base")
# disallowed_special=() treats sequences like "<|endoftext|>" as ordinary
# text instead of raising ValueError. User/source content can legitimately
# contain these substrings, and we only need a token count here.
tokens = encoding.encode(input_string, disallowed_special=())
return len(tokens)
except (ImportError, OSError) as e:
# Fallback: handles ImportError (tiktoken not installed) AND network/OS
# errors such as urllib.error.URLError or ConnectionError raised in
# offline environments when the encoding file cannot be downloaded.
from loguru import logger
logger.warning(
"tiktoken unavailable, falling back to word-count estimation: {}", e
)
return int(len(input_string.split()) * 1.3)
def token_cost(token_count: int, cost_per_million: float = 0.150) -> float:
"""
Calculate the cost of tokens based on the token count and cost per million tokens.
Args:
token_count (int): The number of tokens.
cost_per_million (float): The cost per million tokens. Default is 0.150.
Returns:
float: The calculated cost for the given token count.
"""
return cost_per_million * (token_count / 1_000_000)
+136
View File
@@ -0,0 +1,136 @@
"""
URL validation shared by the credentials API and the AI provisioning layer.
Lives in open_notebook.utils (not api/) so it can be imported both by the API
layer (credential create/update, source-URL ingestion) and by open_notebook's
own AI layer (connection_tester.py, ModelManager) to re-validate a
provider-configured URL immediately before it is actually used - closing most
of the DNS-rebinding TOCTOU window left by validating only once, at save time.
"""
import asyncio
import ipaddress
import socket
from urllib.parse import urlparse
# AWS's IMDSv6 metadata endpoint. Unlike the IPv4 metadata address
# (169.254.169.254), this is a Unique Local Address, not link-local, so it is
# not caught by ip.is_link_local and must be checked explicitly.
_AWS_IMDS_V6_ADDRESS = ipaddress.ip_address("fd00:ec2::254")
async def validate_url(url: str, provider: str) -> None:
"""
Validate URL format for API endpoints.
This is a self-hosted application, so we allow:
- Private IPs (10.x, 172.16-31.x, 192.168.x) for self-hosted services
- Localhost for local services (Ollama, LM Studio, etc.)
We only block:
- Invalid schemes (must be http or https)
- Malformed URLs
- Link-local addresses (169.254.x.x) - used for cloud metadata endpoints
- AWS's IPv6 metadata address (fd00:ec2::254)
- Hostnames that resolve to any of the above
Args:
url: The URL to validate
provider: The provider name (for logging/context)
Raises:
ValueError: If the URL is invalid
"""
if not url or not url.strip():
return # Empty URLs handled elsewhere
try:
parsed = urlparse(url.strip())
# Validate scheme - only http/https allowed
if parsed.scheme not in ("http", "https"):
raise ValueError(
f"Invalid URL scheme: '{parsed.scheme}'. Only http and https are allowed."
)
# Extract hostname
hostname = parsed.hostname
if not hostname:
raise ValueError("Invalid URL: hostname could not be determined.")
# Try to parse as IP address to check for dangerous addresses
try:
ip = ipaddress.ip_address(hostname)
_reject_dangerous_ip(ip, hostname)
except ValueError as ve:
# Re-raise our own ValueErrors
if "Link-local" in str(ve) or "Invalid URL" in str(ve) or "metadata" in str(ve):
raise
# Not an IP address, it's a hostname - need to resolve and check
try:
# Resolve hostname to IP address. This is a blocking call -
# run it off the event loop so a slow/hanging DNS lookup
# doesn't stall every other concurrent request (this is
# called on the hot path of model provisioning, potentially
# once per chat message/transformation).
resolved_ips = await asyncio.to_thread(socket.getaddrinfo, hostname, None)
for family, _, _, _, sockaddr in resolved_ips:
ip_addr = sockaddr[0]
try:
parsed_ip = ipaddress.ip_address(ip_addr)
_reject_dangerous_ip(parsed_ip, hostname, resolved=True)
except ValueError as inner_ve:
if "link-local" in str(inner_ve).lower() or "metadata" in str(inner_ve).lower():
raise
# Skip non-IP addresses (e.g., IPv6 zones)
continue
except socket.gaierror:
# Could not resolve hostname - allow it since the URL may be
# valid in the deployment environment (e.g., Azure endpoints,
# internal DNS names). We only block link-local addresses.
pass
except ValueError:
raise
except Exception:
raise ValueError("Invalid URL format. Check server logs for details.")
def _reject_dangerous_ip(
ip: "ipaddress.IPv4Address | ipaddress.IPv6Address",
hostname: str,
resolved: bool = False,
) -> None:
"""Raise ValueError if `ip` is a link-local or cloud-metadata address."""
is_ipv4_mapped_link_local = (
hasattr(ip, "ipv4_mapped") and ip.ipv4_mapped and ip.ipv4_mapped.is_link_local
)
# Block link-local addresses (169.254.x.x / fe80::/10) - used for cloud
# metadata - including IPv4-mapped IPv6 addresses pointing to link-local
# (e.g. ::ffff:169.254.169.254 bypasses IPv6 is_link_local check).
if ip.is_link_local or is_ipv4_mapped_link_local:
if resolved:
raise ValueError(
f"Hostname '{hostname}' resolves to a link-local address (169.254.x.x) "
"which is not allowed for security reasons. These addresses are used "
"for cloud metadata endpoints."
)
raise ValueError(
"Link-local addresses (169.254.x.x) are not allowed for security reasons. "
"These addresses are used for cloud metadata endpoints."
)
# Block AWS's IMDSv6 metadata address - a Unique Local Address, not
# link-local, so it needs its own explicit check.
if ip == _AWS_IMDS_V6_ADDRESS:
if resolved:
raise ValueError(
f"Hostname '{hostname}' resolves to the AWS IMDSv6 metadata address "
"(fd00:ec2::254), which is not allowed for security reasons."
)
raise ValueError(
"The AWS IMDSv6 metadata address (fd00:ec2::254) is not allowed for "
"security reasons."
)
+153
View File
@@ -0,0 +1,153 @@
"""
Version utilities for Open Notebook.
Handles version comparison, GitHub version fetching, and package version management.
"""
from importlib.metadata import PackageNotFoundError, version
from urllib.parse import urlparse
import requests # type: ignore
import tomli
from packaging.version import parse as parse_version
async def get_version_from_github_async(repo_url: str, branch: str = "main") -> str:
"""
Fetch and parse the version from pyproject.toml in a public GitHub repository (async).
"""
from urllib.parse import urlparse
import httpx
import tomli
# Parse the GitHub URL
parsed_url = urlparse(repo_url)
if "github.com" not in parsed_url.netloc:
raise ValueError("Not a GitHub URL")
# Extract owner and repo name from path
path_parts = parsed_url.path.strip("/").split("/")
if len(path_parts) < 2:
raise ValueError("Invalid GitHub repository URL")
owner, repo = path_parts[0], path_parts[1]
# Construct raw content URL for pyproject.toml
raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/pyproject.toml"
# Fetch the file with timeout using httpx
async with httpx.AsyncClient(timeout=10.0) as client:
response = await client.get(raw_url)
response.raise_for_status()
# Parse TOML content
pyproject_data = tomli.loads(response.text)
# Try to find version
try:
# Check tool.poetry.version
version_str = pyproject_data["tool"]["poetry"]["version"]
except KeyError:
try:
# Check project.version
version_str = pyproject_data["project"]["version"]
except KeyError:
raise KeyError("Version not found in pyproject.toml")
return version_str
def get_version_from_github(repo_url: str, branch: str = "main") -> str:
"""
Fetch and parse the version from pyproject.toml in a public GitHub repository.
Args:
repo_url (str): URL of the GitHub repository
branch (str): Branch name to fetch from (defaults to "main")
Returns:
str: Version string from pyproject.toml
Raises:
ValueError: If the URL is not a valid GitHub repository URL
requests.RequestException: If there's an error fetching the file
KeyError: If version information is not found in pyproject.toml
"""
# Parse the GitHub URL
parsed_url = urlparse(repo_url)
if "github.com" not in parsed_url.netloc:
raise ValueError("Not a GitHub URL")
# Extract owner and repo name from path
path_parts = parsed_url.path.strip("/").split("/")
if len(path_parts) < 2:
raise ValueError("Invalid GitHub repository URL")
owner, repo = path_parts[0], path_parts[1]
# Construct raw content URL for pyproject.toml
raw_url = (
f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/pyproject.toml"
)
# Fetch the file with timeout
response = requests.get(raw_url, timeout=10)
response.raise_for_status()
# Parse TOML content
pyproject_data = tomli.loads(response.text)
# Try to find version in different possible locations
try:
# Check project.version first (poetry style)
version = pyproject_data["tool"]["poetry"]["version"]
except KeyError:
try:
# Check project.version (standard style)
version = pyproject_data["project"]["version"]
except KeyError:
raise KeyError("Version not found in pyproject.toml")
return version
def get_installed_version(package_name: str) -> str:
"""
Get the version of an installed package.
Args:
package_name (str): Name of the installed package
Returns:
str: Version string of the installed package
Raises:
PackageNotFoundError: If the package is not installed
"""
try:
return version(package_name)
except PackageNotFoundError:
raise PackageNotFoundError(f"Package '{package_name}' not found")
def compare_versions(version1: str, version2: str) -> int:
"""
Compare two semantic versions.
Args:
version1 (str): First version string
version2 (str): Second version string
Returns:
int: -1 if version1 < version2
0 if version1 == version2
1 if version1 > version2
"""
v1 = parse_version(version1)
v2 = parse_version(version2)
if v1 < v2:
return -1
elif v1 > v2:
return 1
else:
return 0