diff --git a/business_objects/project.py b/business_objects/project.py index 49e1b1f4..b19fb6b9 100644 --- a/business_objects/project.py +++ b/business_objects/project.py @@ -212,7 +212,9 @@ def get_zero_shot_project_config(project_id: str, payload_id: str) -> Any: def get_or_create_queue_project( - org_id: str, user_id: str, with_commit: bool = False + org_id: str, + user_id: str, + with_commit: bool = False, ) -> Project: ## user_id is a "last used by" indicator @@ -244,12 +246,37 @@ def get_or_create_queue_project( return prj +def create_dataset_project( + org_id: str, + user_id: str, + name: str, + description: str, + tokenizer: str, + with_commit: bool = False, +) -> Project: + ## user_id is a "last used by" indicator + + prj = create( + organization_id=org_id, + name=name, + description=description, + created_by=user_id, + tokenizer=tokenizer, + tokenizer_blank=tokenizer[:2], + with_commit=with_commit, + ) + + return prj + + def create( organization_id: str, name: str, description: str, created_by: str, created_at: Optional[str] = None, + tokenizer: Optional[str] = None, + tokenizer_blank: Optional[str] = None, with_commit: bool = False, status: enums.ProjectStatus = enums.ProjectStatus.INIT_UPLOAD, ) -> Project: @@ -260,6 +287,8 @@ def create( created_by=created_by, created_at=created_at, status=status.value, + tokenizer=tokenizer, + tokenizer_blank=tokenizer_blank, ) general.add(project, with_commit) return project diff --git a/business_objects/user.py b/business_objects/user.py index b44a74d6..ab83c035 100644 --- a/business_objects/user.py +++ b/business_objects/user.py @@ -101,6 +101,14 @@ def update_organization( general.flush_or_commit(with_commit) +def update_language_display( + user_id: str, language_display: str, with_commit: bool = False +) -> None: + user = get(user_id) + user.language_display = language_display + general.flush_or_commit(with_commit) + + def __create_migration_user() -> str: organization_item = organization.get_by_name("migration") if not organization_item: diff --git a/cognition_objects/consumption_log.py b/cognition_objects/consumption_log.py new file mode 100644 index 00000000..db72b559 --- /dev/null +++ b/cognition_objects/consumption_log.py @@ -0,0 +1,90 @@ +from typing import Dict, List, Optional, Tuple, Any + +from datetime import datetime + +from submodules.model import enums + +from ..cognition_objects import message +from ..business_objects import general +from ..session import session +from ..models import CognitionConsumptionLog + + +def get(project_id: str, log_id: str) -> CognitionConsumptionLog: + return ( + session.query(CognitionConsumptionLog) + .filter( + CognitionConsumptionLog.project_id == project_id, + CognitionConsumptionLog.id == log_id, + ) + .first() + ) + + +def get_all_by_project_id( + project_id: str, +) -> List[CognitionConsumptionLog]: + return ( + session.query(CognitionConsumptionLog) + .filter( + CognitionConsumptionLog.project_id == project_id, + ) + .order_by(CognitionConsumptionLog.created_at.asc()) + .all() + ) + + +def get_all_by_project_id_for_year( + project_id: str, + year: int, +) -> List[CognitionConsumptionLog]: + return ( + session.query(CognitionConsumptionLog) + .filter( + CognitionConsumptionLog.project_id == project_id, + CognitionConsumptionLog.created_at >= datetime(year, 1, 1), + CognitionConsumptionLog.created_at < datetime(year + 1, 1, 1), + ) + .order_by(CognitionConsumptionLog.created_at.asc()) + .all() + ) + + +def create( + project_id: str, + user_id: str, + conversation_id: str, + message_id: str, + state: str, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionConsumptionLog: + conversation: CognitionConsumptionLog = CognitionConsumptionLog( + project_id=project_id, + created_by=user_id, + conversation_id=conversation_id, + message_id=message_id, + created_at=created_at, + state=state, + ) + general.add(conversation, with_commit) + return conversation + + +def update( + project_id: str, + log_id: str, + strategy_id: Optional[str] = None, + complexity: Optional[str] = None, + state: Optional[str] = None, + with_commit: bool = True, +) -> CognitionConsumptionLog: + log = get(project_id=project_id, log_id=log_id) + if strategy_id is not None: + log.strategy_id = strategy_id + if complexity is not None: + log.complexity = complexity + if state is not None: + log.state = state + general.flush_or_commit(with_commit) + return log diff --git a/cognition_objects/conversation.py b/cognition_objects/conversation.py index d1a7f491..878f428d 100644 --- a/cognition_objects/conversation.py +++ b/cognition_objects/conversation.py @@ -19,6 +19,19 @@ def get(project_id: str, conversation_id: str) -> CognitionConversation: ) +def get_all_by_spreadsheet_row_id( + spreadsheet_row_id: str, +) -> List[CognitionConversation]: + return ( + session.query(CognitionConversation) + .filter( + CognitionConversation.synopsis_spreadsheet_row_id == spreadsheet_row_id, + ) + .order_by(CognitionConversation.created_at.asc()) + .all() + ) + + def get_all_paginated_by_project_id( project_id: str, page: int, limit: int ) -> Tuple[int, int, List[CognitionConversation]]: @@ -42,7 +55,7 @@ def get_all_paginated_by_project_id( paginated_result = ( session.query(CognitionConversation) .filter(CognitionConversation.project_id == project_id) - .order_by(CognitionConversation.created_at.asc()) + .order_by(CognitionConversation.created_at.desc()) .limit(limit) .offset((page - 1) * limit) .all() @@ -56,6 +69,8 @@ def create( project_id: str, user_id: str, with_commit: bool = True, + synopsis_column: Optional[str] = None, + spreadsheet_row_id: Optional[str] = None, created_at: Optional[datetime] = None, ) -> CognitionConversation: conversation: CognitionConversation = CognitionConversation( @@ -63,6 +78,8 @@ def create( created_by=user_id, created_at=created_at, scope_dict={}, + synopsis_column=synopsis_column, + synopsis_spreadsheet_row_id=spreadsheet_row_id, ) general.add(conversation, with_commit) return conversation @@ -90,11 +107,28 @@ def update( project_id: str, conversation_id: str, scope_dict: Optional[Dict[str, Any]] = None, + header: Optional[str] = None, + error: Optional[str] = None, with_commit: bool = True, ) -> CognitionConversation: conversation_entity = get(project_id, conversation_id) if scope_dict is not None: conversation_entity.scope_dict = scope_dict + if header is not None: + conversation_entity.header = header + if error is not None: + conversation_entity.error = error + general.flush_or_commit(with_commit) + return conversation_entity + + +def clear_error( + project_id: str, + conversation_id: str, + with_commit: bool = True, +) -> CognitionConversation: + conversation_entity = get(project_id, conversation_id) + conversation_entity.error = None general.flush_or_commit(with_commit) return conversation_entity @@ -104,6 +138,9 @@ def update_message( conversation_id: str, message_id: str, answer: Optional[str] = None, + feedback_value: Optional[str] = None, + feedback_category: Optional[str] = None, + feedback_message: Optional[str] = None, strategy_id: Optional[str] = None, scope_dict_diff_previous_conversation: Optional[Dict[str, Any]] = None, with_commit: bool = True, @@ -113,6 +150,12 @@ def update_message( message_entity.strategy_id = strategy_id if answer is not None: message_entity.answer = answer + if feedback_value is not None: + message_entity.feedback_value = feedback_value + if feedback_category is not None: + message_entity.feedback_category = feedback_category + if feedback_message is not None: + message_entity.feedback_message = feedback_message if scope_dict_diff_previous_conversation is not None: message_entity.scope_dict_diff_previous_conversation = ( scope_dict_diff_previous_conversation diff --git a/cognition_objects/markdown_dataset.py b/cognition_objects/markdown_dataset.py index 64468490..00e5ab9b 100644 --- a/cognition_objects/markdown_dataset.py +++ b/cognition_objects/markdown_dataset.py @@ -52,17 +52,21 @@ def get_enriched(org_id: str, id: str) -> Dict[str, Any]: def get_all_paginated_for_category_origin( org_id: str, - category_origin: str, page: int, limit: int, + category_origin: Optional[str] = None, ) -> Tuple[int, int, List[CognitionMarkdownDataset]]: - total_count = ( - session.query(CognitionMarkdownDataset.id) - .filter(CognitionMarkdownDataset.organization_id == org_id) - .filter(CognitionMarkdownDataset.category_origin == category_origin) - .count() + query = session.query(CognitionMarkdownDataset.id).filter( + CognitionMarkdownDataset.organization_id == org_id ) + if category_origin is not None: + query = query.filter( + CognitionMarkdownDataset.category_origin == category_origin + ) + + total_count = query.count() + num_pages = int(total_count / limit) if total_count % limit > 0: num_pages += 1 diff --git a/cognition_objects/message.py b/cognition_objects/message.py index da306188..5bde633f 100644 --- a/cognition_objects/message.py +++ b/cognition_objects/message.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import Any, Dict, List, Optional from datetime import datetime from ..business_objects import general from ..session import session @@ -92,6 +92,36 @@ def create( return message +def update( + project_id: str, + message_id: str, + answer: Optional[str] = None, + facts: Optional[List[Dict[str, Any]]] = None, + selection_widget: Optional[List[Dict[str, Any]]] = None, + feedback_value: Optional[str] = None, + feedback_category: Optional[str] = None, + feedback_message: Optional[str] = None, + with_commit: bool = True, +) -> CognitionMessage: + message = get(project_id, message_id) + if answer is not None: + message.answer = answer + if facts is not None: + message.facts = facts + if selection_widget is not None: + message.selection_widget = selection_widget + if feedback_value is not None: + message.feedback_value = feedback_value + if feedback_category is not None: + message.feedback_category = feedback_category + if feedback_message is not None: + message.feedback_message = feedback_message + + general.flush_or_commit(with_commit) + + return message + + def delete(project_id: str, message_id: str, with_commit: bool = True) -> None: session.query(CognitionMessage).filter( CognitionMessage.project_id == project_id, diff --git a/cognition_objects/pipeline_log.py b/cognition_objects/pipeline_log.py index 787cff0f..2acfd504 100644 --- a/cognition_objects/pipeline_log.py +++ b/cognition_objects/pipeline_log.py @@ -64,6 +64,7 @@ def create( time_elapsed: float, record_dict_diff_previous: Dict[str, Any], scope_dict_diff_previous: Dict[str, Any], + skipped_step: Optional[bool] = None, with_commit: bool = True, created_at: Optional[datetime] = None, ) -> CognitionPipelineLogs: @@ -80,6 +81,7 @@ def create( time_elapsed=time_elapsed, record_dict_diff_previous_message=record_dict_diff_previous, scope_dict_diff_previous_message=scope_dict_diff_previous, + skipped_step=skipped_step, ) general.add(log, with_commit) diff --git a/cognition_objects/project.py b/cognition_objects/project.py index 6d2d398a..efb9215e 100644 --- a/cognition_objects/project.py +++ b/cognition_objects/project.py @@ -44,6 +44,7 @@ def create( color: str, org_id: str, user_id: str, + interface_type: str, refinery_references_project_id: str, refinery_queries_project_id: str, refinery_relevances_project_id: str, @@ -61,6 +62,15 @@ def routing( record_dict['routing'] = 'Low-code strategy' return record_dict, scope_dict +""" + + execute_query_enrichment_if_source_code = """from typing import Dict, Any, Tuple + +def check_execute( + record_dict: Dict[str, Any], scope_dict: Dict[str, Any] +) -> bool: + return True + """ project: CognitionProject = CognitionProject( @@ -70,11 +80,13 @@ def routing( organization_id=org_id, created_by=user_id, created_at=created_at, + interface_type=interface_type, refinery_references_project_id=refinery_references_project_id, refinery_question_project_id=refinery_queries_project_id, refinery_relevance_project_id=refinery_relevances_project_id, operator_routing_source_code=operator_routing_source_code, refinery_synchronization_interval_option=enums.RefinerySynchronizationIntervalOption.NEVER.value, + execute_query_enrichment_if_source_code=execute_query_enrichment_if_source_code, ) general.add(project, with_commit) return project @@ -84,8 +96,12 @@ def update( project_id: str, name: Optional[str] = None, description: Optional[str] = None, + customer_color_primary: Optional[str] = None, + customer_color_primary_only_accent: Optional[bool] = None, + customer_color_secondary: Optional[str] = None, operator_routing_source_code: Optional[str] = None, refinery_synchronization_interval_option: Optional[str] = None, + execute_query_enrichment_if_source_code: Optional[str] = None, with_commit: bool = True, ) -> CognitionProject: project: CognitionProject = get(project_id) @@ -93,12 +109,22 @@ def update( project.name = name if description is not None: project.description = description + if customer_color_primary is not None: + project.customer_color_primary = customer_color_primary + if customer_color_primary_only_accent is not None: + project.customer_color_primary_only_accent = customer_color_primary_only_accent + if customer_color_secondary is not None: + project.customer_color_secondary = customer_color_secondary if operator_routing_source_code is not None: project.operator_routing_source_code = operator_routing_source_code if refinery_synchronization_interval_option is not None: project.refinery_synchronization_interval_option = ( refinery_synchronization_interval_option ) + if execute_query_enrichment_if_source_code is not None: + project.execute_query_enrichment_if_source_code = ( + execute_query_enrichment_if_source_code + ) general.flush_or_commit(with_commit) return project diff --git a/cognition_objects/spreadsheet.py b/cognition_objects/spreadsheet.py new file mode 100644 index 00000000..ab6a4d27 --- /dev/null +++ b/cognition_objects/spreadsheet.py @@ -0,0 +1,81 @@ +from typing import Dict, List, Optional, Tuple, Any + +from datetime import datetime + +from ..business_objects import general +from ..session import session +from ..models import CognitionSynopsisSpreadsheet + + +def get(project_id: str, spreadsheet_id: str) -> CognitionSynopsisSpreadsheet: + return ( + session.query(CognitionSynopsisSpreadsheet) + .filter( + CognitionSynopsisSpreadsheet.id == spreadsheet_id, + CognitionSynopsisSpreadsheet.project_id == project_id, + ) + .first() + ) + + +def get_all_by_project_id( + project_id: str, +) -> List[CognitionSynopsisSpreadsheet]: + return ( + session.query(CognitionSynopsisSpreadsheet) + .filter( + CognitionSynopsisSpreadsheet.project_id == project_id, + ) + .order_by(CognitionSynopsisSpreadsheet.created_at.asc()) + .all() + ) + + +def create( + project_id: str, + user_id: str, + name: str, + synopsis_type: str, + dataset_id: Optional[str] = None, + filter_name: Optional[str] = None, + task_scope_dict: Optional[Dict[str, Any]] = None, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionSynopsisSpreadsheet: + spreadsheet = CognitionSynopsisSpreadsheet( + project_id=project_id, + dataset_id=dataset_id, + name=name, + filter_attribute_name=filter_name, + synopsis_type=synopsis_type, + task_scope_dict=task_scope_dict, + created_at=created_at, + created_by=user_id, + ) + general.add(spreadsheet, with_commit) + return spreadsheet + + +def update( + project_id: str, + spreadsheet_id: str, + name: Optional[str] = None, + filter_name: Optional[str] = None, + with_commit: bool = True, +) -> CognitionSynopsisSpreadsheet: + spreadsheet = get(project_id, spreadsheet_id) + if name is not None: + spreadsheet.name = name + if filter_name is not None: + spreadsheet.filter_name = filter_name + + general.flush_or_commit(with_commit) + return spreadsheet + + +def delete(project_id: str, spreadsheet_id: str, with_commit: bool = True) -> None: + session.query(CognitionSynopsisSpreadsheet).filter( + CognitionSynopsisSpreadsheet.id == spreadsheet_id, + CognitionSynopsisSpreadsheet.project_id == project_id, + ).delete() + general.flush_or_commit(with_commit) diff --git a/cognition_objects/spreadsheet_row.py b/cognition_objects/spreadsheet_row.py new file mode 100644 index 00000000..f868ca86 --- /dev/null +++ b/cognition_objects/spreadsheet_row.py @@ -0,0 +1,61 @@ +from typing import Dict, List, Optional, Tuple, Any + +from datetime import datetime + +from ..business_objects import general +from ..session import session +from ..models import CognitionSynopsisSpreadsheetRow + + +def get(spreadsheet_row_id: str) -> CognitionSynopsisSpreadsheetRow: + return ( + session.query(CognitionSynopsisSpreadsheetRow) + .filter( + CognitionSynopsisSpreadsheetRow.id == spreadsheet_row_id, + ) + .first() + ) + + +def get_all_by_spreadsheet_id( + spreadsheet_id: str, +) -> List[CognitionSynopsisSpreadsheetRow]: + return ( + session.query(CognitionSynopsisSpreadsheetRow) + .filter( + CognitionSynopsisSpreadsheetRow.spreadsheet_id == spreadsheet_id, + ) + .order_by(CognitionSynopsisSpreadsheetRow.created_at.asc()) + .all() + ) + + +def create( + spreadsheet_id: str, + project_id: str, + user_id: str, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionSynopsisSpreadsheetRow: + spreadsheet_row = CognitionSynopsisSpreadsheetRow( + project_id=project_id, + spreadsheet_id=spreadsheet_id, + created_at=created_at, + created_by=user_id, + ) + general.add(spreadsheet_row, with_commit) + return spreadsheet_row + + +def delete( + project_id: str, + spreadsheet_id: str, + spreadsheet_row_id: str, + with_commit: bool = True, +) -> None: + session.query(CognitionSynopsisSpreadsheetRow).filter( + CognitionSynopsisSpreadsheetRow.id == spreadsheet_row_id, + CognitionSynopsisSpreadsheetRow.spreadsheet_id == spreadsheet_id, + CognitionSynopsisSpreadsheetRow.project_id == project_id, + ).delete() + general.flush_or_commit(with_commit) diff --git a/cognition_objects/strategy_step.py b/cognition_objects/strategy_step.py index 92ee0846..7935c9e3 100644 --- a/cognition_objects/strategy_step.py +++ b/cognition_objects/strategy_step.py @@ -54,9 +54,18 @@ def create( step_type: str, position: int, config: Dict, + progress_text: str, with_commit: bool = True, created_at: Optional[datetime] = None, ) -> CognitionStrategyStep: + execute_if_source_code = """from typing import Dict, Any, Tuple + +def check_execute( + record_dict: Dict[str, Any], scope_dict: Dict[str, Any] +) -> bool: + return True +""" + strategy: CognitionStrategyStep = CognitionStrategyStep( project_id=project_id, strategy_id=strategy_id, @@ -67,6 +76,8 @@ def create( step_type=step_type, position=position, config=config, + progress_text=progress_text, + execute_if_source_code=execute_if_source_code, ) general.add(strategy, with_commit) @@ -80,6 +91,9 @@ def update( description: Optional[str] = None, position: Optional[int] = None, config: Optional[Dict] = None, + progress_text: Optional[str] = None, + enable_emissions: Optional[bool] = None, + execute_if_source_code: Optional[str] = None, with_commit: bool = True, ) -> CognitionStrategyStep: strategy_step: CognitionStrategyStep = get(project_id, strategy_step_id) @@ -93,6 +107,12 @@ def update( if config is not None: strategy_step.config = config flag_modified(strategy_step, "config") + if progress_text is not None: + strategy_step.progress_text = progress_text + if enable_emissions is not None: + strategy_step.enable_emissions = enable_emissions + if execute_if_source_code is not None: + strategy_step.execute_if_source_code = execute_if_source_code general.flush_or_commit(with_commit) return strategy_step diff --git a/cognition_objects/team.py b/cognition_objects/team.py new file mode 100644 index 00000000..e1f19978 --- /dev/null +++ b/cognition_objects/team.py @@ -0,0 +1,71 @@ +from typing import List, Optional, Dict +from datetime import datetime +from ..business_objects import general +from ..session import session +from ..models import ( + CognitionTeam, +) +from ..util import prevent_sql_injection +from sqlalchemy import or_ + + +def get(team_id: str) -> CognitionTeam: + return ( + session.query(CognitionTeam) + .filter( + CognitionTeam.id == team_id, + ) + .first() + ) + + +def get_all_by_org_id(org_id: str) -> List[CognitionTeam]: + # Note, atm this doesn't mean all but all on org level + return ( + session.query(CognitionTeam) + .filter(CognitionTeam.organization_id == org_id) + .order_by(CognitionTeam.created_at.asc()) + .all() + ) + + +def create( + organization_id: str, + user_id: str, + name: str, + description: str, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionTeam: + team: CognitionTeam = CognitionTeam( + created_by=user_id, + organization_id=organization_id, + name=name, + description=description, + created_at=created_at, + ) + general.add(team, with_commit) + return team + + +def update( + team_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + with_commit: bool = True, +) -> CognitionTeam: + team: CognitionTeam = get(team_id) + + if name is not None: + team.name = name + if description is not None: + team.description = description + general.flush_or_commit(with_commit) + return team + + +def delete(team_id: str, with_commit: bool = True) -> None: + session.query(CognitionTeam).filter( + CognitionTeam.id == team_id, + ).delete() + general.flush_or_commit(with_commit) diff --git a/cognition_objects/team_member.py b/cognition_objects/team_member.py new file mode 100644 index 00000000..1e8631e8 --- /dev/null +++ b/cognition_objects/team_member.py @@ -0,0 +1,80 @@ +from typing import List, Optional, Dict +from datetime import datetime +from ..business_objects import general +from ..session import session +from ..models import ( + CognitionTeamMember, +) + + +def get(team_member_id: str) -> CognitionTeamMember: + return ( + session.query(CognitionTeamMember) + .filter( + CognitionTeamMember.id == team_member_id, + ) + .first() + ) + + +def get_all_by_team_id(team_id: str) -> List[CognitionTeamMember]: + return ( + session.query(CognitionTeamMember) + .filter(CognitionTeamMember.team_id == team_id) + .order_by(CognitionTeamMember.created_at.asc()) + .all() + ) + + +def get_all_by_user_id(user_id: str) -> List[CognitionTeamMember]: + return ( + session.query(CognitionTeamMember) + .filter(CognitionTeamMember.user_id == user_id) + .order_by(CognitionTeamMember.created_at.asc()) + .all() + ) + + +def get_by_team_and_user_id(team_id: str, user_id: str) -> CognitionTeamMember: + return ( + session.query(CognitionTeamMember) + .filter( + CognitionTeamMember.team_id == team_id, + CognitionTeamMember.user_id == user_id, + ) + .first() + ) + + +def create( + team_id: str, + user_id: str, + created_by: str, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionTeamMember: + team_member: CognitionTeamMember = CognitionTeamMember( + created_by=user_id, + team_id=team_id, + user_id=user_id, + created_at=created_at, + ) + general.add(team_member, with_commit) + return team_member + + +def delete(team_member_id: str, with_commit: bool = True) -> None: + session.query(CognitionTeamMember).filter( + CognitionTeamMember.id == team_member_id, + ).delete() + general.flush_or_commit(with_commit) + + +def delete_by_team_and_user_id( + team_id: str, user_id: str, with_commit: bool = True +) -> None: + session.query(CognitionTeamMember).filter( + CognitionTeamMember.team_id == team_id, + CognitionTeamMember.user_id == user_id, + ).delete() + general.flush_or_commit(with_commit) diff --git a/cognition_objects/team_project_access.py b/cognition_objects/team_project_access.py new file mode 100644 index 00000000..6db65ca9 --- /dev/null +++ b/cognition_objects/team_project_access.py @@ -0,0 +1,69 @@ +from typing import List, Optional, Dict +from datetime import datetime +from ..business_objects import general +from ..session import session +from ..models import ( + CognitionTeamProjectAccess, +) + + +def get(team_project_access_id: str) -> CognitionTeamProjectAccess: + return ( + session.query(CognitionTeamProjectAccess) + .filter( + CognitionTeamProjectAccess.id == team_project_access_id, + ) + .first() + ) + + +def get_all_by_team_id(team_id: str) -> List[CognitionTeamProjectAccess]: + return ( + session.query(CognitionTeamProjectAccess) + .filter(CognitionTeamProjectAccess.team_id == team_id) + .order_by(CognitionTeamProjectAccess.created_at.asc()) + .all() + ) + + +def get_all_by_project_id(project_id: str) -> List[CognitionTeamProjectAccess]: + return ( + session.query(CognitionTeamProjectAccess) + .filter(CognitionTeamProjectAccess.project_id == project_id) + .order_by(CognitionTeamProjectAccess.created_at.asc()) + .all() + ) + + +def create( + user_id: str, + team_id: str, + project_id: str, + with_commit: bool = True, + created_at: Optional[datetime] = None, +) -> CognitionTeamProjectAccess: + team_project_access: CognitionTeamProjectAccess = CognitionTeamProjectAccess( + created_by=user_id, + team_id=team_id, + project_id=project_id, + created_at=created_at, + ) + general.add(team_project_access, with_commit) + return team_project_access + + +def delete(team_project_access_id: str, with_commit: bool = True) -> None: + session.query(CognitionTeamProjectAccess).filter( + CognitionTeamProjectAccess.id == team_project_access_id, + ).delete() + general.flush_or_commit(with_commit) + + +def delete_by_team_and_project_id( + team_id: str, project_id: str, with_commit: bool = True +) -> None: + session.query(CognitionTeamProjectAccess).filter( + CognitionTeamProjectAccess.team_id == team_id, + CognitionTeamProjectAccess.project_id == project_id, + ).delete() + general.flush_or_commit(with_commit) diff --git a/enums.py b/enums.py index 8f67194e..5c383982 100644 --- a/enums.py +++ b/enums.py @@ -130,6 +130,12 @@ class Tablenames(Enum): LLM_STEP = "llm_step" MARKDOWN_LLM_LOGS = "markdown_llm_logs" MARKDOWN_DATASET = "markdown_dataset" + CONSUMPTION_LOG = "consumption_log" + TEAM = "team" + TEAM_MEMBER = "team_member" + TEAM_PROJECT_ACCESS = "team_project_access" + SYNOPSIS_SPREADSHEET = "synopsis_spreadsheet" + SYNOPSIS_SPREADSHEET_ROW = "synopsis_spreadsheet_row" def snake_case_to_pascal_case(self): # the type name of a table is needed to create backrefs @@ -488,17 +494,59 @@ class StrategyStepType(Enum): NONE = "NONE" PYTHON = "PYTHON" LLM = "LLM" + SELECTION = "SELECTION" + QUERY_REPHRASING = "QUERY_REPHRASING" + WEBSEARCH = "WEBSEARCH" + TRUNCATE_CONTEXT = "TRUNCATE_CONTEXT" + HEADER = "HEADER" def get_description(self): return STEP_DESCRIPTIONS.get(self, "No description available") + def get_when_to_use(self): + return STEP_WHEN_TO_USE.get(self, "No description available") + + def get_progress_text(self): + return STEP_PROGRESS_TEXTS.get(self, "No progress text available") + STEP_DESCRIPTIONS = { StrategyStepType.RETRIEVAL: "Fetch facts from a DB", StrategyStepType.RELEVANCE: "Classify retrieved facts", - StrategyStepType.NONE: "Dummy step", StrategyStepType.PYTHON: "Custom python function", StrategyStepType.LLM: "Run a LLM", + StrategyStepType.NONE: "Dummy step", + StrategyStepType.SELECTION: "Select data", + StrategyStepType.QUERY_REPHRASING: "Rephrase query", + StrategyStepType.WEBSEARCH: "Search the web", + StrategyStepType.TRUNCATE_CONTEXT: "Truncate context", + StrategyStepType.HEADER: "Writing header", +} + +STEP_WHEN_TO_USE = { + StrategyStepType.RETRIEVAL: "When you want to retrieve facts from a database", + StrategyStepType.RELEVANCE: "When you want to classify retrieved facts", + StrategyStepType.PYTHON: "When you want to run a custom python function", + StrategyStepType.LLM: "When you want to run a LLM", + StrategyStepType.NONE: "Dummy step", + StrategyStepType.SELECTION: "When you want to select data", + StrategyStepType.QUERY_REPHRASING: "When you want to rephrase a query", + StrategyStepType.WEBSEARCH: "When you want to search the web", + StrategyStepType.TRUNCATE_CONTEXT: "When you want to truncate context", + StrategyStepType.HEADER: "When you want to write a header", +} + +STEP_PROGRESS_TEXTS = { + StrategyStepType.RETRIEVAL: "Retrieving facts", + StrategyStepType.RELEVANCE: "Classifying facts", + StrategyStepType.NONE: "Dummy step", + StrategyStepType.PYTHON: "Running custom python function", + StrategyStepType.LLM: "Running LLM", + StrategyStepType.SELECTION: "Selecting data", + StrategyStepType.QUERY_REPHRASING: "Rephrasing query", + StrategyStepType.WEBSEARCH: "Searching the web", + StrategyStepType.TRUNCATE_CONTEXT: "Truncating context", + StrategyStepType.HEADER: "Headline generation", } @@ -560,3 +608,34 @@ def try_parse_enum_value(string: str, enumType: Enum, raise_me: bool = True) -> raise ValueError(f"Invalid value {string} for enum {enumType}") return return parsed + + +class EmitType(Enum): + ANSWER = "ANSWER" + RETRIEVAL_RESULTS = "RETRIEVAL_RESULTS" + FOLLOW_UPS = "FOLLOW_UPS" + SELECTION = "SELECTION" + QUERY_REPHRASING = "QUERY_REPHRASING" + + +class CognitionLLMStepUsageType(Enum): + BASE = "BASE" + QUERY_REPHRASING = "QUERY_REPHRASING" + + +class StrategyComplexity(Enum): + SIMPLE = "SIMPLE" + MEDIUM = "MEDIUM" + COMPLEX = "COMPLEX" + + +class StrategyComplexityPrice(Enum): + SIMPLE = 0.25 + MEDIUM = 0.5 + COMPLEX = 0.75 + + +class ConsumptionLogState(Enum): + CREATED = "CREATED" + FINISHED = "FINISHED" + FAILED = "FAILED" diff --git a/models.py b/models.py index c4935711..41aad4d6 100644 --- a/models.py +++ b/models.py @@ -157,6 +157,7 @@ class User(Base): index=True, ) role = Column(String, default=UserRoles.ENGINEER.value) # enum UserRoles + language_display = Column(String, default="en") notifications = parent_to_child_relationship( Tablenames.USER, Tablenames.NOTIFICATION, @@ -1010,6 +1011,27 @@ class TaskQueue(Base): # --- COGNITION TABLES +class CognitionTeam(Base): + __tablename__ = Tablenames.TEAM.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + + organization_id = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.ORGANIZATION.value}.id", ondelete="CASCADE"), + index=True, + ) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + ) + + created_at = Column(DateTime, default=sql.func.now()) + name = Column(String) + description = Column(String) + + class CognitionProject(Base): __tablename__ = Tablenames.PROJECT.value __table_args__ = {"schema": "cognition"} @@ -1046,6 +1068,12 @@ class CognitionProject(Base): operator_routing_source_code = Column(String) wizard_running = Column(Boolean, default=False) refinery_synchronization_interval_option = Column(String) + interface_type = Column(String) + execute_query_enrichment_if_source_code = Column(String) + + customer_color_primary = Column(String, default="#18181b") + customer_color_primary_only_accent = Column(Boolean, default=False) + customer_color_secondary = Column(String, default="#9333ea") class CognitionStrategy(Base): @@ -1093,6 +1121,9 @@ class CognitionStrategyStep(Base): step_type = Column(String) position = Column(Integer) config = Column(JSON) + progress_text = Column(String) + enable_emissions = Column(Boolean, default=True) + execute_if_source_code = Column(String) class CognitionConversation(Base): @@ -1111,6 +1142,12 @@ class CognitionConversation(Base): ) created_at = Column(DateTime, default=sql.func.now()) scope_dict = Column(JSON) + header = Column(String) + error = Column(String) + + # synopsis-specific + synopsis_spreadsheet_row_id = Column(String) + synopsis_column = Column(String) class CognitionMessage(Base): @@ -1140,10 +1177,11 @@ class CognitionMessage(Base): created_at = Column(DateTime, default=sql.func.now()) question = Column(String) facts = Column(ARRAY(JSON)) + selection_widget = Column(ARRAY(JSON)) answer = Column(String) - # None = not yet answered, True = positive, false = negative - positive_feedback = Column(Boolean) + feedback_value = Column(String) + feedback_category = Column(String) feedback_message = Column(String) scope_dict_diff_previous_conversation = Column(JSON) @@ -1184,6 +1222,7 @@ class CognitionPipelineLogs(Base): record_dict_diff_previous_message = Column(JSON) content = Column(ARRAY(String)) time_elapsed = Column(Float) + skipped_step = Column(Boolean, default=False) class CognitionEnvironmentVariable(Base): @@ -1344,3 +1383,139 @@ class CognitionRefinerySynchronizationTask(Base): state = Column(String) # e.g. CREATED, FINISHED, FAILED logs = Column(ARRAY(String)) num_records_created = Column(Integer) + + +class CognitionConsumptionLog(Base): + __tablename__ = Tablenames.CONSUMPTION_LOG.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + project_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.PROJECT.value}.id", ondelete="CASCADE"), + index=True, + ) + strategy_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.STRATEGY.value}.id", ondelete="CASCADE"), + index=True, + ) + conversation_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.CONVERSATION.value}.id", ondelete="CASCADE"), + index=True, + ) + message_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.MESSAGE.value}.id", ondelete="CASCADE"), + index=True, + ) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + ) + created_at = Column(DateTime, default=sql.func.now()) + complexity = Column(String) + state = Column(String) + + +class CognitionTeamMember(Base): + __tablename__ = Tablenames.TEAM_MEMBER.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + + team_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.TEAM.value}.id", ondelete="CASCADE"), + index=True, + ) + + user_id = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="CASCADE"), + index=True, + ) + + created_at = Column(DateTime, default=sql.func.now()) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + ) + + +class CognitionTeamProjectAccess(Base): + __tablename__ = Tablenames.TEAM_PROJECT_ACCESS.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + + team_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.TEAM.value}.id", ondelete="CASCADE"), + index=True, + ) + + project_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.PROJECT.value}.id", ondelete="CASCADE"), + index=True, + ) + + created_at = Column(DateTime, default=sql.func.now()) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + ) + + +class CognitionSynopsisSpreadsheet(Base): + __tablename__ = Tablenames.SYNOPSIS_SPREADSHEET.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + dataset_id = Column( + UUID(as_uuid=True), + ForeignKey( + f"cognition.{Tablenames.MARKDOWN_DATASET.value}.id", ondelete="CASCADE" + ), + index=True, + ) + project_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.PROJECT.value}.id", ondelete="CASCADE"), + index=True, + ) + name = Column(String) + synopsis_type = Column(String) + task_scope_dict = Column(JSON) + filter_attribute_name = Column(String) + created_at = Column(DateTime, default=sql.func.now()) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + ) + + +class CognitionSynopsisSpreadsheetRow(Base): + __tablename__ = Tablenames.SYNOPSIS_SPREADSHEET_ROW.value + __table_args__ = {"schema": "cognition"} + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + project_id = Column( + UUID(as_uuid=True), + ForeignKey(f"cognition.{Tablenames.PROJECT.value}.id", ondelete="CASCADE"), + index=True, + ) + spreadsheet_id = Column( + UUID(as_uuid=True), + ForeignKey( + f"cognition.{Tablenames.SYNOPSIS_SPREADSHEET.value}.id", ondelete="CASCADE" + ), + index=True, + ) + created_at = Column(DateTime, default=sql.func.now()) + created_by = Column( + UUID(as_uuid=True), + ForeignKey(f"{Tablenames.USER.value}.id", ondelete="SET NULL"), + index=True, + )