add TreeManagerConfiguration as sub-module to global pydantic settings class

backend/main.py

@@ -17,7 +17,7 @@ from oasst_backend.config import settings
 from oasst_backend.database import engine
 from oasst_backend.models import message_tree_state
 from oasst_backend.prompt_repository import PromptRepository, TaskRepository, UserRepository
-from oasst_backend.tree_manager import TreeManager, TreeManagerConfiguration
+from oasst_backend.tree_manager import TreeManager
 from oasst_shared.exceptions import OasstError, OasstErrorCode
 from oasst_shared.schemas import protocol as protocol_schema
 from pydantic import BaseModel
@@ -120,7 +120,7 @@ if settings.DEBUG_USE_SEED_DATA:
                 pr = PromptRepository(
                     db=db, api_client=api_client, client_user=dummy_user, user_repository=ur, task_repository=tr
                 )
-                tm = TreeManager(db, pr, TreeManagerConfiguration())
+                tm = TreeManager(db, pr)
 
                 with open(settings.DEBUG_USE_SEED_DATA_PATH) as f:
                     dummy_messages_raw = json.load(f)
@@ -187,9 +187,8 @@ if settings.DEBUG_USE_SEED_DATA:
 def ensure_tree_states():
     try:
         logger.info("Startup: TreeManager.ensure_tree_states()")
-        cfg = TreeManagerConfiguration()  # TODO: decide where config is stored, e.g. load form json/yaml file
         with Session(engine) as db:
-            tm = TreeManager(db, None, configuration=cfg)
+            tm = TreeManager(db, None)
             tm.ensure_tree_states()
 
     except Exception:

backend/oasst_backend/api/v1/tasks.py

@@ -6,7 +6,7 @@ from fastapi.security.api_key import APIKey
 from loguru import logger
 from oasst_backend.api import deps
 from oasst_backend.prompt_repository import PromptRepository, TaskRepository
-from oasst_backend.tree_manager import TreeManager, TreeManagerConfiguration
+from oasst_backend.tree_manager import TreeManager
 from oasst_shared.exceptions import OasstError, OasstErrorCode
 from oasst_shared.schemas import protocol as protocol_schema
 from sqlmodel import Session
