feat: adding conversation profile

basilisk/config/__init__.py

@@ -1,4 +1,4 @@
-from .account_config import Account, AccountOrganization
+from .account_config import Account, AccountManager, AccountOrganization
 from .account_config import get_account_config as accounts
 from .config_enums import (
 	AccountSource,
@@ -8,19 +8,26 @@
 	ReleaseChannelEnum,
 	get_account_source_labels,
 )
+from .conversation_profile import ConversationProfile
+from .conversation_profile import (
+	get_conversation_profile_config as conversation_profiles,
+)
 from .main_config import BasiliskConfig
 from .main_config import get_basilisk_config as conf
 
 __all__ = [
-	"accounts",
 	"Account",
+	"AccountManager",
 	"AccountOrganization",
-	"get_account_source_labels",
+	"AccountSource",
+	"accounts",
+	"AutomaticUpdateModeEnum",
+	"BasiliskConfig",
 	"conf",
+	"ConversationProfile",
+	"conversation_profiles",
+	"get_account_source_labels",
 	"KeyStorageMethodEnum",
-	"AccountSource",
 	"LogLevelEnum",
 	"ReleaseChannelEnum",
-	"AutomaticUpdateModeEnum",
-	"BasiliskConfig",
 ]

basilisk/config/account_config.py

@@ -3,7 +3,7 @@
 import logging
 from functools import cache, cached_property
 from os import getenv
-from typing import Any, Iterable, Optional, Union
+from typing import Annotated, Any, Iterable, Optional, Union
 from uuid import UUID, uuid4
 
 import keyring
@@ -82,6 +82,10 @@ def delete_keyring_password(self):
 			keyring.delete_password(APP_NAME, str(self.id))
 
 
