From 0a868fceb5abfbbf1769b5619dee9b25463d4071 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:06:41 -0300
Subject: [PATCH 01/14] remove model info from docs

---
 .env.example  | 15 ---------------
 docs/SETUP.md |  1 -
 2 files changed, 16 deletions(-)

diff --git a/.env.example b/.env.example
index 7a4a27b..a964062 100644
--- a/.env.example
+++ b/.env.example
@@ -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=
 
diff --git a/docs/SETUP.md b/docs/SETUP.md
index c459caa..52753da 100644
--- a/docs/SETUP.md
+++ b/docs/SETUP.md
@@ -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

From 63a568490e6a221da310cbe675262b699b61f5a0 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:08:05 -0300
Subject: [PATCH 02/14] refactor domain into a module

---
 open_notebook/domain/__init__.py              |   0
 open_notebook/domain/base.py                  | 147 +++++++++++++++++
 .../{domain.py => domain/notebook.py}         | 152 ++----------------
 open_notebook/repository.py                   |   2 +-
 pages/2_📒_Notebooks.py                       |   2 +-
 pages/3_🔍_Search.py                          |   2 +-
 stream_app/chat.py                            |   2 +-
 stream_app/note.py                            |   2 +-
 8 files changed, 161 insertions(+), 148 deletions(-)
 create mode 100644 open_notebook/domain/__init__.py
 create mode 100644 open_notebook/domain/base.py
 rename open_notebook/{domain.py => domain/notebook.py} (65%)

diff --git a/open_notebook/domain/__init__.py b/open_notebook/domain/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py
new file mode 100644
index 0000000..da9b0ac
--- /dev/null
+++ b/open_notebook/domain/base.py
@@ -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.exceptions import (
+    DatabaseOperationError,
+    InvalidInputError,
+    NotFoundError,
+)
+from open_notebook.repository import (
+    repo_create,
+    repo_delete,
+    repo_query,
+    repo_relate,
+    repo_update,
+)
+
+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:
+        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
diff --git a/open_notebook/domain.py b/open_notebook/domain/notebook.py
similarity index 65%
rename from open_notebook/domain.py
rename to open_notebook/domain/notebook.py
index ef238eb..b917548 100644
--- a/open_notebook/domain.py
+++ b/open_notebook/domain/notebook.py
@@ -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.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:
diff --git a/open_notebook/repository.py b/open_notebook/repository.py
index 2375387..ca6f47a 100644
--- a/open_notebook/repository.py
+++ b/open_notebook/repository.py
@@ -34,7 +34,7 @@ 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
 
diff --git a/pages/2_📒_Notebooks.py b/pages/2_📒_Notebooks.py
index 7d39d08..b6e5701 100644
--- a/pages/2_📒_Notebooks.py
+++ b/pages/2_📒_Notebooks.py
@@ -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
diff --git a/pages/3_🔍_Search.py b/pages/3_🔍_Search.py
index bd10233..6c69ea2 100644
--- a/pages/3_🔍_Search.py
+++ b/pages/3_🔍_Search.py
@@ -1,6 +1,6 @@
 import streamlit as st
 
-from open_notebook.domain import text_search, vector_search
+from open_notebook.domain.notebook import text_search, vector_search
 from open_notebook.utils import get_embedding
 from stream_app.note import note_list_item
 from stream_app.source import source_list_item
diff --git a/stream_app/chat.py b/stream_app/chat.py
index 44936da..ae2e732 100644
--- a/stream_app/chat.py
+++ b/stream_app/chat.py
@@ -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
diff --git a/stream_app/note.py b/stream_app/note.py
index b085cad..50c35cc 100644
--- a/stream_app/note.py
+++ b/stream_app/note.py
@@ -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
 

From 8bb5db158f233b0dce5da8fed6cc92809fc41f0a Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:09:24 -0300
Subject: [PATCH 03/14] implement model config

---
 open_notebook/config.py                       |  13 ++
 open_notebook/domain/models.py                |  38 ++++
 open_notebook/graphs/chat.py                  |   8 +-
 .../graphs/content_processing/audio.py        |  11 +-
 open_notebook/graphs/doc_query.py             |   2 +-
 open_notebook/graphs/multipattern.py          |   4 +-
 open_notebook/graphs/pattern.py               |   5 +-
 open_notebook/graphs/recursive_toc.py         |   3 +-
 open_notebook/graphs/summary.py               |   5 +-
 open_notebook/graphs/utils.py                 |   7 +-
 open_notebook/llm_router.py                   |  35 ---
 open_notebook/models/__init__.py              |   0
 open_notebook/models/embedding_models.py      |  62 +++++
 open_notebook/{ => models}/llms.py            |  50 ++---
 open_notebook/models/speech_to_text_models.py |  62 +++++
 open_notebook/plugins/podcasts.py             |   2 +-
 open_notebook/utils.py                        |  18 --
 pages/9_⚙️_Settings.py                         | 212 ++++++++++++++++++
 stream_app/source.py                          |   2 +-
 19 files changed, 434 insertions(+), 105 deletions(-)
 create mode 100644 open_notebook/domain/models.py
 delete mode 100644 open_notebook/llm_router.py
 create mode 100644 open_notebook/models/__init__.py
 create mode 100644 open_notebook/models/embedding_models.py
 rename open_notebook/{ => models}/llms.py (86%)
 create mode 100644 open_notebook/models/speech_to_text_models.py
 create mode 100644 pages/9_⚙️_Settings.py

diff --git a/open_notebook/config.py b/open_notebook/config.py
index 4543a55..38a6972 100644
--- a/open_notebook/config.py
+++ b/open_notebook/config.py
@@ -3,6 +3,10 @@ import os
 import yaml
 from loguru import logger
 
+from open_notebook.domain.models import DefaultModels
+from open_notebook.models.embedding_models import get_embedding_model
+from open_notebook.models.speech_to_text_models import get_speech_to_text_model
+
 # todo: enable config file overwrite with env vars
 current_dir = os.path.dirname(os.path.abspath(__file__))
 project_root = os.path.dirname(current_dir)
