admin-valentin/open-notebook
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #14 from lfnovo/model_mgmt

Browse source 
Model mgmt
This commit is contained in:
Luis Novo 2024-10-30 18:30:17 -03:00 committed by GitHub
commit 53883aab4d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 1038 additions and 533 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
.env.example
View file

@ -1,37 +1,22 @@
# DEFAULT MODEL_CONFIGURATIONS
DEFAULT_MODEL="openai/gpt-4o-mini"
SUMMARIZATION_MODEL="openai/gpt-4o-mini"
# OPENAI
# USE MODEL NAMES AS "openai/<modelname>"
# EXAMPLE - openai/gpt-4o-mini
OPENAI_API_KEY=
# ANTHROPIC
# USE MODEL NAMES AS "anthropic/<modelname>"
# EXAMPLE - anthropic/claude-3-5-sonnet-20240620
# ANTHROPIC_API_KEY=
# GEMINI
# USE MODEL NAMES AS "gemini/<modelname>"
# EXAMPLE - gemini/gemini-1.5-pro-002
# GEMINI_API_KEY=
# VERTEXAI
# USE MODEL NAMES AS "vertexai/<modelname>"
# EXAMPLE - vertexai/gemini-1.5-pro-002
# VERTEX_PROJECT=my-google-cloud-project-name
# GOOGLE_APPLICATION_CREDENTIALS=./google-credentials.json
# OLLAMA
# USE MODEL NAMES AS "ollama/<modelname>"
# EXAMPLE - ollama/gemma2
# OLLAMA_API_BASE="http://10.20.30.20:11434"
# OPEN ROUTER
# USE MODEL NAMES AS "openrouter/<modelname>"
# EXAMPLE - openrouter/nvidia/llama-3.1-nemotron-70b-instruct
# OPENROUTER_BASE_URL="https://openrouter.ai/api/v1"
# OPENROUTER_API_KEY=

4
Makefile
View file

@ -51,3 +51,7 @@ docker-update-latest: docker-buildx-prepare
# Release with latest
docker-release-all: docker-release docker-update-latest
dev:
docker compose -f docker-compose.dev.yml up --build

23
README.md
View file

@ -29,7 +29,7 @@ services:
open_notebook:
image: lfnovo/open_notebook:latest
ports:
- "8502:8502"
- "8080:8502"
env_file:
- ./docker.env
depends_on:
@ -52,16 +52,29 @@ Go to the [Usage](docs/USAGE.md) page to learn how to use all features.
- **Multi-Notebook Support**: Organize your research across multiple notebooks effortlessly.
- **Broad Content Integration**: Works with links, PDFs, TXT files, PowerPoint presentations, YouTube videos, and pasted text (audio/video support coming soon).
- **Multi-model support**: Open AI, Anthropic, Gemini, Vertex AI, Open Router, Ollama.
- **Podcast Generator**: Automatically convert your notes into a podcast format.
- **Broad Content Integration**: Works with links, PDFs, EPUB, Office, TXT, Markdown files, YouTube videos, Audio files, Video files and pasted text.
- **AI-Powered Notes**: Write notes yourself or let the AI assist you in generating insights.
- **Recursive Summarization**: Tackle large content by recursively summarizing it.
- **Integrated Search Engines**: Built-in full-text and vector search for faster information retrieval.
- **Fine-Grained Context Management**: Choose exactly what to share with the AI to maintain control.
- **Podcast Generator**: Automatically convert your notes into a podcast format.
- **Multi-model support**: Open AI, Anthropic, Gemini, Vertex AI, Open Router, Ollama.
## 🚀 New Features
### v0.0.7 - Model Management 🗂️
- Manage your AI models and providers in a single interface
- Define default models for several tasks such as chat, transformation, embedding, etc
- Enabled support for Embedding models from Gemini, Vertex and Ollama
### v0.0.6 - ePub and Office files support 📄
You can now process ePub and Office files (Word, Excel, PowerPoint), extracting text and insights from them. Perfect for books, reports, presentations, and more.
### v0.0.5 - Audio and Video support 📽️
You can now process audio and video files, extracting transcripts and insights from them. Perfect for podcasts, interviews, lectures, and more.
### v0.0.4 - Podcasts 🎙️
You can now build amazing custom podcasts based on your own data. Customize your speakers, episode structure, cadence, voices, etc.

62
app_home.py
View file

@ -1,29 +1,43 @@
import streamlit as st
from open_notebook.exceptions import InvalidDatabaseSchema, NoSchemaFound
from open_notebook.repository import check_database_version, execute_migration
from open_notebook.database.migrate import MigrationManager
# from open_notebook.config import DEFAULT_MODELS
from open_notebook.domain.models import DefaultModels
from stream_app.utils import version_sidebar
try:
version_sidebar()
check_database_version()
default_models = DefaultModels.load()
version_sidebar()
mm = MigrationManager()
if mm.needs_migration:
st.warning("The Open Notebook database needs a migration to run properly.")
if st.button("Run Migration"):
mm.run_migration_up()
st.success("Migration successful")
st.rerun()
elif (
not default_models.default_chat_model
or not default_models.default_transformation_model
):
st.warning(
"You don't have default chat and transformation models selected. Please, select them on the settings page."
)
elif not default_models.default_embedding_model:
st.warning(
"You don't have a default embedding model selected. Vector search will not be possible and your assistant will be less able to answer your queries. Please, select one on the settings page."
)
elif not default_models.default_speech_to_text_model:
st.warning(
"You don't have a default speech to text model selected. Your assistant will not be able to transcribe audio. Please, select one on the settings page."
)
elif not default_models.default_text_to_speech_model:
st.warning(
"You don't have a default text to speech model selected. Your assistant will not be able to generate audio and podcasts. Please, select one on the settings page."
)
elif not default_models.large_context_model:
st.warning(
"You don't have a large context model selected. Your assistant will not be able to process large documents. Please, select one on the settings page."
)
else:
st.switch_page("pages/2_📒_Notebooks.py")
except NoSchemaFound as e:
st.warning(e)
if st.button("Create Schema.."):
try:
execute_migration("db_setup.surrealql")
st.success("Schema created successfully")
st.rerun()
except Exception as e:
st.error(e)
except InvalidDatabaseSchema as e:
st.warning(e)
if st.button("Execute Migration.."):
try:
execute_migration("0_0_1_to_0_0_2.surrealql")
st.success("Migration executed successfully")
st.rerun()
except Exception as e:
st.error(e)
st.stop()

82
database/0_0_1_to_0_0_2.surrealql
View file

@ -1,82 +0,0 @@
DEFINE FIELD full_text ON TABLE source TYPE option<string>;
REMOVE TABLE IF EXISTS source_chunk;
REMOVE INDEX IF EXISTS idx_source_full ON TABLE source_chunk;
DEFINE FIELD IF NOT EXISTS archived ON TABLE notebook TYPE option<bool> DEFAULT False;
DEFINE INDEX idx_source_full ON TABLE source_chunk COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
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 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 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 podcast_config SCHEMALESS;
UPDATE open_notebook:database_info SET
version= "0.0.2";

40
docs/SETUP.md
View file