@@ -36,8 +36,7 @@ def request_task(
 
     try:
         pr = PromptRepository(db, api_client, client_user=request.user)
-        tree_manager_config = TreeManagerConfiguration()
-        tm = TreeManager(db, pr, tree_manager_config)
+        tm = TreeManager(db, pr)
         task, message_tree_id, parent_message_id = tm.next_task(request.type)
         pr.task_repository.store_task(task, message_tree_id, parent_message_id, request.collective)
 
@@ -113,8 +112,7 @@ async def tasks_interaction(
 
     try:
         pr = PromptRepository(db, api_client, client_user=interaction.user)
-        tree_manager_config = TreeManagerConfiguration()
-        tm = TreeManager(db, pr, tree_manager_config)
+        tm = TreeManager(db, pr)
         return await tm.handle_interaction(interaction)
 
     except OasstError:

backend/oasst_backend/config.py

@@ -1,7 +1,54 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import AnyHttpUrl, BaseSettings, FilePath, PostgresDsn, validator
+from oasst_shared.schemas import protocol as protocol_schema
+from pydantic import AnyHttpUrl, BaseModel, BaseSettings, FilePath, PostgresDsn, validator
+
+
+class TreeManagerConfiguration(BaseModel):
+    """TreeManager configuration settings"""
+
+    max_active_trees: int = 10
+    """Maximum number of concurrently active message trees in the database.
+    No new initial prompt tasks are handed out to users if this
+    number is reached."""
+
+    max_tree_depth: int = 6
+    """Maximum depth of message tree."""
+
+    max_children_count: int = 5
+    """Maximum number of reply messages per tree node."""
+
+    goal_tree_size: int = 15
+    """Total number of messages to gather per tree"""
+
+    num_reviews_initial_prompt: int = 3
+    """Number of peer review checks to collect in INITIAL_PROMPT_REVIEW state."""
+
+    num_reviews_reply: int = 3
+    """Number of peer review checks to collect per reply (other than initial_prompt)"""
+
+    p_full_labeling_review_prompt: float = 0.1
+    """Probability of full text-labeling (instead of mandatory only) for initial prompts"""
+
+    p_full_labeling_review_reply_assistant: float = 0.1
+    """Probability of full text-labeling (instead of mandatory only) for assistant replies"""
+
+    p_full_labeling_review_reply_prompter: float = 0.1
+    """Probability of full text-labeling (instead of mandatory only) for prompter replies"""
+
+    acceptance_threshold_initial_prompt: float = 0.6
+    """Threshold for accepting an initial prompt"""
+
+    acceptance_threshold_reply: float = 0.6
+    """Threshold for accepting a reply"""
+
+    num_required_rankings: int = 3
+    """Number of rankings in which the message participated."""
+
+    mandatory_labels_initial_prompt: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
+    mandatory_labels_assistant_reply: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
+    mandatory_labels_prompter_reply: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
 
 
 class Settings(BaseSettings):
@@ -54,5 +101,13 @@ class Settings(BaseSettings):
             return v
         raise ValueError(v)
 
+    tree_manager: Optional[TreeManagerConfiguration] = TreeManagerConfiguration()
 
-settings = Settings(_env_file=".env")
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+        case_sensitive = False
+        env_nested_delimiter = "__"
+
+
+settings = Settings()

backend/oasst_backend/tree_manager.py

@@ -8,7 +8,7 @@ import numpy as np
 import pydantic
 from loguru import logger
 from oasst_backend.api.v1.utils import prepare_conversation, prepare_conversation_message_list
-from oasst_backend.config import settings
+from oasst_backend.config import TreeManagerConfiguration, settings
 from oasst_backend.models import Message, MessageReaction, MessageTreeState, TextLabels, message_tree_state
 from oasst_backend.prompt_repository import PromptRepository
 from oasst_backend.utils.hugging_face import HfEmbeddingModel, HfUrl, HuggingFaceAPI
@@ -18,52 +18,6 @@ from sqlalchemy.sql import text
 from sqlmodel import Session, func
 
 
-class TreeManagerConfiguration(pydantic.BaseModel):
-    """TreeManager configuration settings"""
-
-    max_active_trees: int = 10
-    """Maximum number of concurrently active message trees in the database.
-    No new initial prompt tasks are handed out to users if this
-    number is reached."""
-
-    max_tree_depth: int = 6
-    """Maximum depth of message tree."""
-
-    max_children_count: int = 5
-    """Maximum number of reply messages per tree node."""
-
-    goal_tree_size: int = 15
-    """Total number of messages to gather per tree"""
-
-    num_reviews_initial_prompt: int = 3
-    """Number of peer review checks to collect in INITIAL_PROMPT_REVIEW state."""
-
-    num_reviews_reply: int = 3
-    """Number of peer review checks to collect per reply (other than initial_prompt)"""
-
-    p_full_labeling_review_prompt: float = 0.1
-    """Probability of full text-labeling (instead of mandatory only) for initial prompts"""
-
-    p_full_labeling_review_reply_assistant: float = 0.1
-    """Probability of full text-labeling (instead of mandatory only) for assistant replies"""
-
-    p_full_labeling_review_reply_prompter: float = 0.1
-    """Probability of full text-labeling (instead of mandatory only) for prompter replies"""
-
-    acceptance_threshold_initial_prompt: float = 0.6
-    """Threshold for accepting an initial prompt"""
-
-    acceptance_threshold_reply: float = 0.6
-    """Threshold for accepting a reply"""
-
-    num_required_rankings: int = 3
-    """Number of rankings in which the message participated."""
-
-    mandatory_labels_initial_prompt: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
-    mandatory_labels_assistant_reply: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
-    mandatory_labels_prompter_reply: Optional[list[protocol_schema.TextLabel]] = [protocol_schema.TextLabel.spam]
-
-
 class TaskType(Enum):
     NONE = -1
     RANKING = 0
@@ -114,9 +68,9 @@ class IncompleteRankingsRow(pydantic.BaseModel):
 class TreeManager:
     _all_text_labels = list(map(lambda x: x.value, protocol_schema.TextLabel))
 
-    def __init__(self, db: Session, prompt_repository: PromptRepository, configuration: TreeManagerConfiguration):
+    def __init__(self, db: Session, prompt_repository: PromptRepository):
         self.db = db
-        self.cfg = configuration
+        self.cfg = settings.tree_manager
         self.pr = prompt_repository
 
     def _task_selection(