@@ -32,3 +36,12 @@ 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()
+
+EMBEDDING_MODEL = get_embedding_model(DEFAULT_MODELS.default_embedding_model)
+
+SPEECH_TO_TEXT_MODEL = get_speech_to_text_model(
+    DEFAULT_MODELS.default_speech_to_text_model
+)
diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py
new file mode 100644
index 0000000..12074ff
--- /dev/null
+++ b/open_notebook/domain/models.py
@@ -0,0 +1,38 @@
+from typing import ClassVar, Optional
+
+from loguru import logger
+from pydantic import BaseModel
+
+from open_notebook.domain.base import ObjectModel
+from open_notebook.repository import (
+    repo_query,
+    repo_update,
+)
+
+
+class Model(ObjectModel):
+    table_name: ClassVar[str] = "model"
+    name: str
+    provider: str
+    type: str
+
+
+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:
+            logger.debug(result)
+            return DefaultModels(**result[0])
+
+    @classmethod
+    def update(self, data):
+        repo_update("open_notebook:default_models", data)
diff --git a/open_notebook/graphs/chat.py b/open_notebook/graphs/chat.py
index 37c6f04..84a30be 100644
--- a/open_notebook/graphs/chat.py
+++ b/open_notebook/graphs/chat.py
@@ -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,
diff --git a/open_notebook/graphs/content_processing/audio.py b/open_notebook/graphs/content_processing/audio.py
index 5afafb7..4921403 100644
--- a/open_notebook/graphs/content_processing/audio.py
+++ b/open_notebook/graphs/content_processing/audio.py
@@ -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)}
 
diff --git a/open_notebook/graphs/doc_query.py b/open_notebook/graphs/doc_query.py
index 2c673db..9a2ffb9 100644
--- a/open_notebook/graphs/doc_query.py
+++ b/open_notebook/graphs/doc_query.py
@@ -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
 
 
diff --git a/open_notebook/graphs/multipattern.py b/open_notebook/graphs/multipattern.py
index 8f2b6d7..e74d7a3 100644
--- a/open_notebook/graphs/multipattern.py
+++ b/open_notebook/graphs/multipattern.py
@@ -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)
diff --git a/open_notebook/graphs/pattern.py b/open_notebook/graphs/pattern.py
index b7a9bd0..c47cc14 100644
--- a/open_notebook/graphs/pattern.py
+++ b/open_notebook/graphs/pattern.py
@@ -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(
diff --git a/open_notebook/graphs/recursive_toc.py b/open_notebook/graphs/recursive_toc.py
index 9cffc5d..a9cb795 100644
--- a/open_notebook/graphs/recursive_toc.py
+++ b/open_notebook/graphs/recursive_toc.py
@@ -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(
diff --git a/open_notebook/graphs/summary.py b/open_notebook/graphs/summary.py
index a262c5e..df54ff5 100644
--- a/open_notebook/graphs/summary.py
+++ b/open_notebook/graphs/summary.py
@@ -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 {
diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py
index 67f6866..9e95fcd 100644
--- a/open_notebook/graphs/utils.py
+++ b/open_notebook/graphs/utils.py
@@ -1,8 +1,7 @@
-import os
-
 from langchain.output_parsers import OutputFixingParser
 
-from open_notebook.llm_router import get_langchain_model
+from open_notebook.config import DEFAULT_MODELS
+from open_notebook.models.llms import get_langchain_model
 from open_notebook.prompter import Prompter
 
 
@@ -15,7 +14,7 @@ def run_pattern(
     output_fixing_model_name=None,
 ) -> dict:
     if not model_name:
-        model_name = os.environ["DEFAULT_MODEL"]
+        model_name = DEFAULT_MODELS.default_transformation_model
 
     chain = get_langchain_model(model_name)
     if parser:
diff --git a/open_notebook/llm_router.py b/open_notebook/llm_router.py
deleted file mode 100644
index 9e1096e..0000000
--- a/open_notebook/llm_router.py
+++ /dev/null
@@ -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()
diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/open_notebook/models/embedding_models.py b/open_notebook/models/embedding_models.py
new file mode 100644
index 0000000..aaaec65
--- /dev/null
+++ b/open_notebook/models/embedding_models.py
@@ -0,0 +1,62 @@
+"""
+Classes for supporting different embedding models
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional
+
+from openai import OpenAI
+
+from open_notebook.domain.models import Model
+
+
+@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 OpenAIEmbeddingModel(EmbeddingModel):
+    model_name: str
+
+    def embed(self, text: str) -> List[float]:
+        """
+        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
+        )
+
+
+EMBEDDING_CLASS_MAP = {
+    "openai": OpenAIEmbeddingModel,
+}
+
+
+def get_embedding_model(model_id):
+    assert model_id, "Model ID cannot be empty"
+    model = Model.get(model_id)
+    if not model:
+        raise ValueError(f"Model with ID {model_id} not found")
+    if model.provider not in EMBEDDING_CLASS_MAP.keys():
+        raise ValueError(
+            f"Provider {model.provider} not compatible with Embedding Models"
+        )
+    return EMBEDDING_CLASS_MAP[model.provider](model_name=model.name)
diff --git a/open_notebook/llms.py b/open_notebook/models/llms.py
similarity index 86%
rename from open_notebook/llms.py
rename to open_notebook/models/llms.py
index efcc431..64c0213 100644
--- a/open_notebook/llms.py
+++ b/open_notebook/models/llms.py
@@ -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
+from open_notebook.domain.models import Model
 
 
 @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,
         )
 
@@ -240,26 +240,26 @@ class OpenAILanguageModel(LanguageModel):
         )
 
 
-# @dataclass
-# class EmbeddingModel(ABC):
-#     model_name: str
-#     dimensions: int
-
-#     def to_redis_vectorizer(self) -> BaseVectorizer:
-#         raise NotImplementedError
+# 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,
+}
 
 
-# @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)
+# todo: make the provider check type specific
+def get_langchain_model(model_id, json=False):
+    model = Model.get(model_id)
+    if not model:
+        raise ValueError(f"Model {model_id} not found")
+    if model.provider not in PROVIDER_CLASS_MAP.keys():
+        raise ValueError(f"Provider {model.provider} not found")
+    return PROVIDER_CLASS_MAP[model.provider](
+        model_name=model.name, json=json
+    ).to_langchain()
diff --git a/open_notebook/models/speech_to_text_models.py b/open_notebook/models/speech_to_text_models.py
new file mode 100644
index 0000000..4191812
--- /dev/null
+++ b/open_notebook/models/speech_to_text_models.py
@@ -0,0 +1,62 @@
+"""
+Classes for supporting different transcription models
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from openai import OpenAI
+
+from open_notebook.domain.models import Model
+
+
+@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
+        """
+        # 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
+
+
+SPEECH_TO_TEXT_CLASS_MAP = {
+    "openai": OpenAISpeechToTextModel,
+}
+
+
+# todo: acho que dá pra juntar todos os get models em uma coisa só
+def get_speech_to_text_model(model_id):
+    assert model_id, "Model ID cannot be empty"
+    model = Model.get(model_id)
+    if not model:
+        raise ValueError(f"Model with ID {model_id} not found")
+    if model.provider not in SPEECH_TO_TEXT_CLASS_MAP.keys():
+        raise ValueError(
+            f"Provider {model.provider} not compatible with Embedding Models"
+        )
+    return SPEECH_TO_TEXT_CLASS_MAP[model.provider](model_name=model.name)
diff --git a/open_notebook/plugins/podcasts.py b/open_notebook/plugins/podcasts.py
index 327a969..2645474 100644
--- a/open_notebook/plugins/podcasts.py
+++ b/open_notebook/plugins/podcasts.py
@@ -4,7 +4,7 @@ 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):
diff --git a/open_notebook/utils.py b/open_notebook/utils.py
index ab43172..340762e 100644
--- a/open_notebook/utils.py
+++ b/open_notebook/utils.py
@@ -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)
 
