diff --git a/backend/main.py b/backend/main.py index 76b0a35b..cf0b1f76 100644 --- a/backend/main.py +++ b/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: diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 4a976f45..17df814f 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/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: diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py index 2e2fbf95..0aeea390 100644 --- a/backend/oasst_backend/config.py +++ b/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() diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py index 8be96aae..6181d62e 100644 --- a/backend/oasst_backend/tree_manager.py +++ b/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(