@ -63,7 +63,6 @@ services:
- "8080:8502"
environment:
- OPENAI_API_KEY=API_KEY
- DEFAULT_MODEL=gpt-4o-mini
- SURREAL_ADDRESSsurrealdb
- SURREAL_PORT=8000
- SURREAL_USER=root
@ -105,7 +104,7 @@ or the shourcut
make run
```
## Setting up the providers
## Setting up the providers and models
Several new providers are supported now:
@ -121,30 +120,33 @@ All providers are installed out of the box. All you need to do is to setup the e
Please refer to the `.env.example` file for instructions on which ENV variables are necessary for each.
### Use provider-modelname convention
### Create models on the Settings page
You should prepend the provider name to the model_name when setting up your env variables, examples:
Go to the settings page and create your different models.
- openai/gpt-4o-mini
- anthropic/claude-3-5-sonnet-20240620
- ollama/gemma2
- openrouter/nvidia/llama-3.1-nemotron-70b-instruct
- vertexai/gemini-1.5-flash-001
- gemini/gemini-1.5-flash-001
| Model Type | Supported Providers |
|------------|-----------|
| Language | OpenAI, Anthropic, Open Router, LiteLLM, Vertex AI, Vertex AI, Anthropic, Gemini, Ollama |
| Embedding | OpenAI, Gemini, Vertex AI, Ollama |
| Speech to Text | OpenAI |
| Text to Speech | OpenAI, ElevenLabs |
__There will be a UI configuration for models in the coming days.__
## Setup 2 models for more flexibility
> 📝 **Notice:** For complete usage of all the features, you need to setup at least 4 models (one of each type).
There are 2 configurations for models at this point:
After setting up the models, head to the Model Defaults tab to define the default models. There are several defaults to setup.
```
DEFAULT_MODEL="openai/gpt-4o-mini"
SUMMARIZATION_MODEL="openrouter/nvidia/llama-3.1-nemotron-70b-instruct"
```
- **DEFAULT_MODEL** is used by the chat tool
- **SUMMARIZATION_MODEL (optional)** is used on the content summarization
| Model Default | Purpose |
|------------|-----------|
| Chat Model | Will be used on all chats |
| Transformation Model | Will be used for summaries, insights, etc |
| Large Context | For content higher then 110k tokens (use Gemini here) |
| Speech to Text | For transcribing text from your audio/video uploads |
| Text to Speech | For generating podcasts |
| Embedding | For creating vector representation of content |
All model types and defaults are required for now. If you are not sure which to pick, go with OpenAI, the only one that covers all possible model types.
The reason for opting for this route is because different LLMs, will behave better/worse depending on the type of request and type of tools offered. So it makes sense to build a more refined system to decide which model should process which task.

107
database/db_setup.surrealql → migrations/1.surrealql
View file

@ -1,92 +1,78 @@
REMOVE table IF EXISTS source;
REMOVE table IF EXISTS reference;
REMOVE table IF EXISTS notebook;
REMOVE table IF EXISTS note;
REMOVE table IF EXISTS artifact;
REMOVE table IF EXISTS source_chunk;
REMOVE table IF EXISTS source_insight;
REMOVE ANALYZER IF EXISTS my_analyzer;
REMOVE FUNCTION IF EXISTS fn::text_search;
REMOVE INDEX IF EXISTS idx_source_full ON TABLE source_chunk;
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_source_title ON TABLE source;
REMOVE INDEX IF EXISTS idx_note_title ON TABLE note;
DEFINE TABLE IF NOT EXISTS source SCHEMAFULL;
DEFINE FIELD asset
DEFINE FIELD IF NOT EXISTS
asset
ON TABLE source
FLEXIBLE TYPE option<object>;
DEFINE FIELD title ON TABLE source TYPE option<string>;
DEFINE FIELD full_text ON TABLE source TYPE option<string>;
DEFINE FIELD topics ON TABLE source TYPE option<array<string>>;
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 created ON source DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD updated ON source DEFAULT time::now() VALUE time::now();
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 source ON TABLE source_embedding TYPE record<source>;
DEFINE FIELD order ON TABLE source_embedding TYPE int;
DEFINE FIELD content ON TABLE source_embedding TYPE string;
DEFINE FIELD embedding ON TABLE source_embedding TYPE array<float>;
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 source ON TABLE source_insight TYPE record<source>;
DEFINE FIELD insight_type ON TABLE source_insight TYPE string;
DEFINE FIELD content ON TABLE source_insight TYPE string;
DEFINE FIELD embedding ON TABLE source_insight TYPE array<float>;
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 source_delete ON TABLE source WHEN ($after == NONE) THEN {
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 title ON TABLE note TYPE option<string>;
DEFINE FIELD summary ON TABLE note TYPE option<string>;
DEFINE FIELD content ON TABLE note TYPE option<string>;
DEFINE FIELD embedding ON TABLE note TYPE array<float>;
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 created ON note DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD updated ON note DEFAULT time::now() VALUE time::now();
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 name ON TABLE notebook TYPE option<string>;
DEFINE FIELD description ON TABLE notebook TYPE option<string>;
DEFINE FIELD archived ON TABLE notebook TYPE option<bool> DEFAULT False;
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 created ON notebook DEFAULT time::now() VALUE $before OR time::now();
DEFINE FIELD updated ON notebook DEFAULT time::now() VALUE time::now();
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 reference
DEFINE TABLE IF NOT EXISTS reference
TYPE RELATION
FROM source TO notebook;
DEFINE TABLE artifact
DEFINE TABLE IF NOT EXISTS artifact
TYPE RELATION
FROM note TO notebook;
-- entender o analyzer
DEFINE ANALYZER my_analyzer TOKENIZERS blank,class,camel,punct FILTERS snowball(english), lowercase;
DEFINE TABLE IF NOT EXISTS podcast_config SCHEMALESS;
DEFINE INDEX idx_source_title ON TABLE source COLUMNS title SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX idx_source_full_text ON TABLE source COLUMNS full_text SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX idx_source_embed_chunk ON TABLE source_embedding COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX idx_source_insight ON TABLE source_insight COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX idx_note ON TABLE note COLUMNS content SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
DEFINE INDEX idx_note_title ON TABLE note COLUMNS title SEARCH ANALYZER my_analyzer BM25 HIGHLIGHTS;
-- 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
@ -150,8 +136,6 @@ DEFINE FUNCTION IF NOT EXISTS fn::text_search($query_text: string, $match_count:
};
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 =
@ -188,10 +172,7 @@ DEFINE FUNCTION IF NOT EXISTS fn::vector_search($query: array<float>, $match_cou
};
CREATE open_notebook:database_info SET
version= "0.0.2";
UPDATE open_notebook:database_info SET
version= "0.0.2";
IF array::len(select * from open_notebook:default_models) == 0 THEN
CREATE open_notebook:default_models SET
default_chat_model= ""
END;

24
migrations/1_down.surrealql Normal file
View file

@ -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;

21
open_notebook/config.py
View file

@ -3,7 +3,9 @@ import os
import yaml
from loguru import logger
# todo: enable config file overwrite with env vars
from open_notebook.domain.models import DefaultModels
from open_notebook.models import get_model
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
config_path = os.path.join(project_root, "open_notebook_config.yaml")
@ -32,3 +34,20 @@ os.makedirs(UPLOADS_FOLDER, exist_ok=True)
# PODCASTS FOLDER
PODCASTS_FOLDER = f"{DATA_FOLDER}/podcasts"
os.makedirs(PODCASTS_FOLDER, exist_ok=True)
DEFAULT_MODELS = DefaultModels.load()
if DEFAULT_MODELS.default_embedding_model:
EMBEDDING_MODEL = get_model(
DEFAULT_MODELS.default_embedding_model, model_type="embedding"
)
else:
EMBEDDING_MODEL = None
if DEFAULT_MODELS.default_speech_to_text_model:
SPEECH_TO_TEXT_MODEL = get_model(
DEFAULT_MODELS.default_speech_to_text_model, model_type="speech_to_text"
)
else:
SPEECH_TO_TEXT_MODEL = None

56
open_notebook/database/migrate.py Normal file
View file

@ -0,0 +1,56 @@
import os
from loguru import logger
from sblpy.connection import SurrealSyncConnection
from sblpy.migrations.db_processes import get_latest_version
from sblpy.migrations.migrations import Migration
from sblpy.migrations.runner import MigrationRunner
class MigrationManager:
def __init__(self):
self.connection = SurrealSyncConnection(
host=os.environ["SURREAL_ADDRESS"],
port=int(os.environ["SURREAL_PORT"]),
user=os.environ["SURREAL_USER"],
password=os.environ["SURREAL_PASS"],
namespace=os.environ["SURREAL_NAMESPACE"],
database=os.environ["SURREAL_DATABASE"],
encrypted=False, # Set to True if using SSL
)
self.up_migrations = [Migration.from_file("migrations/1.surrealql")]
self.down_migrations = [Migration.from_file("migrations/1_down.surrealql")]
self.runner = MigrationRunner(
up_migrations=self.up_migrations,
down_migrations=self.down_migrations,
connection=self.connection,
)
def get_current_version(self) -> int:
return get_latest_version(
self.connection.host,
self.connection.port,
self.connection.user,
self.connection.password,
self.connection.namespace,
self.connection.database,
)
@property
def needs_migration(self) -> bool:
current_version = self.get_current_version()
return current_version < len(self.up_migrations)
def run_migration_up(self):
current_version = self.get_current_version()
logger.debug(f"Current version before migration: {current_version}")
if self.needs_migration:
try:
self.runner.run()
new_version = self.get_current_version()
logger.debug(f"Migration successful. New version: {new_version}")
except Exception as e:
logger.error(f"Migration failed: {str(e)}")
else:
logger.debug("Database is already at the latest version")

32
open_notebook/repository.py → open_notebook/database/repository.py
View file

@ -5,10 +5,6 @@ from typing import Any, Dict, Optional
from loguru import logger
from sblpy.connection import SurrealSyncConnection
from open_notebook.exceptions import InvalidDatabaseSchema, NoSchemaFound
EXPECTED_VERSION = "0.0.2"
@contextmanager
def db_connection():
@ -34,30 +30,11 @@ def repo_query(query_str: str, vars: Optional[Dict[str, Any]] = None):
result = connection.query(query_str, vars)
return result
except Exception as e:
# logger.debug(f"Query: {query_str}, Variables: {vars}")
logger.critical(f"Query: {query_str}, Variables: {vars}")
logger.exception(e)
raise
def check_database_version():
try:
result = repo_query("SELECT * FROM open_notebook:database_info;")
if not result:
raise NoSchemaFound("Database schema not found")
version = result[0]["version"]
logger.info(f"Connected to SurrealDB, using schema version {version}")
if version != EXPECTED_VERSION:
raise InvalidDatabaseSchema(
f"Version mismatch. Expected {EXPECTED_VERSION}, got {version}"
)
except Exception as e:
logger.error(e)
raise e
def repo_create(table: str, data: Dict[str, Any]):
query = f"CREATE {table} CONTENT {data};"
# vars = {"table": table, "data": data}
@ -89,10 +66,3 @@ def repo_relate(source: str, relationship: str, target: str):
result = repo_query(query)
logger.debug(f"RELATE query result: {result}")
return result
def execute_migration(script: str):
with open(f"database/{script}", "r") as file:
content = file.read()
return repo_query(content)

0
open_notebook/domain/__init__.py Normal file
View file

147
open_notebook/domain/base.py Normal file
View file

@ -0,0 +1,147 @@
from datetime import datetime
from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar
from loguru import logger
from pydantic import BaseModel, ValidationError, field_validator
from open_notebook.database.repository import (
repo_create,
repo_delete,
repo_query,
repo_relate,
repo_update,
)
from open_notebook.exceptions import (
DatabaseOperationError,
InvalidInputError,
NotFoundError,
)
T = TypeVar("T", bound="ObjectModel")
class ObjectModel(BaseModel):
id: Optional[str] = None
table_name: ClassVar[str] = ""
created: Optional[datetime] = None
updated: Optional[datetime] = None
@classmethod
def get_all(cls: Type[T], order_by=None) -> List[T]:
try:
if order_by:
order = f" ORDER BY {order_by}"
else:
order = ""
result = repo_query(f"SELECT * FROM {cls.table_name} {order}")
objects = []
for obj in result:
try:
objects.append(cls(**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
def get(cls: Type[T], id: str) -> T:
if not id:
raise InvalidInputError("ID cannot be empty")
try:
result = repo_query(f"SELECT * FROM {id}")
if result:
return cls(**result[0])
return None
except Exception as e:
logger.error(f"Error fetching {cls.table_name} with id {id}: {str(e)}")
logger.exception(e)
raise NotFoundError(f"{cls.table_name} with id {id} not found")
def needs_embedding(self) -> bool:
return False
def get_embedding_content(self) -> Optional[str]:
return None
def save(self) -> None:
from open_notebook.config import EMBEDDING_MODEL
try:
logger.debug(f"Validating {self.__class__.__name__}")
self.model_validate(self.model_dump(), strict=True)
data = self._prepare_save_data()
data["updated"] = datetime.now().isoformat()
if self.needs_embedding():
embedding_content = self.get_embedding_content()
if embedding_content:
data["embedding"] = EMBEDDING_MODEL.embed(embedding_content)
if self.id is None:
data["created"] = datetime.now().isoformat()
logger.debug("Creating new record")
repo_result = repo_create(self.__class__.table_name, data)
else:
logger.debug(f"Updating record with id {self.id}")
repo_result = repo_update(self.id, data)
# Update the current instance with the result
for key, value in repo_result[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 Exception as e:
logger.error(f"Error saving record: {e}")
raise
except Exception as e:
logger.error(f"Error saving {self.__class__.table_name}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
def _prepare_save_data(self) -> Dict[str, Any]:
data = self.model_dump()
# del data["created"]
# del data["updated"]
return {key: value for key, value in data.items() if value is not None}
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 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}"
)
def relate(self, relationship: str, target_id: str) -> Any:
if not relationship or not target_id or not self.id:
raise InvalidInputError("Relationship and target ID must be provided")
try:
return repo_relate(self.id, relationship, target_id)
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

46
open_notebook/domain/models.py Normal file
View file

@ -0,0 +1,46 @@
from typing import ClassVar, Optional
from pydantic import BaseModel
from open_notebook.database.repository import (
repo_query,
repo_update,
)
from open_notebook.domain.base import ObjectModel
class Model(ObjectModel):
table_name: ClassVar[str] = "model"
name: str
provider: str
type: str
@classmethod
def get_models_by_type(cls, model_type):
models = repo_query(
"SELECT * FROM model WHERE type=$model_type;", {"model_type": model_type}
)
return [Model(**model) for model in models]
class DefaultModels(BaseModel):
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] = None
default_embedding_model: Optional[str] = None
@classmethod
def load(self):
result = repo_query("SELECT * FROM open_notebook:default_models;")
if result:
result = result[0]
dm = DefaultModels(**result)
return dm
return DefaultModels()
@classmethod
def update(self, data):
repo_update("open_notebook:default_models", data)

160
open_notebook/domain.py → open_notebook/domain/notebook.py
View file

@ -1,153 +1,23 @@
import os
from datetime import datetime
from typing import Any, ClassVar, Dict, List, Literal, Optional, Type, TypeVar
from typing import Any, ClassVar, Dict, List, Literal, Optional
from langchain_core.runnables.config import RunnableConfig
from loguru import logger
from pydantic import BaseModel, Field, ValidationError, field_validator
from pydantic import BaseModel, Field, field_validator
from open_notebook.config import EMBEDDING_MODEL
from open_notebook.database.repository import (
repo_create,
repo_query,
)
from open_notebook.domain.base import ObjectModel
from open_notebook.exceptions import (
DatabaseOperationError,
InvalidInputError,
NotFoundError,
)
from open_notebook.graphs.multipattern import graph as pattern_graph
from open_notebook.graphs.recursive_toc import graph as toc_graph
from open_notebook.repository import (
repo_create,
repo_delete,
repo_query,
repo_relate,
repo_update,
)
from open_notebook.utils import get_embedding, split_text, surreal_clean
T = TypeVar("T", bound="ObjectModel")
class ObjectModel(BaseModel):
id: Optional[str] = None
table_name: ClassVar[str] = ""
created: Optional[datetime] = None
updated: Optional[datetime] = None
@classmethod
def get_all(cls: Type[T], order_by=None) -> List[T]:
try:
if order_by:
order = f" ORDER BY {order_by}"
else:
order = ""
result = repo_query(f"SELECT * FROM {cls.table_name} {order}")
objects = []
for obj in result:
try:
objects.append(cls(**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
def get(cls: Type[T], id: str) -> Optional[T]:
if not id:
raise InvalidInputError("ID cannot be empty")
try:
result = repo_query(f"SELECT * FROM {id}")
if result:
return cls(**result[0])
return None
except Exception as e:
logger.error(f"Error fetching {cls.table_name} with id {id}: {str(e)}")
logger.exception(e)
raise NotFoundError(f"{cls.table_name} with id {id} not found")
def needs_embedding(self) -> bool:
return False
def get_embedding_content(self) -> Optional[str]:
return None
def save(self) -> None:
try:
logger.debug(f"Validating {self.__class__.__name__}")
self.model_validate(self.model_dump(), strict=True)
data = self._prepare_save_data()
data["updated"] = datetime.now().isoformat()
if self.needs_embedding():
embedding_content = self.get_embedding_content()
if embedding_content:
data["embedding"] = get_embedding(embedding_content)
if self.id is None:
data["created"] = datetime.now().isoformat()
logger.debug("Creating new record")
repo_result = repo_create(self.__class__.table_name, data)
else:
logger.debug(f"Updating record with id {self.id}")
repo_result = repo_update(self.id, data)
# Update the current instance with the result
for key, value in repo_result[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 Exception as e:
logger.error(f"Error saving record: {e}")
raise
except Exception as e:
logger.error(f"Error saving {self.__class__.table_name}: {str(e)}")
logger.exception(e)
raise DatabaseOperationError(e)
def _prepare_save_data(self) -> Dict[str, Any]:
data = self.model_dump()
# del data["created"]
# del data["updated"]
return {key: value for key, value in data.items() if value is not None}
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 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}"
)
def relate(self, relationship: str, target_id: str) -> Any:
if not relationship or not target_id or not self.id:
raise InvalidInputError("Relationship and target ID must be provided")
try:
return repo_relate(self.id, relationship, target_id)
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
from open_notebook.utils import split_text, surreal_clean
class Notebook(ObjectModel):
@ -288,7 +158,7 @@ class Source(ObjectModel):
"source": {self.id},
"order": {i},
"content": $content,
"embedding": {get_embedding(chunk)},
"embedding": {EMBEDDING_MODEL.embed(chunk)},
}};""",
{"content": surreal_clean(chunk)},
)
@ -322,7 +192,7 @@ class Source(ObjectModel):
if not insight_type or not content:
raise InvalidInputError("Insight type and content must be provided")
try:
embedding = get_embedding(content)
embedding = EMBEDDING_MODEL.embed(content)
return repo_query(
f"""
CREATE source_insight CONTENT {{
@ -396,9 +266,7 @@ class Note(ObjectModel):
return self.content
def text_search(
keyword: str, results: int, source: bool = True, note: bool = True
) -> List[Dict[str, Any]]:
def text_search(keyword: str, results: int, source: bool = True, note: bool = True):
if not keyword:
raise InvalidInputError("Search keyword cannot be empty")
try:
@ -415,9 +283,7 @@ def text_search(
raise DatabaseOperationError("Failed to perform text search")
def vector_search(
keyword: str, results: int, source: bool = True, note: bool = True
) -> List[Dict[str, Any]]:
def vector_search(keyword: str, results: int, source: bool = True, note: bool = True):
if not keyword:
raise InvalidInputError("Search keyword cannot be empty")
try:

12
open_notebook/exceptions.py
View file

@ -16,12 +16,6 @@ class UnsupportedTypeException(OpenNotebookError):
pass
class NoSchemaFound(OpenNotebookError):
"""Raised when a database schema is not found."""
pass
class InvalidInputError(OpenNotebookError):
"""Raised when invalid input is provided."""
@ -70,12 +64,6 @@ class NetworkError(OpenNotebookError):
pass
class InvalidDatabaseSchema(OpenNotebookError):
"""Raised when the database is not under the expected schema."""
pass
class NoTranscriptFound(OpenNotebookError):
"""Raised when no transcript is found for a video."""

8
open_notebook/graphs/chat.py
View file

@ -9,8 +9,8 @@ from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict
from open_notebook.config import LANGGRAPH_CHECKPOINT_FILE
from open_notebook.domain import Notebook
from open_notebook.config import DEFAULT_MODELS, LANGGRAPH_CHECKPOINT_FILE
from open_notebook.domain.notebook import Notebook
from open_notebook.graphs.utils import run_pattern
@ -22,7 +22,9 @@ class ThreadState(TypedDict):
def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict:
model_name = config.get("configurable", {}).get("model_name", None)
model_name = config.get("configurable", {}).get(
"model_name", DEFAULT_MODELS.default_chat_model
)
ai_message = run_pattern(
"chat",
model_name,

11
open_notebook/graphs/content_processing/audio.py
View file

@ -4,9 +4,9 @@ from math import ceil
from loguru import logger
from pydub import AudioSegment
from open_notebook.config import SPEECH_TO_TEXT_MODEL
from open_notebook.graphs.content_processing.state import SourceState
# todo: add a speechtotext model to the config
# future: parallelize the transcription process
@ -73,9 +73,6 @@ def split_audio(input_file, segment_length_minutes=15, output_prefix=None):
def extract_audio(data: SourceState):
input_audio_path = data.get("file_path")
from openai import OpenAI
client = OpenAI()
audio_files = []
try:
@ -83,11 +80,7 @@ def extract_audio(data: SourceState):
transcriptions = []
for audio_file in audio_files:
with open(audio_file, "rb") as audio:
transcription = client.audio.transcriptions.create(
model="whisper-1", file=audio
)
transcriptions.append(transcription.text)
transcriptions.append(SPEECH_TO_TEXT_MODEL.transcribe(audio_file))
return {"content": " ".join(transcriptions)}

2
open_notebook/graphs/doc_query.py
View file

@ -6,7 +6,7 @@ from langchain_core.runnables import (
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict
from open_notebook.domain import Note, Notebook, Source
from open_notebook.domain.notebook import Note, Notebook, Source
from open_notebook.graphs.utils import run_pattern

4
open_notebook/graphs/multipattern.py
View file

@ -1,5 +1,4 @@
import operator
import os
from typing import List, Literal, Sequence
from langchain_core.runnables import (
@ -8,6 +7,7 @@ from langchain_core.runnables import (
from langgraph.graph import END, START, StateGraph
from typing_extensions import Annotated, TypedDict
from open_notebook.config import DEFAULT_MODELS
from open_notebook.graphs.utils import run_pattern
@ -19,7 +19,7 @@ class PatternChainState(TypedDict):
def call_model(state: dict, config: RunnableConfig) -> dict:
model_name = config.get("configurable", {}).get(
"model_name", os.environ.get("DEFAULT_MODEL")
"model_name", DEFAULT_MODELS.default_transformation_model
)
transformations = state["transformations"]
current_transformation = transformations.pop(0)

5
open_notebook/graphs/pattern.py
View file

@ -1,11 +1,10 @@
import os
from langchain_core.runnables import (
RunnableConfig,
)
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict
from open_notebook.config import DEFAULT_MODELS
from open_notebook.graphs.utils import run_pattern
@ -17,7 +16,7 @@ class PatternState(TypedDict):
def call_model(state: dict, config: RunnableConfig) -> dict:
model_name = config.get("configurable", {}).get(
"model_name", os.environ.get("DEFAULT_MODEL")
"model_name", DEFAULT_MODELS.default_transformation_model
)
return {
"output": run_pattern(

3
open_notebook/graphs/recursive_toc.py
View file

@ -7,6 +7,7 @@ from langchain_core.runnables import (
from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict
from open_notebook.config import DEFAULT_MODELS
from open_notebook.graphs.utils import run_pattern
from open_notebook.utils import split_text
@ -49,7 +50,7 @@ def chunk_condition(state: TocState) -> Literal["get_chunk", END]: # type: igno
def call_model(state: TocState, config: RunnableConfig) -> dict:
model_name = config.get("configurable", {}).get(
"model_name", os.environ.get("SUMMARIZATION_MODEL")
"model_name", DEFAULT_MODELS.default_transformation_model
)
return {
"toc": run_pattern(

5
open_notebook/graphs/summary.py
View file

@ -9,6 +9,7 @@ from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel
from typing_extensions import TypedDict
from open_notebook.config import DEFAULT_MODELS
from open_notebook.graphs.utils import run_pattern
from open_notebook.utils import split_text
@ -57,9 +58,9 @@ def chunk_condition(state: SummaryState) -> Literal["get_chunk", END]: # type:
return END
def call_model(state: SummaryState, config: RunnableConfig) -> dict:
def call_model(state: dict, config: RunnableConfig) -> dict:
model_name = config.get("configurable", {}).get(
"model_name", os.environ.get("SUMMARIZATION_MODEL")
"model_name", DEFAULT_MODELS.default_transformation_model
)
parser = PydanticOutputParser(pydantic_object=SummaryResponse)
return {

34
open_notebook/graphs/utils.py
View file

@ -1,9 +1,10 @@
import os
from langchain.output_parsers import OutputFixingParser
from loguru import logger
from open_notebook.llm_router import get_langchain_model
from open_notebook.config import DEFAULT_MODELS
from open_notebook.models import get_model
from open_notebook.prompter import Prompter
from open_notebook.utils import token_count
def run_pattern(
@ -14,24 +15,35 @@ def run_pattern(
parser=None,
output_fixing_model_name=None,
) -> dict:
if not model_name:
model_name = os.environ["DEFAULT_MODEL"]
system_prompt = Prompter(prompt_template=pattern_name, parser=parser).render(
data=state
)
chain = get_langchain_model(model_name)
tokens = token_count(str(system_prompt) + str(messages))
if tokens > 105_000 and DEFAULT_MODELS.large_context_model:
model_name = DEFAULT_MODELS.large_context_model
logger.debug(
f"Using large context model ({model_name}) because the content has {tokens} tokens"
)
logger.warning(system_prompt)
elif tokens > 105_000 and not DEFAULT_MODELS.large_context_model:
logger.critical(
f"Content has {tokens} tokens, but no large context model is configured"
)
elif not model_name:
model_name = DEFAULT_MODELS.default_transformation_model
chain = get_model(model_name, model_type="language")
if parser:
chain = chain | parser
if output_fixing_model_name and parser:
output_fix_model = get_langchain_model(output_fixing_model_name)
output_fix_model = get_model(output_fixing_model_name, model_type="language")
chain = chain | OutputFixingParser.from_llm(
parser=parser,
llm=output_fix_model,
)
system_prompt = Prompter(prompt_template=pattern_name, parser=parser).render(
data=state
)
if len(messages) > 0:
response = chain.invoke([system_prompt] + messages)
else:

35
open_notebook/llm_router.py
View file

@ -1,35 +0,0 @@
from open_notebook.llms import (
AnthropicLanguageModel,
GeminiLanguageModel,
LiteLLMLanguageModel,
OllamaLanguageModel,
OpenAILanguageModel,
OpenRouterLanguageModel,
VertexAILanguageModel,
VertexAnthropicLanguageModel,
)
# Map provider names to classes
PROVIDER_CLASS_MAP = {
"ollama": OllamaLanguageModel,
"openrouter": OpenRouterLanguageModel,
"vertexai-anthropic": VertexAnthropicLanguageModel,
"litellm": LiteLLMLanguageModel,
"vertexai": VertexAILanguageModel,
"anthropic": AnthropicLanguageModel,
"openai": OpenAILanguageModel,
"gemini": GeminiLanguageModel,
}
def get_langchain_model(model_name, json=False):
parts = model_name.split("/")
provider = parts[0]
model_name_wihout_provider = "/".join(parts[1:])
if provider not in PROVIDER_CLASS_MAP.keys():
raise ValueError(
f"Provider {provider} not found in config. Make sure you use the correct format for model names, example: openai/gpt-4o-mini"
)
return PROVIDER_CLASS_MAP[provider](
model_name=model_name_wihout_provider, json=json
).to_langchain()

83
open_notebook/models/__init__.py Normal file
View file

@ -0,0 +1,83 @@
from open_notebook.domain.models import Model
from open_notebook.models.embedding_models import (
GeminiEmbeddingModel,
OllamaEmbeddingModel,
OpenAIEmbeddingModel,
VertexEmbeddingModel,
)
from open_notebook.models.llms import (
AnthropicLanguageModel,
GeminiLanguageModel,
LiteLLMLanguageModel,
OllamaLanguageModel,
OpenAILanguageModel,
OpenRouterLanguageModel,
VertexAILanguageModel,
VertexAnthropicLanguageModel,
)
from open_notebook.models.speech_to_text_models import OpenAISpeechToTextModel
from open_notebook.models.text_to_speech_models import (
ElevenLabsTextToSpeechModel,
OpenAITextToSpeechModel,
)
# Unified model class map with type information
MODEL_CLASS_MAP = {
"language": {
"ollama": OllamaLanguageModel,
"openrouter": OpenRouterLanguageModel,
"vertexai-anthropic": VertexAnthropicLanguageModel,
"litellm": LiteLLMLanguageModel,
"vertexai": VertexAILanguageModel,
"anthropic": AnthropicLanguageModel,
"openai": OpenAILanguageModel,
"gemini": GeminiLanguageModel,
},
"embedding": {
"openai": OpenAIEmbeddingModel,
"gemini": GeminiEmbeddingModel,
"vertexai": VertexEmbeddingModel,
"ollama": OllamaEmbeddingModel,
},
"speech_to_text": {
"openai": OpenAISpeechToTextModel,
},
"text_to_speech": {
"openai": OpenAITextToSpeechModel,
"elevenlabs": ElevenLabsTextToSpeechModel,
},
}
def get_model(model_id, model_type="language", **kwargs):
"""
Get a model instance based on model_id and type.
Args:
model_id: The ID of the model to retrieve
model_type: Type of model ('language', 'embedding', or 'speech_to_text')
**kwargs: Additional arguments to pass to the model constructor
"""
assert model_id, "Model ID cannot be empty"
model: Model = Model.get(model_id)
if not model:
raise ValueError(f"Model with ID {model_id} not found")
if model_type not in MODEL_CLASS_MAP:
raise ValueError(f"Invalid model type: {model_type}")
provider_map = MODEL_CLASS_MAP[model_type]
if model.provider not in provider_map:
raise ValueError(
f"Provider {model.provider} not compatible with {model_type} models"
)
model_class = provider_map[model.provider]
model_instance = model_class(model_name=model.name, **kwargs)
# Special handling for language models that need langchain conversion
if model_type == "language":
return model_instance.to_langchain()
return model_instance

104
open_notebook/models/embedding_models.py Normal file
View file

@ -0,0 +1,104 @@
"""
Classes for supporting different embedding models
"""
from __future__ import annotations
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Optional
import requests
# todo: add support for multiple embeddings (array)
@dataclass
class EmbeddingModel(ABC):
"""
Abstract base class for language models.
"""
model_name: Optional[str] = None
@abstractmethod
def embed(self, text: str) -> List[float]:
"""
Generates an embedding
"""
raise NotImplementedError
@dataclass
class OllamaEmbeddingModel(EmbeddingModel):
model_name: str
base_url: str = os.environ.get("OLLAMA_API_BASE", "http://localhost:11434")
def embed(self, text: str) -> List[float]:
"""
Embeds the content using Open AI embedding
"""
text = text.replace("\n", " ")
response = requests.post(
f"{self.base_url}/api/embed",
json={"model": self.model_name, "input": [text]},
)
return response.json()["embeddings"][0]
@dataclass
class GeminiEmbeddingModel(EmbeddingModel):
model_name: str
def embed(self, text: str) -> List[float]:
import google.generativeai as genai
"""
Embeds the content using Open AI embedding
"""
model_name = (
self.model_name
if self.model_name.startswith("models/")
else f"models/{self.model_name}"
)
result = genai.embed_content(model=model_name, content=text)
return result["embedding"]
@dataclass
class VertexEmbeddingModel(EmbeddingModel):
model_name: str
def embed(self, text: str) -> List[float]:
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
texts = [text]
# The dimensionality of the output embeddings.
# dimensionality = 256
# The task type for embedding. Check the available tasks in the model's documentation.
model = TextEmbeddingModel.from_pretrained(self.model_name)
inputs = [TextEmbeddingInput(text) for text in texts]
embeddings = model.get_embeddings(inputs)
return embeddings[0].values
@dataclass
class OpenAIEmbeddingModel(EmbeddingModel):
model_name: str
def embed(self, text: str) -> List[float]:
from openai import OpenAI
"""
Embeds the content using Open AI embedding
"""
# todo: make this Singleton
client = OpenAI()
text = text.replace("\n", " ")
return (
client.embeddings.create(input=[text], model=self.model_name)
.data[0]
.embedding
)

33
open_notebook/llms.py → open_notebook/models/llms.py
View file

@ -1,5 +1,5 @@
"""
Classes for supporting different language and vector models
Classes for supporting different language models
"""
import os
@ -15,9 +15,9 @@ from langchain_google_vertexai import ChatVertexAI
from langchain_google_vertexai.model_garden import ChatAnthropicVertex
from langchain_ollama.chat_models import ChatOllama
from langchain_openai.chat_models import ChatOpenAI
from pydantic import SecretStr
# from redisvl.utils.vectorize import BaseVectorizer
# from redisvl.utils.vectorize.text.openai import OpenAITextVectorizer
# future: is there a value on returning langchain specific models?
@dataclass
@ -186,7 +186,7 @@ class OpenRouterLanguageModel(LanguageModel):
max_tokens=self.max_tokens,
model_kwargs=kwargs,
streaming=self.streaming,
api_key=os.environ.get("OPENROUTER_API_KEY", "openrouter"),
api_key=SecretStr(os.environ.get("OPENROUTER_API_KEY", "openrouter")),
top_p=self.top_p,
)
@ -238,28 +238,3 @@ class OpenAILanguageModel(LanguageModel):
streaming=self.streaming,
top_p=self.top_p,
)
# @dataclass
# class EmbeddingModel(ABC):
# model_name: str
# dimensions: int
# def to_redis_vectorizer(self) -> BaseVectorizer:
# raise NotImplementedError
# @dataclass
# class OpenAIEmbeddingModel(EmbeddingModel):
# """
# Embedding model that uses the OpenAI text embedding model.
# """
# model_name: str
# dimensions: int
# def to_redis_vectorizer(self) -> OpenAITextVectorizer:
# """
# Convert the embedding model to a Redis vectorizer.
# """
# return OpenAITextVectorizer(model=self.model_name)

42
open_notebook/models/speech_to_text_models.py Normal file
View file

@ -0,0 +1,42 @@
"""
Classes for supporting different transcription models
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional
@dataclass
class SpeechToTextModel(ABC):
"""
Abstract base class for speech to text models.
"""
model_name: Optional[str] = None
@abstractmethod
def transcribe(self, audio_file_path: str) -> str:
"""
Generates a text transcription from audio
"""
raise NotImplementedError
@dataclass
class OpenAISpeechToTextModel(SpeechToTextModel):
model_name: str
def transcribe(self, audio_file_path: str) -> str:
"""
Transcribes an audio file into text
"""
from openai import OpenAI
# todo: make this Singleton
client = OpenAI()
with open(audio_file_path, "rb") as audio:
transcription = client.audio.transcriptions.create(
model=self.model_name, file=audio
)
return transcription.text

26
open_notebook/models/text_to_speech_models.py Normal file
View file

@ -0,0 +1,26 @@
"""
Classes for supporting different text to speech models
"""
from abc import ABC
from dataclasses import dataclass
from typing import Optional
@dataclass
class TextToSpeechModel(ABC):
"""
Abstract base class for text to speech models.
"""
model_name: Optional[str] = None
@dataclass
class OpenAITextToSpeechModel(TextToSpeechModel):
model_name: str
@dataclass
class ElevenLabsTextToSpeechModel(TextToSpeechModel):
model_name: str

6
open_notebook/plugins/podcasts.py
View file

@ -1,10 +1,10 @@
from typing import ClassVar, List, Literal, Optional
from typing import ClassVar, List, Optional
from loguru import logger
from podcastfy.client import generate_podcast
from pydantic import Field, field_validator
from open_notebook.domain import ObjectModel
from open_notebook.domain.notebook import ObjectModel
class PodcastEpisode(ObjectModel):
@ -31,7 +31,7 @@ class PodcastConfig(ObjectModel):
ending_message: Optional[str] = None
wordcount: int = Field(ge=400, le=10000)
creativity: float = Field(ge=0, le=1)
provider: Literal["openai", "elevenlabs", "edge"] = Field(default="openai")
provider: str = Field(default="openai")
voice1: Optional[str] = None
voice2: Optional[str] = None
model: str

18
open_notebook/utils.py
View file

@ -6,11 +6,8 @@ from urllib.parse import urlparse
import requests
import tomli
from langchain_text_splitters import CharacterTextSplitter
from openai import OpenAI
from packaging.version import parse as parse_version
client = OpenAI()
def split_text(txt: str, chunk=1000, overlap=0, separator=" "):
"""
@ -63,21 +60,6 @@ def token_cost(token_count, cost_per_million=0.150):
return cost_per_million * (token_count / 1_000_000)
def get_embedding(text, model="text-embedding-3-small"):
"""
Get the embedding for the input text using the specified model.
Args:
text (str): The input text to get the embedding for.
model (str): The name of the embedding model to use. Default is "text-embedding-3-small".
Returns:
list: The embedding vector for the input text.
"""
text = text.replace("\n", " ")
return client.embeddings.create(input=[text], model=model).data[0].embedding
def remove_non_ascii(text):
return re.sub(r"[^\x00-\x7F]+", "", text)

2
pages/2_📒_Notebooks.py
View file

@ -1,7 +1,7 @@
import streamlit as st
from humanize import naturaltime
from open_notebook.domain import Notebook
from open_notebook.domain.notebook import Notebook
from stream_app.chat import chat_sidebar
from stream_app.note import add_note, note_card
from stream_app.source import add_source, source_card

6
pages/3_🔍_Search.py
View file

@ -1,7 +1,7 @@
import streamlit as st
from open_notebook.domain import text_search, vector_search
from open_notebook.utils import get_embedding
from open_notebook.config import EMBEDDING_MODEL
from open_notebook.domain.notebook import text_search, vector_search
from stream_app.note import note_list_item
from stream_app.source import source_list_item
from stream_app.utils import version_sidebar
@ -33,7 +33,7 @@ with st.container(border=True):
)
elif search_type == "Vector Search":
st.write(f"Searching for {search_term}")
embed_query = get_embedding(search_term)
embed_query = EMBEDDING_MODEL.embed(search_term)
st.session_state["search_results"] = vector_search(
embed_query, 100, search_sources, search_notes
)

34
pages/5_🎙️_Podcasts.py
View file

@ -1,6 +1,9 @@
from typing import Dict, List
import streamlit as st
from streamlit_tags import st_tags
from open_notebook.domain.models import Model
from open_notebook.plugins.podcasts import (
PodcastConfig,
PodcastEpisode,
@ -17,6 +20,21 @@ st.set_page_config(
version_sidebar()
text_to_speech_models = Model.get_models_by_type("text_to_speech")
provider_models: Dict[str, List[str]] = {}
for model in text_to_speech_models:
if model.provider not in provider_models:
provider_models[model.provider] = []
provider_models[model.provider].append(model.name)
if len(text_to_speech_models) == 0:
st.error("No text to speech models found. Please set one up in the Settings page.")
st.stop()
episodes_tab, templates_tab = st.tabs(["Episodes", "Templates"])
with episodes_tab:
@ -76,7 +94,7 @@ with templates_tab:
pd_cfg["ending_message"] = st.text_input(
"Ending Message", placeholder="Thank you for listening!"
)
pd_cfg["provider"] = st.selectbox("Provider", ["openai", "elevenlabs", "edge"])
pd_cfg["provider"] = st.selectbox("Provider", provider_models.keys())
pd_cfg["voice1"] = st.text_input(
"Voice 1", help="You can use Elevenlabs voice ID"
)
@ -86,7 +104,8 @@ with templates_tab:
pd_cfg["voice2"] = st.text_input(
"Voice 2", help="You can use Elevenlabs voice ID"
)
pd_cfg["model"] = st.text_input("Model")
pd_cfg["model"] = st.selectbox("Model", provider_models[pd_cfg["provider"]])
st.caption(
"OpenAI: tts-1 or tts-1-hd, Elevenlabs: eleven_multilingual_v2, eleven_turbo_v2_5"
)
@ -183,8 +202,8 @@ with templates_tab:
)
pd_config.provider = st.selectbox(
"Provider",
["openai", "elevenlabs", "edge"],
index=["openai", "elevenlabs", "edge"].index(pd_config.provider),
list(provider_models.keys()),
index=list(provider_models.keys()).index(pd_config.provider),
key=f"provider_{pd_config.id}",
)
pd_config.voice1 = st.text_input(
@ -202,8 +221,11 @@ with templates_tab:
key=f"voice2_{pd_config.id}",
help="You can use Elevenlabs voice ID",
)
pd_config.model = st.text_input(
"Model", value=pd_config.model, key=f"model_{pd_config.id}"
pd_config.model = st.selectbox(
"Model",
provider_models[pd_config.provider],
index=provider_models[pd_config.provider].index(pd_config.model),
key=f"model_{pd_config.id}",
)
st.caption(
"OpenAI: tts-1 or tts-1-hd, Elevenlabs: eleven_multilingual_v2, eleven_turbo_v2_5"

222
pages/9_⚙️_Settings.py Normal file
View file

@ -0,0 +1,222 @@
import os
import streamlit as st
from open_notebook.domain.models import DefaultModels, Model
from open_notebook.models import MODEL_CLASS_MAP
from stream_app.utils import version_sidebar
st.set_page_config(
layout="wide", page_title="⚙️ Settings", initial_sidebar_state="expanded"
)
version_sidebar()
st.title("Settings")
model_tab, model_defaults_tab = st.tabs(["Models", "Model Defaults"])
provider_status = {}
model_types = [
# "vision",
"language",
"embedding",
"text_to_speech",
"speech_to_text",
]
provider_status["ollama"] = os.environ.get("OLLAMA_API_BASE") is not None
provider_status["openai"] = os.environ.get("OPENAI_API_KEY") is not None
provider_status["vertexai"] = (
os.environ.get("VERTEX_PROJECT") is not None
and os.environ.get("VERTEX_LOCATION") is not None
and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is not None
)
provider_status["vertexai-anthropic"] = (
os.environ.get("VERTEX_PROJECT") is not None
and os.environ.get("VERTEX_LOCATION") is not None
and os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is not None
)
provider_status["gemini"] = os.environ.get("GEMINI_API_KEY") is not None
provider_status["openrouter"] = (
os.environ.get("OPENROUTER_API_KEY") is not None
and os.environ.get("OPENAI_API_KEY") is not None
and os.environ.get("OPENROUTER_BASE_URL") is not None
)
provider_status["anthropic"] = os.environ.get("ANTHROPIC_API_KEY") is not None
provider_status["elevenlabs"] = os.environ.get("ELEVENLABS_API_KEY") is not None
provider_status["litellm"] = (
provider_status["ollama"]
or provider_status["vertexai"]
or provider_status["vertexai-anthropic"]
or provider_status["anthropic"]
or provider_status["openai"]
or provider_status["gemini"]
)
available_providers = [k for k, v in provider_status.items() if v]
unavailable_providers = [k for k, v in provider_status.items() if not v]
with model_tab:
st.subheader("Add Model")
provider = st.selectbox("Provider", available_providers)
if len(unavailable_providers) > 0:
st.caption(
f"Unavailable Providers: {', '.join(unavailable_providers)}. Please check docs page if you wish to enable them."
)
# Filter model types based on provider availability in MODEL_CLASS_MAP
available_model_types = []
for model_type in model_types:
if model_type in MODEL_CLASS_MAP and provider in MODEL_CLASS_MAP[model_type]:
available_model_types.append(model_type)
if not available_model_types:
st.error(f"No compatible model types available for provider: {provider}")
else:
model_type = st.selectbox(
"Model Type",
available_model_types,
help="Use language for text generation models, text_to_speech for TTS models for generating podcasts, etc.",
)
model_name = st.text_input(
"Model Name", "", help="gpt-4o-mini, claude, gemini, llama3, etc"
)
if st.button("Save"):
model = Model(name=model_name, provider=provider, type=model_type)
model.save()
st.success("Saved")
st.divider()
all_models = Model.get_all()
st.subheader("Configured Models")
model_types_available = {
# "vision": False,
"language": False,
"embedding": False,
"text_to_speech": False,
"speech_to_text": False,
}
for model in all_models:
model_types_available[model.type] = True
with st.container(border=True):
st.markdown(f"{model.name} ({model.provider}, {model.type})")
if st.button("Delete", key=f"delete_{model.id}"):
model.delete()
st.rerun()
for model_type, available in model_types_available.items():
if not available:
st.warning(f"No models available for {model_type}")
def get_selected_index(models, model_id, default=0):
"""Returns the index of the selected model in the list of models"""
if not model_id or not models:
return default
for i, model in enumerate(models):
if model.id == model_id:
return i
return default
with model_defaults_tab:
default_models = DefaultModels.load().model_dump()
all_models = Model.get_all()
text_generation_models = [model for model in all_models if model.type == "language"]
text_to_speech_models = [
model for model in all_models if model.type == "text_to_speech"
]
speech_to_text_models = [
model for model in all_models if model.type == "speech_to_text"
]
vision_models = [model for model in all_models if model.type == "vision"]
embedding_models = [model for model in all_models if model.type == "embedding"]
st.write(
"In this section, you can select the default models to be used on the various content operations done by Open Notebook. Some of these can be overriden in the different modules."
)
defs = {}
defs["default_chat_model"] = st.selectbox(
"Default Chat Model",
text_generation_models,
format_func=lambda x: x.name,
help="This model will be used for chat.",
index=get_selected_index(
text_generation_models, default_models.get("default_chat_model")
),
)
st.divider()
defs["default_transformation_model"] = st.selectbox(
"Default Transformation Model",
text_generation_models,
format_func=lambda x: x.name,
help="This model will be used for text transformations such as summaries, insights, etc.",
index=get_selected_index(
text_generation_models, default_models.get("default_transformation_model")
),
)
st.caption("You can override this model on individual transformations")
st.divider()
defs["large_context_model"] = st.selectbox(
"Large Context Model",
text_generation_models,
format_func=lambda x: x.name,
help="This model will be used for larger context generation -- recommended: Gemini",
index=get_selected_index(
text_generation_models, default_models.get("large_context_model")
),
)
st.caption("Recommended to use Gemini models for larger context processing")
st.divider()
defs["default_text_to_speech_model"] = st.selectbox(
"Default Text to Speech Model",
text_to_speech_models,
format_func=lambda x: x.name,
help="This is the default model for converting text to speech (podcasts, etc)",
index=get_selected_index(
text_to_speech_models, default_models.get("default_text_to_speech_model")
),
)
st.caption("You can override this model on different podcasts")
st.divider()
defs["default_speech_to_text_model"] = st.selectbox(
"Default Speech to Text Model",
speech_to_text_models,
format_func=lambda x: x.name,
help="This is the default model for converting speech to text (audio transcriptions, etc)",
index=get_selected_index(
speech_to_text_models, default_models.get("default_speech_to_text_model")
),
)
st.divider()
# defs["default_vision_model"] = st.selectbox(
# "Default Vision Model",
# vision_models,
# format_func=lambda x: x.name,
# help="This is the default model for vision tasks (image recognition, PDF recognition, etc)",
# index=get_selected_index(
# vision_models, default_models.get("default_vision_model")
# ),
# )
# st.divider()
defs["default_embedding_model"] = st.selectbox(
"Default Embedding Model",
embedding_models,
format_func=lambda x: x.name,
help="This is the default model for embeddings (semantic search, etc)",
index=get_selected_index(
embedding_models, default_models.get("default_embedding_model")
),
)
st.caption(
"Caution: you cannot change the embedding model once there is embeddings or they will need to be regenerated"
)
# if st.button("Save Defaults", key="save_defaults"):
for k, v in defs.items():
if v:
defs[k] = v.id
DefaultModels.update(defs)

45
poetry.lock generated
View file

@ -1932,6 +1932,27 @@ qtconsole = ["qtconsole"]
test = ["packaging", "pickleshare", "pytest", "pytest-asyncio (<0.22)", "testpath"]
test-extra = ["curio", "ipython[test]", "matplotlib (!=3.2.0)", "nbformat", "numpy (>=1.23)", "pandas", "trio"]
[[package]]
name = "ipywidgets"
version = "8.1.5"
description = "Jupyter interactive widgets"
optional = false
python-versions = ">=3.7"
files = [
{file = "ipywidgets-8.1.5-py3-none-any.whl", hash = "sha256:3290f526f87ae6e77655555baba4f36681c555b8bdbbff430b70e52c34c86245"},
{file = "ipywidgets-8.1.5.tar.gz", hash = "sha256:870e43b1a35656a80c18c9503bbf2d16802db1cb487eec6fab27d683381dde17"},
]
[package.dependencies]
comm = ">=0.1.3"
ipython = ">=6.1.0"
jupyterlab-widgets = ">=3.0.12,<3.1.0"
traitlets = ">=4.3.1"
widgetsnbextension = ">=4.0.12,<4.1.0"
[package.extras]
test = ["ipykernel", "jsonschema", "pytest (>=3.6.0)", "pytest-cov", "pytz"]
[[package]]
name = "jedi"
version = "0.19.1"
@ -2163,6 +2184,17 @@ files = [
{file = "jupyterlab_pygments-0.3.0.tar.gz", hash = "sha256:721aca4d9029252b11cfa9d185e5b5af4d54772bb8072f9b7036f4170054d35d"},
]
[[package]]
name = "jupyterlab-widgets"
version = "3.0.13"
description = "Jupyter interactive widgets for JupyterLab"
optional = false
python-versions = ">=3.7"
files = [
{file = "jupyterlab_widgets-3.0.13-py3-none-any.whl", hash = "sha256:e3cda2c233ce144192f1e29914ad522b2f4c40e77214b0cc97377ca3d323db54"},
{file = "jupyterlab_widgets-3.0.13.tar.gz", hash = "sha256:a2966d385328c1942b683a8cd96b89b8dd82c8b8f81dda902bb2bc06d46f5bed"},
]
[[package]]
name = "langchain"
version = "0.3.4"
@ -6167,6 +6199,17 @@ files = [
[package.extras]
test = ["pytest (>=6.0.0)", "setuptools (>=65)"]
[[package]]
name = "widgetsnbextension"
version = "4.0.13"
description = "Jupyter interactive widgets for Jupyter Notebook"
optional = false
python-versions = ">=3.7"
files = [
{file = "widgetsnbextension-4.0.13-py3-none-any.whl", hash = "sha256:74b2692e8500525cc38c2b877236ba51d34541e6385eeed5aec15a70f88a6c71"},
{file = "widgetsnbextension-4.0.13.tar.gz", hash = "sha256:ffcb67bc9febd10234a362795f643927f4e0c05d9342c727b65d2384f8feacb6"},
]
[[package]]
name = "win32-setctime"
version = "1.1.0"
@ -6324,4 +6367,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.0"
python-versions = "^3.11"
content-hash = "4fa191c6df5a7a355eb0d61f9560ec70e4671ac49cd54fa3a166c1e25c325671"
content-hash = "265ed7b26b19c54847b8e549f09ccbf8be68120b34f392fb5b8afc9ffccd62ac"

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "open-notebook"
version = "0.0.6"
version = "0.0.7"
description = "An open source implementation of a research assistant, inspired by Google Notebook LM"
authors = ["Luis Novo <lfnovo@gmail.com>"]
license = "MIT"
@ -46,12 +46,14 @@ bs4 = "^0.0.2"
python-docx = "^1.1.2"
python-pptx = "^1.0.2"
openpyxl = "^3.1.5"
google-generativeai = "^0.8.3"
[tool.poetry.group.dev.dependencies]
ipykernel = "^6.29.5"
ruff = "^0.5.5"
mypy = "^1.11.1"
types-requests = "^2.32.0.20241016"
ipywidgets = "^8.1.5"
[build-system]
requires = ["poetry-core"]

4
stream_app/chat.py
View file

@ -1,7 +1,7 @@
import streamlit as st
from langchain_core.runnables import RunnableConfig
from open_notebook.domain import Note, Source
from open_notebook.domain.notebook import Note, Source
from open_notebook.graphs.chat import graph as chat_graph
from open_notebook.plugins.podcasts import PodcastConfig
from open_notebook.utils import token_count
@ -54,8 +54,6 @@ def execute_chat(txt_input, session_id):
return result
# todo: se eu for usar o token count, preciso deixar configuravel
# seria bom ter um total de tokens no admin em algum lugar
def chat_sidebar(session_id):
context = build_context(session_id=session_id)
tokens = token_count(str(context) + str(st.session_state[session_id]["messages"]))

2
stream_app/note.py
View file

@ -3,7 +3,7 @@ from humanize import naturaltime
from loguru import logger
from streamlit_monaco import st_monaco # type: ignore
from open_notebook.domain import Note
from open_notebook.domain.notebook import Note
from open_notebook.graphs.multipattern import graph as pattern_graph
from open_notebook.utils import surreal_clean

2
stream_app/source.py
View file

@ -8,7 +8,7 @@ from humanize import naturaltime
from loguru import logger
from open_notebook.config import UPLOADS_FOLDER
from open_notebook.domain import Asset, Source
from open_notebook.domain.notebook import Asset, Source
from open_notebook.exceptions import UnsupportedTypeException
from open_notebook.graphs.content_processing import graph
from open_notebook.graphs.multipattern import graph as transform_graph