diff --git a/pages/9_⚙️_Settings.py b/pages/9_⚙️_Settings.py
new file mode 100644
index 0000000..110bf7c
--- /dev/null
+++ b/pages/9_⚙️_Settings.py
@@ -0,0 +1,212 @@
+import os
+
+import streamlit as st
+
+from open_notebook.domain.models import DefaultModels, Model
+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",
+    "text generation",
+    "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["eleven_labs"] = 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."
+        )
+    model_name = st.text_input("Model Name", "")
+    model_type = st.selectbox("Model Type", model_types)
+    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,
+        "text generation": 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}")
+
+
+# todo: check for each type of model
+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 == "text generation"
+    ]
+
+    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"
+    )
+
+    for k, v in defs.items():
+        defs[k] = v.id
+
+    if st.button("Save Defaults", key="save_defaults"):
+        DefaultModels.update(defs)
+        st.rerun()
+
+# todo: return an error if a selected model is no longer supported
+# todo: do this check on the app homepage as well
diff --git a/stream_app/source.py b/stream_app/source.py
index d8f316a..96109bb 100644
--- a/stream_app/source.py
+++ b/stream_app/source.py
@@ -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

From 859b7f6e7ec47e54003341fd5b282dfd282c86f8 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:30:29 -0300
Subject: [PATCH 04/14] simplify the model selector

---
 open_notebook/config.py                       |  13 +-
 open_notebook/graphs/utils.py                 |   6 +-
 open_notebook/models/__init__.py              | 141 ++++++++++++++++++
 open_notebook/models/embedding_models.py      |  19 ---
 open_notebook/models/llms.py                  |  27 ----
 open_notebook/models/speech_to_text_models.py |  20 ---
 6 files changed, 151 insertions(+), 75 deletions(-)

diff --git a/open_notebook/config.py b/open_notebook/config.py
index 38a6972..9b15710 100644
--- a/open_notebook/config.py
+++ b/open_notebook/config.py
@@ -4,8 +4,7 @@ import yaml
 from loguru import logger
 
 from open_notebook.domain.models import DefaultModels
-from open_notebook.models.embedding_models import get_embedding_model
-from open_notebook.models.speech_to_text_models import get_speech_to_text_model
+from open_notebook.models import get_model
 
 # todo: enable config file overwrite with env vars
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -40,8 +39,10 @@ os.makedirs(PODCASTS_FOLDER, exist_ok=True)
 
 DEFAULT_MODELS = DefaultModels.load()
 
-EMBEDDING_MODEL = get_embedding_model(DEFAULT_MODELS.default_embedding_model)
-
-SPEECH_TO_TEXT_MODEL = get_speech_to_text_model(
-    DEFAULT_MODELS.default_speech_to_text_model
+EMBEDDING_MODEL = get_model(
+    DEFAULT_MODELS.default_embedding_model, model_type="embedding"
+)
+
+SPEECH_TO_TEXT_MODEL = get_model(
+    DEFAULT_MODELS.default_speech_to_text_model, model_type="speech_to_text"
 )
diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py
index 9e95fcd..921afa5 100644
--- a/open_notebook/graphs/utils.py
+++ b/open_notebook/graphs/utils.py
@@ -1,7 +1,7 @@
 from langchain.output_parsers import OutputFixingParser
 
 from open_notebook.config import DEFAULT_MODELS
-from open_notebook.models.llms import get_langchain_model
+from open_notebook.models import get_model
 from open_notebook.prompter import Prompter
 
 
@@ -16,12 +16,12 @@ def run_pattern(
     if not model_name:
         model_name = DEFAULT_MODELS.default_transformation_model
 
-    chain = get_langchain_model(model_name)
+    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,
diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
index e69de29..84907bc 100644
--- a/open_notebook/models/__init__.py
+++ b/open_notebook/models/__init__.py
@@ -0,0 +1,141 @@
+from open_notebook.domain.models import Model
+from open_notebook.models.embedding_models import OpenAIEmbeddingModel
+from open_notebook.models.llms import (
+    AnthropicLanguageModel,
+    GeminiLanguageModel,
+    LiteLLMLanguageModel,
+    OllamaLanguageModel,
+    OpenAILanguageModel,
+    OpenRouterLanguageModel,
+    VertexAILanguageModel,
+    VertexAnthropicLanguageModel,
+)
+from open_notebook.models.speech_to_text_models import OpenAISpeechToTextModel
+
+# 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,
+    },
+    "speech_to_text": {
+        "openai": OpenAISpeechToTextModel,
+    },
+}
+
+
+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.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
+
+
+# from open_notebook.domain.models import Model
+# from open_notebook.models.embedding_models import OpenAIEmbeddingModel
+# from open_notebook.models.llms import (
+#     AnthropicLanguageModel,
+#     GeminiLanguageModel,
+#     LiteLLMLanguageModel,
+#     OllamaLanguageModel,
+#     OpenAILanguageModel,
+#     OpenRouterLanguageModel,
+#     VertexAILanguageModel,
+#     VertexAnthropicLanguageModel,
+# )
+# from open_notebook.models.speech_to_text_models import OpenAISpeechToTextModel
+
+# SPEECH_TO_TEXT_CLASS_MAP = {
+#     "openai": OpenAISpeechToTextModel,
+# }
+
+
+# # todo: acho que dá pra juntar todos os get models em uma coisa só
+# def get_speech_to_text_model(model_id):
+#     assert model_id, "Model ID cannot be empty"
+#     model = Model.get(model_id)
+#     if not model:
+#         raise ValueError(f"Model with ID {model_id} not found")
+#     if model.provider not in SPEECH_TO_TEXT_CLASS_MAP.keys():
+#         raise ValueError(
+#             f"Provider {model.provider} not compatible with Embedding Models"
+#         )
+#     return SPEECH_TO_TEXT_CLASS_MAP[model.provider](model_name=model.name)
+
+
+# # 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,
+# }
+
+
+# # todo: make the provider check type specific
+# def get_langchain_model(model_id, json=False):
+#     model = Model.get(model_id)
+#     if not model:
+#         raise ValueError(f"Model {model_id} not found")
+#     if model.provider not in PROVIDER_CLASS_MAP.keys():
+#         raise ValueError(f"Provider {model.provider} not found")
+#     return PROVIDER_CLASS_MAP[model.provider](
+#         model_name=model.name, json=json
+#     ).to_langchain()
+
+
+# EMBEDDING_CLASS_MAP = {
+#     "openai": OpenAIEmbeddingModel,
+# }
+
+
+# def get_embedding_model(model_id):
+#     assert model_id, "Model ID cannot be empty"
+#     model = Model.get(model_id)
+#     if not model:
+#         raise ValueError(f"Model with ID {model_id} not found")
+#     if model.provider not in EMBEDDING_CLASS_MAP.keys():
+#         raise ValueError(
+#             f"Provider {model.provider} not compatible with Embedding Models"
+#         )
+#     return EMBEDDING_CLASS_MAP[model.provider](model_name=model.name)
diff --git a/open_notebook/models/embedding_models.py b/open_notebook/models/embedding_models.py
index aaaec65..ad8f141 100644
--- a/open_notebook/models/embedding_models.py
+++ b/open_notebook/models/embedding_models.py
@@ -8,8 +8,6 @@ from typing import List, Optional
 
 from openai import OpenAI
 