+AccountInfoStr = Annotated[str, Field(pattern="^env:[a-zA-Z]+")]
+AccountInfo = Union[UUID4, AccountInfoStr]
+
+
 class Account(BaseModel):
 	"""
 	Manage API key and organization key
@@ -196,8 +200,8 @@ def active_organization(self) -> Optional[AccountOrganization]:
 
 	def reset_active_organization(self):
 		try:
-			del self.active_organization
-		except AttributeError:
+			del self.__dict__["active_organization"]
+		except KeyError:
 			pass
 
 	@property
@@ -219,9 +223,24 @@ def delete_keyring_password(self):
 		if self.api_key_storage_method == KeyStorageMethodEnum.system:
 			keyring.delete_password(APP_NAME, str(self.id))
 
+	def get_account_info(self) -> AccountInfo:
+		if self.source == AccountSource.ENV_VAR:
+			return f"env:{self.provider.name}"
+		return self.id
+
 	def __eq__(self, value: Account) -> bool:
 		return self.id == value.id
 
+	@property
+	def display_name(self) -> str:
+		organization = (
+			self.active_organization.name
+			if self.active_organization
+			else _("Personal")
+		)
+		provider_name = self.provider.name
+		return f"{self.name} ({organization}) - {provider_name}"
+
 
 config_file_name = "accounts.yml"
 
@@ -236,24 +255,28 @@ class AccountManager(BasiliskBaseSettings):
 
 	accounts: list[OnErrorOmit[Account]] = Field(default=list())
 
-	default_account_info: Optional[Union[UUID, str]] = Field(
+	default_account_info: Optional[AccountInfo] = Field(
 		default=None, union_mode="left_to_right"
 	)
 
 	@cached_property
 	def default_account(self) -> Account:
-		index = 0
-		if isinstance(self.default_account_info, UUID):
+		account = self.get_account_from_info(self.default_account_info)
+		if not account:
+			log.warning(
+				f"Default account not found for id {self.default_account_info} using the first account"
+			)
+			account = self[0]
+		return account
+
+	def get_account_from_info(self, value: AccountInfo) -> Optional[Account]:
+		if isinstance(value, UUID):
 			try:
-				return self[self.default_account_info]
+				return self[value]
 			except KeyError:
-				log.warning(
-					f"Default account not found for id {self.default_account_info} using the first account"
-				)
-		elif isinstance(
-			self.default_account_info, str
-		) and self.default_account_info.startswith("env:"):
-			provider_name = self.default_account_info[4:]
+				return None
+		elif isinstance(value, str):
+			provider_name = value[4:]
 			index = next(
 				locate(
 					self.accounts,
@@ -263,21 +286,15 @@ def default_account(self) -> Account:
 				None,
 			)
 			if index is None:
-				log.warning(
-					f"Default account not found in env variable for provider {provider_name} using the first account"
-				)
-				index = 0
-		return self.accounts[index]
+				return self.accounts[0]
+			return self.accounts[index]
 
 	def set_default_account(self, value: Optional[Account]):
 		if not value or not isinstance(value, Account):
 			self.default_account_info = None
 			del self.__dict__["default_account"]
 			return
-		if value.source == AccountSource.ENV_VAR:
-			self.default_account_info = f"env:{value.provider.name}"
-		else:
-			self.default_account_info = value.id
+		self.default_account_info = value.get_account_info()
 		self.__dict__["default_account"] = value
 
 	@field_validator("accounts", mode="after")

basilisk/config/conversation_profile.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import logging
+from functools import cache, cached_property
+from typing import Any, Iterable, Optional, Union
+from uuid import UUID, uuid4
+
+from pydantic import (
+	UUID4,
+	BaseModel,
+	ConfigDict,
+	Field,
+	OnErrorOmit,
+	field_validator,
+	model_validator,
+)
+
+from basilisk.provider import Provider, get_provider
+
+from .account_config import Account, AccountInfo, get_account_config
+from .config_helper import (
+	BasiliskBaseSettings,
+	get_settings_config_dict,
+	save_config_file,
+)
+
+log = logging.getLogger(__name__)
+
+
+class ConversationProfile(BaseModel):
+	model_config = ConfigDict(revalidate_instances="always")
+	id: UUID4 = Field(default_factory=uuid4)
+	name: str
+	system_prompt: str = Field(default="")
+	account_info: Optional[AccountInfo] = Field(default=None)
+	ai_model_info: Optional[str] = Field(
+		default=None, pattern=r"^[a-zA-Z]+/.+$"
+	)
+	max_tokens: Optional[int] = Field(default=None)
+	temperature: Optional[float] = Field(default=None)
+	top_p: Optional[float] = Field(default=None)
+	stream_mode: bool = Field(default=True)
+
+	def __init__(self, **data: Any):
+		try:
+			super().__init__(**data)
+		except Exception as e:
+			log.error(
+				f"Error in conversation profile {e}; the profile will not be accessible",
+				exc_info=e,
+			)
+			raise e
+
+	@field_validator("ai_model_info")
+	@classmethod
+	def provider_must_exist(cls, value: str):
+		if value is None:
+			return None
+		provider_id, model_id = value.split("/", 1)
+		get_provider(id=provider_id)
+		return value
+
+	@classmethod
+	def get_default(cls) -> ConversationProfile:
+		return cls(name="default", system_prompt="")
+
+	@cached_property
+	def account(self) -> Optional[Account]:
+		if self.account_info is None:
+			return None
+		return get_account_config().get_account_from_info(self.account_info)
+
+	def set_account(self, account: Optional[Account]):
+		if account is None:
+			self.account_info = None
+			if "account" in self.__dict__:
+				del self.__dict__["account"]
+		else:
+			self.account_info = account.get_account_info()
+			self.__dict__["account"] = account
+
+	@property
+	def ai_model_id(self) -> Optional[str]:
+		if self.ai_model_info is None:
+			return None
+		return self.ai_model_info.split("/", 1)[1]
+
+	@property
+	def ai_provider(self) -> Optional[Provider]:
+		if self.account is None and self.ai_model_info is None:
+			return None
+		if self.account:
+			return self.account.provider
+		if self.ai_model_info:
+			provider_id, model_id = self.ai_model_info.split("/", 1)
+			return get_provider(id=provider_id)
+
+	def set_model_info(self, provider_id: str, model_id: str):
+		self.ai_model_info = f"{provider_id}/{model_id}"
+
+	def __eq__(self, value: Optional[ConversationProfile]) -> bool:
+		if value is None:
+			return False
+		return self.id == value.id
+
+	@model_validator(mode="after")
+	def check_same_provider(self) -> ConversationProfile:
+		if self.account is not None and self.ai_model_info is not None:
+			provider_id, model_id = self.ai_model_info.split("/", 1)
+			if provider_id != self.account.provider.id:
+				raise ValueError(
+					"Model provider must be the same as account provider"
+				)
+		return self
+
+	@model_validator(mode="after")
+	def check_model_params(self) -> ConversationProfile:
+		if self.ai_model_info is None:
+			if self.max_tokens is not None:
+				raise ValueError("Max tokens must be None without model")
+			if self.temperature is not None:
+				raise ValueError("Temperature must be None without model")
+			if self.top_p is not None:
+				raise ValueError("Top P must be None without model")
+		return self
+
+
+config_file_name = "profiles.yml"
+
+
+class ConversationProfileManager(BasiliskBaseSettings):
+	model_config = get_settings_config_dict(config_file_name)
+
+	profiles: list[OnErrorOmit[ConversationProfile]] = Field(
+		default_factory=list
+	)
+
+	default_profile_id: Optional[UUID4] = Field(default=None)
+
+	def get_profile(self, **kwargs: dict) -> Optional[ConversationProfile]:
+		return next(
+			filter(
+				lambda p: all(getattr(p, k) == v for k, v in kwargs.items()),
+				self.profiles,
+			),
+			None,
+		)
+
+	@cached_property
+	def default_profile(self) -> Optional[ConversationProfile]:
+		if self.default_profile_id is None:
+			return None
+		return self.get_profile(id=self.default_profile_id)
+
+	def set_default_profile(self, value: Optional[ConversationProfile]):
+		if value is None:
+			self.default_profile_id = None
+		else:
+			self.default_profile_id = value.id
+		if "default_profile" in self.__dict__:
+			del self.__dict__["default_profile"]
+
+	@model_validator(mode="after")
+	def check_default_profile(self) -> ConversationProfileManager:
+		if self.default_profile_id is None:
+			return self
+		if self.default_profile is None:
+			raise ValueError("Default profile not found")
+		return self
+
+	def __iter__(self) -> Iterable[ConversationProfile]:
+		return iter(self.profiles)
+
+	def add(self, profile: ConversationProfile):
+		self.profiles.append(profile)
+
+	def remove(self, profile: ConversationProfile):
+		if profile == self.default_profile:
+			self.default_profile_id = None
+			if "default_profile" in self.__dict__:
+				del self.__dict__["default_profile"]
+		self.profiles.remove(profile)
+
+	def __len__(self) -> int:
+		return len(self.profiles)
+
+	def __getitem__(self, index: Union[int, UUID]) -> ConversationProfile:
+		if isinstance(index, int):
+			return self.profiles[index]
+		elif isinstance(index, UUID):
+			profile = self.get_profile(id=index)
+			if profile is None:
+				raise KeyError(f"No profile found with id {index}")
+			return profile
+		else:
+			raise TypeError(f"Invalid index type: {type(index)}")
+
+	def __delitem__(self, index: int):
+		profile = self.profiles[index]
+		self.remove(profile)
+
+	def __setitem__(self, index: Union[int, UUID], value: ConversationProfile):
+		if isinstance(index, int):
+			self.profiles[index] = value
+		elif isinstance(index, UUID):
+			profile = self.get_profile(id=index)
+			if not profile:
+				self.add(value)
+			else:
+				idx = self.profiles.index(profile)
+				self.profiles[idx] = value
+		else:
+			raise TypeError(f"Invalid index type: {type(index)}")
+
+	def save(self):
+		save_config_file(
+			self.model_dump(
+				mode="json", exclude_defaults=True, exclude_none=True
+			),
+			config_file_name,
+		)
+
+
+@cache
+def get_conversation_profile_config() -> ConversationProfileManager:
+	log.debug("Loading conversation profile config")
+	return ConversationProfileManager()

basilisk/gui/account_dialog.py

@@ -618,7 +618,7 @@ def init_ui(self):
 		)
 		btn = wx.Button(panel, wx.ID_CLOSE)
 		btn.Bind(wx.EVT_BUTTON, self.on_close)
-
+		self.SetEscapeId(btn.GetId())
 		sizer.Add(btn, 0, wx.ALL, 5)
 
 	def init_data(self):