chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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)
|
||||
@@ -0,0 +1 @@
|
||||
@AGENTS.md
|
||||
@@ -0,0 +1,2 @@
|
||||
# AI infrastructure module
|
||||
# Contains model configuration, provisioning, and management
|
||||
Binary file not shown.
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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},
|
||||
)
|
||||
@@ -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);
|
||||
|
||||
};
|
||||
@@ -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")
|
||||
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
Domain models for Open Notebook.
|
||||
|
||||
This module exports the core domain models used throughout the application.
|
||||
"""
|
||||
|
||||
__all__: list[str] = []
|
||||
@@ -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()
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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]
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
@@ -0,0 +1,2 @@
|
||||
# Podcasts module
|
||||
# Contains podcast episode models, profiles, and generation logic
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)}")
|
||||
@@ -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
|
||||
@@ -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)}")
|
||||
@@ -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] + "..."
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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."
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user