-from open_notebook.domain.models import Model
-
 
 @dataclass
 class EmbeddingModel(ABC):
@@ -43,20 +41,3 @@ class OpenAIEmbeddingModel(EmbeddingModel):
             .data[0]
             .embedding
         )
-
-
-EMBEDDING_CLASS_MAP = {
-    "openai": OpenAIEmbeddingModel,
-}
-
-
-def get_embedding_model(model_id):
-    assert model_id, "Model ID cannot be empty"
-    model = Model.get(model_id)
-    if not model:
-        raise ValueError(f"Model with ID {model_id} not found")
-    if model.provider not in EMBEDDING_CLASS_MAP.keys():
-        raise ValueError(
-            f"Provider {model.provider} not compatible with Embedding Models"
-        )
-    return EMBEDDING_CLASS_MAP[model.provider](model_name=model.name)
diff --git a/open_notebook/models/llms.py b/open_notebook/models/llms.py
index 64c0213..ac95181 100644
--- a/open_notebook/models/llms.py
+++ b/open_notebook/models/llms.py
@@ -17,8 +17,6 @@ from langchain_ollama.chat_models import ChatOllama
 from langchain_openai.chat_models import ChatOpenAI
 from pydantic import SecretStr
 
-from open_notebook.domain.models import Model
-
 
 @dataclass
 class LanguageModel(ABC):
@@ -238,28 +236,3 @@ class OpenAILanguageModel(LanguageModel):
             streaming=self.streaming,
             top_p=self.top_p,
         )
-
-
-# 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,
-}
-
-
-# todo: make the provider check type specific
-def get_langchain_model(model_id, json=False):
-    model = Model.get(model_id)
-    if not model:
-        raise ValueError(f"Model {model_id} not found")
-    if model.provider not in PROVIDER_CLASS_MAP.keys():
-        raise ValueError(f"Provider {model.provider} not found")
-    return PROVIDER_CLASS_MAP[model.provider](
-        model_name=model.name, json=json
-    ).to_langchain()
diff --git a/open_notebook/models/speech_to_text_models.py b/open_notebook/models/speech_to_text_models.py
index 4191812..214bb45 100644
--- a/open_notebook/models/speech_to_text_models.py
+++ b/open_notebook/models/speech_to_text_models.py
@@ -8,8 +8,6 @@ from typing import Optional
 
 from openai import OpenAI
 
-from open_notebook.domain.models import Model
-
 
 @dataclass
 class SpeechToTextModel(ABC):
@@ -42,21 +40,3 @@ class OpenAISpeechToTextModel(SpeechToTextModel):
                 model=self.model_name, file=audio
             )
             return transcription.text
-
-
-SPEECH_TO_TEXT_CLASS_MAP = {
-    "openai": OpenAISpeechToTextModel,
-}
-
-
-# todo: acho que dá pra juntar todos os get models em uma coisa só
-def get_speech_to_text_model(model_id):
-    assert model_id, "Model ID cannot be empty"
-    model = Model.get(model_id)
-    if not model:
-        raise ValueError(f"Model with ID {model_id} not found")
-    if model.provider not in SPEECH_TO_TEXT_CLASS_MAP.keys():
-        raise ValueError(
-            f"Provider {model.provider} not compatible with Embedding Models"
-        )
-    return SPEECH_TO_TEXT_CLASS_MAP[model.provider](model_name=model.name)

From aca75221a71b334680dafe7b9c95facbd2952e5b Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:36:59 -0300
Subject: [PATCH 05/14] remove optional model get

---
 open_notebook/domain/base.py     |  2 +-
 open_notebook/models/__init__.py | 76 +-------------------------------
 2 files changed, 2 insertions(+), 76 deletions(-)

diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py
index da9b0ac..4fca352 100644
--- a/open_notebook/domain/base.py
+++ b/open_notebook/domain/base.py
@@ -48,7 +48,7 @@ class ObjectModel(BaseModel):
             raise DatabaseOperationError(e)
 
     @classmethod
-    def get(cls: Type[T], id: str) -> Optional[T]:
+    def get(cls: Type[T], id: str) -> T:
         if not id:
             raise InvalidInputError("ID cannot be empty")
         try:
diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
index 84907bc..dc2cf36 100644
--- a/open_notebook/models/__init__.py
+++ b/open_notebook/models/__init__.py
@@ -43,7 +43,7 @@ def get_model(model_id, model_type="language", **kwargs):
         **kwargs: Additional arguments to pass to the model constructor
     """
     assert model_id, "Model ID cannot be empty"
-    model = Model.get(model_id)
+    model: Model = Model.get(model_id)
 
     if not model:
         raise ValueError(f"Model with ID {model_id} not found")
@@ -65,77 +65,3 @@ def get_model(model_id, model_type="language", **kwargs):
         return model_instance.to_langchain()
 
     return model_instance
-
-
-# from open_notebook.domain.models import Model
-# from open_notebook.models.embedding_models import OpenAIEmbeddingModel
-# from open_notebook.models.llms import (
-#     AnthropicLanguageModel,
-#     GeminiLanguageModel,
-#     LiteLLMLanguageModel,
-#     OllamaLanguageModel,
-#     OpenAILanguageModel,
-#     OpenRouterLanguageModel,
-#     VertexAILanguageModel,
-#     VertexAnthropicLanguageModel,
-# )
-# from open_notebook.models.speech_to_text_models import OpenAISpeechToTextModel
-
-# SPEECH_TO_TEXT_CLASS_MAP = {
-#     "openai": OpenAISpeechToTextModel,
-# }
-
-
-# # todo: acho que dá pra juntar todos os get models em uma coisa só
-# def get_speech_to_text_model(model_id):
-#     assert model_id, "Model ID cannot be empty"
-#     model = Model.get(model_id)
-#     if not model:
-#         raise ValueError(f"Model with ID {model_id} not found")
-#     if model.provider not in SPEECH_TO_TEXT_CLASS_MAP.keys():
-#         raise ValueError(
-#             f"Provider {model.provider} not compatible with Embedding Models"
-#         )
-#     return SPEECH_TO_TEXT_CLASS_MAP[model.provider](model_name=model.name)
-
-
-# # 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,
-# }
-
-
-# # todo: make the provider check type specific
-# def get_langchain_model(model_id, json=False):
-#     model = Model.get(model_id)
-#     if not model:
-#         raise ValueError(f"Model {model_id} not found")
-#     if model.provider not in PROVIDER_CLASS_MAP.keys():
-#         raise ValueError(f"Provider {model.provider} not found")
-#     return PROVIDER_CLASS_MAP[model.provider](
-#         model_name=model.name, json=json
-#     ).to_langchain()
-
-
-# EMBEDDING_CLASS_MAP = {
-#     "openai": OpenAIEmbeddingModel,
-# }
-
-
-# def get_embedding_model(model_id):
-#     assert model_id, "Model ID cannot be empty"
-#     model = Model.get(model_id)
-#     if not model:
-#         raise ValueError(f"Model with ID {model_id} not found")
-#     if model.provider not in EMBEDDING_CLASS_MAP.keys():
-#         raise ValueError(
-#             f"Provider {model.provider} not compatible with Embedding Models"
-#         )
-#     return EMBEDDING_CLASS_MAP[model.provider](model_name=model.name)

From 99f177044f95c1683fd2bed0c69beb4a3cdeb69b Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:44:48 -0300
Subject: [PATCH 06/14] prevent users from picking a model type not supported
 by the provider

---
 open_notebook/models/__init__.py |  1 +
 pages/9_⚙️_Settings.py            | 41 +++++++++++++++++++-------------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
index dc2cf36..f0a1e27 100644
--- a/open_notebook/models/__init__.py
+++ b/open_notebook/models/__init__.py
@@ -30,6 +30,7 @@ MODEL_CLASS_MAP = {
     "speech_to_text": {
         "openai": OpenAISpeechToTextModel,
     },
+    "text_to_speech": {"openai": None, "elevenlabs": None},
 }
 
 
diff --git a/pages/9_⚙️_Settings.py b/pages/9_⚙️_Settings.py
index 110bf7c..aa76b50 100644
--- a/pages/9_⚙️_Settings.py
+++ b/pages/9_⚙️_Settings.py
@@ -3,6 +3,7 @@ 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(
@@ -19,10 +20,10 @@ provider_status = {}
 
 model_types = [
     # "vision",
-    "text generation",
+    "language",
     "embedding",
-    "text to speech",
-    "speech to text",
+    "text_to_speech",
+    "speech_to_text",
 ]
 
 provider_status["ollama"] = os.environ.get("OLLAMA_API_BASE") is not None
@@ -65,20 +66,30 @@ with model_tab:
             f"Unavailable Providers: {', '.join(unavailable_providers)}. Please check docs page if you wish to enable them."
         )
     model_name = st.text_input("Model Name", "")
-    model_type = st.selectbox("Model Type", model_types)
-    if st.button("Save"):
-        model = Model(name=model_name, provider=provider, type=model_type)
-        model.save()
-        st.success("Saved")
+
+    # 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)
+        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,
-        "text generation": False,
+        "language": False,
         "embedding": False,
-        "text to speech": False,
-        "speech to text": False,
+        "text_to_speech": False,
+        "speech_to_text": False,
     }
     for model in all_models:
         model_types_available[model.type] = True
@@ -107,16 +118,14 @@ def get_selected_index(models, model_id, default=0):
 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 == "text generation"
-    ]
+    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"
+        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"
+        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"]

From 2f07f0ab498b6f47187b7c45a78f4d970d652db5 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 14:46:02 -0300
Subject: [PATCH 07/14] fix search embedding

---
 pages/3_🔍_Search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pages/3_🔍_Search.py b/pages/3_🔍_Search.py
index 6c69ea2..3d6977c 100644
--- a/pages/3_🔍_Search.py
+++ b/pages/3_🔍_Search.py
@@ -1,7 +1,7 @@
 import streamlit as st
 
+from open_notebook.config import EMBEDDING_MODEL
 from open_notebook.domain.notebook import text_search, vector_search
-from open_notebook.utils import get_embedding
 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
             )

From f64897fbf880d4b7cfbea5a05f8d9d041814adb6 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 15:19:16 -0300
Subject: [PATCH 08/14] support ollama, vertex and gemini embedding models

---
 open_notebook/models/__init__.py              | 10 ++-
 open_notebook/models/embedding_models.py      | 64 ++++++++++++++++++-
 open_notebook/models/llms.py                  |  2 +
 open_notebook/models/speech_to_text_models.py |  4 +-
 poetry.lock                                   | 45 ++++++++++++-
 pyproject.toml                                |  2 +
 6 files changed, 122 insertions(+), 5 deletions(-)

diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
index f0a1e27..a1bcd33 100644
--- a/open_notebook/models/__init__.py
+++ b/open_notebook/models/__init__.py
@@ -1,5 +1,10 @@
 from open_notebook.domain.models import Model
-from open_notebook.models.embedding_models import OpenAIEmbeddingModel
+from open_notebook.models.embedding_models import (
+    GeminiEmbeddingModel,
+    OllamaEmbeddingModel,
+    OpenAIEmbeddingModel,
+    VertexEmbeddingModel,
+)
 from open_notebook.models.llms import (
     AnthropicLanguageModel,
     GeminiLanguageModel,
@@ -26,6 +31,9 @@ MODEL_CLASS_MAP = {
     },
     "embedding": {
         "openai": OpenAIEmbeddingModel,
+        "gemini": GeminiEmbeddingModel,
+        "vertexai": VertexEmbeddingModel,
+        "ollama": OllamaEmbeddingModel,
     },
     "speech_to_text": {
         "openai": OpenAISpeechToTextModel,
diff --git a/open_notebook/models/embedding_models.py b/open_notebook/models/embedding_models.py
index ad8f141..ac9466c 100644
--- a/open_notebook/models/embedding_models.py
+++ b/open_notebook/models/embedding_models.py
@@ -2,11 +2,16 @@
 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
 
-from openai import OpenAI
+import requests
+
+# todo: add support for multiple embeddings (array)
 
 
 @dataclass
@@ -25,11 +30,68 @@ class EmbeddingModel(ABC):
         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
+        """
+        # todo: make this Singleton
+        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
         """
diff --git a/open_notebook/models/llms.py b/open_notebook/models/llms.py
index ac95181..3c9046d 100644
--- a/open_notebook/models/llms.py
+++ b/open_notebook/models/llms.py
@@ -17,6 +17,8 @@ from langchain_ollama.chat_models import ChatOllama
 from langchain_openai.chat_models import ChatOpenAI
 from pydantic import SecretStr
 
+# future: is there a value on returning langchain specific models?
+
 
 @dataclass
 class LanguageModel(ABC):
diff --git a/open_notebook/models/speech_to_text_models.py b/open_notebook/models/speech_to_text_models.py
index 214bb45..aa89d51 100644
--- a/open_notebook/models/speech_to_text_models.py
+++ b/open_notebook/models/speech_to_text_models.py
@@ -6,8 +6,6 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional
 
-from openai import OpenAI
-
 
 @dataclass
 class SpeechToTextModel(ABC):
@@ -33,6 +31,8 @@ class OpenAISpeechToTextModel(SpeechToTextModel):
         """
         Transcribes an audio file into text
         """
+        from openai import OpenAI
+
         # todo: make this Singleton
         client = OpenAI()
         with open(audio_file_path, "rb") as audio:
diff --git a/poetry.lock b/poetry.lock
index 2ac5bc1..6151763 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -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"
diff --git a/pyproject.toml b/pyproject.toml
index a8bffdb..c8d5138 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"]

From 2aa453f73f5f223d8126e1d4d79e900518204ed8 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 15:39:45 -0300
Subject: [PATCH 09/14] enable support for large context fallback

---
 open_notebook/graphs/utils.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py
index 921afa5..f004e7f 100644
--- a/open_notebook/graphs/utils.py
+++ b/open_notebook/graphs/utils.py
@@ -1,8 +1,10 @@
 from langchain.output_parsers import OutputFixingParser
+from loguru import logger
 
 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(
@@ -13,7 +15,22 @@ def run_pattern(
     parser=None,
     output_fixing_model_name=None,
 ) -> dict:
-    if not model_name:
+    system_prompt = Prompter(prompt_template=pattern_name, parser=parser).render(
+        data=state
+    )
+
+    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")
@@ -27,10 +44,6 @@ def run_pattern(
             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:

From 2de8520d0c7caad07f4d0de1df3808eb0fa941b7 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 16:33:07 -0300
Subject: [PATCH 10/14] refactor database module and migrations

---
 app_home.py                                   |  34 ++----
 database/0_0_1_to_0_0_2.surrealql             |  82 --------------
 .../1.surrealql                               | 107 +++++++-----------
 migrations/1_down.surrealql                   |  24 ++++
 open_notebook/database/migrate.py             |  56 +++++++++
 open_notebook/{ => database}/repository.py    |  30 -----
 open_notebook/domain/base.py                  |  12 +-
 open_notebook/domain/models.py                |   4 +-
 open_notebook/domain/notebook.py              |   8 +-
 open_notebook/exceptions.py                   |  12 --
 10 files changed, 146 insertions(+), 223 deletions(-)
 delete mode 100644 database/0_0_1_to_0_0_2.surrealql
 rename database/db_setup.surrealql => migrations/1.surrealql (55%)
 create mode 100644 migrations/1_down.surrealql
 create mode 100644 open_notebook/database/migrate.py
 rename open_notebook/{ => database}/repository.py (70%)

diff --git a/app_home.py b/app_home.py
index 8c81175..e18ffaa 100644
--- a/app_home.py
+++ b/app_home.py
@@ -1,29 +1,15 @@
 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 stream_app.utils import version_sidebar
 
-try:
-    version_sidebar()
-    check_database_version()
+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()
+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()
diff --git a/database/0_0_1_to_0_0_2.surrealql b/database/0_0_1_to_0_0_2.surrealql
deleted file mode 100644
index aa2d8c0..0000000
--- a/database/0_0_1_to_0_0_2.surrealql
+++ /dev/null
@@ -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";
diff --git a/database/db_setup.surrealql b/migrations/1.surrealql
similarity index 55%
rename from database/db_setup.surrealql
rename to migrations/1.surrealql
index 32e9492..8ca52ca 100644
--- a/database/db_setup.surrealql
+++ b/migrations/1.surrealql
@@ -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;
diff --git a/migrations/1_down.surrealql b/migrations/1_down.surrealql
new file mode 100644
index 0000000..e53791a
--- /dev/null
+++ b/migrations/1_down.surrealql
@@ -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;
diff --git a/open_notebook/database/migrate.py b/open_notebook/database/migrate.py
new file mode 100644
index 0000000..7d99fbd
--- /dev/null
+++ b/open_notebook/database/migrate.py
@@ -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")
diff --git a/open_notebook/repository.py b/open_notebook/database/repository.py
similarity index 70%
rename from open_notebook/repository.py
rename to open_notebook/database/repository.py
index ca6f47a..59f0cca 100644
--- a/open_notebook/repository.py
+++ b/open_notebook/database/repository.py
@@ -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():
@@ -39,25 +35,6 @@ def repo_query(query_str: str, vars: Optional[Dict[str, Any]] = None):
             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)
diff --git a/open_notebook/domain/base.py b/open_notebook/domain/base.py
index 4fca352..0da80b5 100644
--- a/open_notebook/domain/base.py
+++ b/open_notebook/domain/base.py
@@ -4,18 +4,18 @@ from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar
 from loguru import logger
 from pydantic import BaseModel, ValidationError, field_validator
 
-from open_notebook.exceptions import (
-    DatabaseOperationError,
-    InvalidInputError,
-    NotFoundError,
-)
-from open_notebook.repository import (
+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")
 
diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py
index 12074ff..752fba6 100644
--- a/open_notebook/domain/models.py
+++ b/open_notebook/domain/models.py
@@ -3,11 +3,11 @@ from typing import ClassVar, Optional
 from loguru import logger
 from pydantic import BaseModel
 
-from open_notebook.domain.base import ObjectModel
-from open_notebook.repository import (
+from open_notebook.database.repository import (
     repo_query,
     repo_update,
 )
+from open_notebook.domain.base import ObjectModel
 
 
 class Model(ObjectModel):
diff --git a/open_notebook/domain/notebook.py b/open_notebook/domain/notebook.py
index b917548..e52ad8a 100644
--- a/open_notebook/domain/notebook.py
+++ b/open_notebook/domain/notebook.py
@@ -6,6 +6,10 @@ from loguru import logger
 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,
@@ -13,10 +17,6 @@ from open_notebook.exceptions import (
 )
 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_query,
-)
 from open_notebook.utils import split_text, surreal_clean
 
 
diff --git a/open_notebook/exceptions.py b/open_notebook/exceptions.py
index 501e67a..49be004 100644
--- a/open_notebook/exceptions.py
+++ b/open_notebook/exceptions.py
@@ -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."""
 

From c01b8255fcc4c4663e1c552292748ac1156d6d73 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 17:15:45 -0300
Subject: [PATCH 11/14] check for default models

---
 app_home.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/app_home.py b/app_home.py
index e18ffaa..43890e3 100644
--- a/app_home.py
+++ b/app_home.py
@@ -1,5 +1,6 @@
 import streamlit as st
 
+from open_notebook.config import DEFAULT_MODELS
 from open_notebook.database.migrate import MigrationManager
 from stream_app.utils import version_sidebar
 
@@ -11,5 +12,30 @@ if mm.needs_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."
+    )
+    st.stop()
+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."
+    )
+    st.stop()
+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")

From 9390ea82f209a4c196802afbce9c755f5b7573f7 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 17:16:06 -0300
Subject: [PATCH 12/14] cleanup todos

---
 open_notebook/config.py                  | 1 -
 open_notebook/models/embedding_models.py | 1 -
 stream_app/chat.py                       | 2 --
 3 files changed, 4 deletions(-)

diff --git a/open_notebook/config.py b/open_notebook/config.py
index 9b15710..7a45dc9 100644
--- a/open_notebook/config.py
+++ b/open_notebook/config.py
@@ -6,7 +6,6 @@ from loguru import logger
 from open_notebook.domain.models import DefaultModels
 from open_notebook.models import get_model
 
-# todo: enable config file overwrite with env vars
 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")
diff --git a/open_notebook/models/embedding_models.py b/open_notebook/models/embedding_models.py
index ac9466c..43119cf 100644
--- a/open_notebook/models/embedding_models.py
+++ b/open_notebook/models/embedding_models.py
@@ -39,7 +39,6 @@ class OllamaEmbeddingModel(EmbeddingModel):
         """
         Embeds the content using Open AI embedding
         """
-        # todo: make this Singleton
         text = text.replace("\n", " ")
         response = requests.post(
             f"{self.base_url}/api/embed",
diff --git a/stream_app/chat.py b/stream_app/chat.py
index ae2e732..00d8bd5 100644
--- a/stream_app/chat.py
+++ b/stream_app/chat.py
@@ -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"]))

From 68cd7d7dda1c152c6633a6009424a9fbf5491993 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 17:17:16 -0300
Subject: [PATCH 13/14] setup text to speech models in Podcasts

---
 open_notebook/domain/models.py                |  7 ++++
 open_notebook/models/__init__.py              |  9 ++++-
 open_notebook/models/text_to_speech_models.py | 26 ++++++++++++++
 open_notebook/plugins/podcasts.py             |  4 +--
 pages/5_🎙️_Podcasts.py                         | 34 +++++++++++++++----
 pages/9_⚙️_Settings.py                         |  6 +---
 6 files changed, 72 insertions(+), 14 deletions(-)
 create mode 100644 open_notebook/models/text_to_speech_models.py

diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py
index 752fba6..aea8551 100644
--- a/open_notebook/domain/models.py
+++ b/open_notebook/domain/models.py
@@ -16,6 +16,13 @@ class Model(ObjectModel):
     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
diff --git a/open_notebook/models/__init__.py b/open_notebook/models/__init__.py
index a1bcd33..9473c9e 100644
--- a/open_notebook/models/__init__.py
+++ b/open_notebook/models/__init__.py
@@ -16,6 +16,10 @@ from open_notebook.models.llms import (
     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 = {
@@ -38,7 +42,10 @@ MODEL_CLASS_MAP = {
     "speech_to_text": {
         "openai": OpenAISpeechToTextModel,
     },
-    "text_to_speech": {"openai": None, "elevenlabs": None},
+    "text_to_speech": {
+        "openai": OpenAITextToSpeechModel,
+        "elevenlabs": ElevenLabsTextToSpeechModel,
+    },
 }
 
 
diff --git a/open_notebook/models/text_to_speech_models.py b/open_notebook/models/text_to_speech_models.py
new file mode 100644
index 0000000..a1fb123
--- /dev/null
+++ b/open_notebook/models/text_to_speech_models.py
@@ -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
diff --git a/open_notebook/plugins/podcasts.py b/open_notebook/plugins/podcasts.py
index 2645474..8ad6050 100644
--- a/open_notebook/plugins/podcasts.py
+++ b/open_notebook/plugins/podcasts.py
@@ -1,4 +1,4 @@
-from typing import ClassVar, List, Literal, Optional
+from typing import ClassVar, List, Optional
 
 from loguru import logger
 from podcastfy.client import generate_podcast
@@ -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
diff --git a/pages/5_🎙️_Podcasts.py b/pages/5_🎙️_Podcasts.py
index 4d5e7fa..61ed8df 100644
--- a/pages/5_🎙️_Podcasts.py
+++ b/pages/5_🎙️_Podcasts.py
@@ -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"
diff --git a/pages/9_⚙️_Settings.py b/pages/9_⚙️_Settings.py
index aa76b50..02ca7e7 100644
--- a/pages/9_⚙️_Settings.py
+++ b/pages/9_⚙️_Settings.py
@@ -45,7 +45,7 @@ provider_status["openrouter"] = (
     and os.environ.get("OPENROUTER_BASE_URL") is not None
 )
 provider_status["anthropic"] = os.environ.get("ANTHROPIC_API_KEY") is not None
-provider_status["eleven_labs"] = os.environ.get("ELEVENLABS_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"]
@@ -104,7 +104,6 @@ with model_tab:
             st.warning(f"No models available for {model_type}")
 
 
-# todo: check for each type of model
 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:
@@ -216,6 +215,3 @@ with model_defaults_tab:
     if st.button("Save Defaults", key="save_defaults"):
         DefaultModels.update(defs)
         st.rerun()
-
-# todo: return an error if a selected model is no longer supported
-# todo: do this check on the app homepage as well

From 2f2cdabd2dda9c5b3e57506cbcfa5efc226efa93 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 18:28:29 -0300
Subject: [PATCH 14/14] Fix issue with model defaults and bump version

---
 Makefile                       |  4 ++++
 README.md                      | 23 +++++++++++++++-----
 app_home.py                    | 20 +++++++++--------
 docs/SETUP.md                  | 39 ++++++++++++++++++----------------
 open_notebook/config.py        | 18 ++++++++++------
 open_notebook/domain/models.py |  7 +++---
 pages/9_⚙️_Settings.py          | 19 +++++++++++------
 pyproject.toml                 |  2 +-
 8 files changed, 83 insertions(+), 49 deletions(-)

diff --git a/Makefile b/Makefile
index 0f0ce2f..f9f6152 100644
--- a/Makefile
+++ b/Makefile
@@ -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 
\ No newline at end of file
diff --git a/README.md b/README.md
index 9f9c745..564fbe0 100644
--- a/README.md
+++ b/README.md
@@ -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. 
diff --git a/app_home.py b/app_home.py
index 43890e3..9840c6c 100644
--- a/app_home.py
+++ b/app_home.py
@@ -1,9 +1,13 @@
 import streamlit as st
 
-from open_notebook.config import DEFAULT_MODELS
 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
 
+default_models = DefaultModels.load()
+
 version_sidebar()
 mm = MigrationManager()
 if mm.needs_migration:
@@ -13,27 +17,25 @@ if mm.needs_migration:
         st.success("Migration successful")
         st.rerun()
 elif (
-    not DEFAULT_MODELS.default_chat_model
-    or not DEFAULT_MODELS.default_transformation_model
+    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."
     )
-    st.stop()
-elif not DEFAULT_MODELS.default_embedding_model:
+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."
     )
-    st.stop()
-elif not DEFAULT_MODELS.default_speech_to_text_model:
+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:
+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:
+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."
     )
diff --git a/docs/SETUP.md b/docs/SETUP.md
index 52753da..9a60be2 100644
--- a/docs/SETUP.md
+++ b/docs/SETUP.md
@@ -104,7 +104,7 @@ or the shourcut
 make run
 ```
 
-## Setting up the providers
+## Setting up the providers and models
 
 Several new providers are supported now:
 
@@ -120,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.
 
diff --git a/open_notebook/config.py b/open_notebook/config.py
index 7a45dc9..8c0bdcd 100644
--- a/open_notebook/config.py
+++ b/open_notebook/config.py
@@ -38,10 +38,16 @@ os.makedirs(PODCASTS_FOLDER, exist_ok=True)
 
 DEFAULT_MODELS = DefaultModels.load()
 
-EMBEDDING_MODEL = get_model(
-    DEFAULT_MODELS.default_embedding_model, model_type="embedding"
-)
+if DEFAULT_MODELS.default_embedding_model:
+    EMBEDDING_MODEL = get_model(
+        DEFAULT_MODELS.default_embedding_model, model_type="embedding"
+    )
+else:
+    EMBEDDING_MODEL = None
 
-SPEECH_TO_TEXT_MODEL = get_model(
-    DEFAULT_MODELS.default_speech_to_text_model, model_type="speech_to_text"
-)
+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
diff --git a/open_notebook/domain/models.py b/open_notebook/domain/models.py
index aea8551..5699147 100644
--- a/open_notebook/domain/models.py
+++ b/open_notebook/domain/models.py
@@ -1,6 +1,5 @@
 from typing import ClassVar, Optional
 
-from loguru import logger
 from pydantic import BaseModel
 
 from open_notebook.database.repository import (
@@ -37,8 +36,10 @@ class DefaultModels(BaseModel):
     def load(self):
         result = repo_query("SELECT * FROM open_notebook:default_models;")
         if result:
-            logger.debug(result)
-            return DefaultModels(**result[0])
+            result = result[0]
+            dm = DefaultModels(**result)
+            return dm
+        return DefaultModels()
 
     @classmethod
     def update(self, data):
diff --git a/pages/9_⚙️_Settings.py b/pages/9_⚙️_Settings.py
index 02ca7e7..a7b515d 100644
--- a/pages/9_⚙️_Settings.py
+++ b/pages/9_⚙️_Settings.py
@@ -65,7 +65,6 @@ with model_tab:
         st.caption(
             f"Unavailable Providers: {', '.join(unavailable_providers)}. Please check docs page if you wish to enable them."
         )
-    model_name = st.text_input("Model Name", "")
 
     # Filter model types based on provider availability in MODEL_CLASS_MAP
     available_model_types = []
@@ -76,7 +75,14 @@ with model_tab:
     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)
+        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()
@@ -209,9 +215,8 @@ with model_defaults_tab:
         "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():
-        defs[k] = v.id
-
-    if st.button("Save Defaults", key="save_defaults"):
-        DefaultModels.update(defs)
-        st.rerun()
+        if v:
+            defs[k] = v.id
+    DefaultModels.update(defs)
diff --git a/pyproject.toml b/pyproject.toml
index c8d5138..50584a1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"