diff --git a/flexus_client_kit/ckit_bot_exec.py b/flexus_client_kit/ckit_bot_exec.py index e22cfcc7..68e77272 100644 --- a/flexus_client_kit/ckit_bot_exec.py +++ b/flexus_client_kit/ckit_bot_exec.py @@ -66,7 +66,7 @@ def official_setup_mixing_procedure(marketable_setup_default, persona_setup) -> class RobotContext: - def __init__(self, fclient: ckit_client.FlexusClient, p: ckit_bot_query.FPersonaOutput): + def __init__(self, fclient: ckit_client.FlexusClient, p: ckit_bot_query.FPersonaOutput, external_auth: Optional[Dict[str, Any]] = None): self._handler_updated_message: Optional[Callable[[ckit_ask_model.FThreadMessageOutput], Awaitable[None]]] = None self._handler_upd_thread: Optional[Callable[[ckit_ask_model.FThreadOutput], Awaitable[None]]] = None self._handler_updated_task: Optional[Callable[[ckit_kanban.FPersonaKanbanTaskOutput], Awaitable[None]]] = None @@ -92,6 +92,7 @@ def __init__(self, fclient: ckit_client.FlexusClient, p: ckit_bot_query.FPersona self.workdir = "/tmp/bot_workspace/%s/" % p.persona_id self.running_test_scenario = False self.running_happy_yaml = "" + self.external_auth = external_auth or {} os.makedirs(self.workdir, exist_ok=True) def on_updated_message(self, handler: Callable[[ckit_ask_model.FThreadMessageOutput], Awaitable[None]]): @@ -366,6 +367,7 @@ def __init__( self.running_happy_yaml = running_happy_yaml self.subscribe_to_erp_tables = subscribe_to_erp_tables self.subscribe_to_emsg_types = subscribe_to_emsg_types + self.auth: Dict[str, Dict[str, Any]] = {} async def subscribe_and_produce_callbacks( @@ -408,23 +410,52 @@ async def subscribe_and_produce_callbacks( reassign_threads = False # logger.info("subs %s %s %s" % (upd.news_action, upd.news_about, upd.news_payload_id)) - if upd.news_about == "flexus_persona": + if upd.news_about == "flexus_external_auth": + handled = True + if upd.news_action in ["INSERT", "UPDATE"]: + if upd.news_payload_auth is None: + continue + + persona_id = upd.news_payload_auth.auth_persona_id + ws_id = upd.news_payload_auth.ws_id + provider = upd.news_payload_auth.auth_service_provider + is_workspace_token = not persona_id or persona_id == "" + + logger.info(f"Received auth: persona_id={persona_id}, ws_id={ws_id}, provider={provider}, keys={list(upd.news_payload_auth.auth_key2value.keys())}") + + # Store auth tokens (use persona_id for persona-scoped, ws_id for workspace-scoped) + auth_key = persona_id if not is_workspace_token else ws_id + if auth_key not in bc.auth: + bc.auth[auth_key] = {} + bc.auth[auth_key][provider] = upd.news_payload_auth.auth_key2value + logger.info(f"Stored in bc.auth[{auth_key}][{provider}] with {len(upd.news_payload_auth.auth_key2value)} keys") + bc.bots_running[persona_id].instance_rcx._restart_requested = True + + elif upd.news_action == "DELETE": + if upd.news_payload_auth is None: + continue + + persona_id = upd.news_payload_auth.auth_persona_id + ws_id = upd.news_payload_auth.ws_id + provider = upd.news_payload_auth.auth_service_provider + + # Remove auth tokens + auth_key = persona_id if persona_id else ws_id + if auth_key in bc.auth: + bc.auth[auth_key].pop(provider, None) + logger.info(f"Removed auth {provider} from bc.auth[{auth_key}]") + bc.bots_running[persona_id].instance_rcx._restart_requested = True + + + elif upd.news_about == "flexus_persona": if upd.news_action in ["INSERT", "UPDATE"]: assert upd.news_payload_persona.ws_id assert upd.news_payload_persona.ws_timezone handled = True persona_id = upd.news_payload_id - if bot := bc.bots_running.get(persona_id, None): - if bot.instance_rcx.persona.persona_setup != upd.news_payload_persona.persona_setup: - logger.info("Persona %s setup changed, requesting graceful shutdown" % persona_id) - del bc.bots_running[persona_id] - bc.shutting_down_tasks.add(bot.atask) - bot.atask.add_done_callback(bc.shutting_down_tasks.discard) - bot.instance_rcx._restart_requested = True - bot.instance_rcx._parked_anything_new.set() if persona_id not in bc.bots_running: - rcx = RobotContext(fclient, upd.news_payload_persona) + rcx = RobotContext(fclient, upd.news_payload_persona, bc.auth.get(persona_id, {})) rcx.running_test_scenario = bc.running_test_scenario rcx.running_happy_yaml = bc.running_happy_yaml bc.bots_running[persona_id] = BotInstance( @@ -432,7 +463,6 @@ async def subscribe_and_produce_callbacks( atask=asyncio.create_task(crash_boom_bang(fclient, rcx, bc.bot_main_loop)), instance_rcx=rcx, ) - reassign_threads = True elif upd.news_action == "DELETE": handled = True @@ -538,6 +568,8 @@ async def subscribe_and_produce_callbacks( if bot := bc.bots_running.get(emsg.emsg_persona_id): bot.instance_rcx._parked_emessages[emsg.emsg_id] = emsg bot.instance_rcx._parked_anything_new.set() + else: + logger.warning("External message about persona %s, but no bot is running it." % emsg.emsg_persona_id) elif upd.news_action == "INITIAL_UPDATES_OVER": if len(bc.bots_running) == 0: diff --git a/flexus_client_kit/ckit_bot_install.py b/flexus_client_kit/ckit_bot_install.py index 3901c198..f885126d 100644 --- a/flexus_client_kit/ckit_bot_install.py +++ b/flexus_client_kit/ckit_bot_install.py @@ -78,6 +78,8 @@ async def marketplace_upsert_dev_bot( marketable_tags: List[str] = [], marketable_forms: Optional[Dict[str, str]] = None, marketable_required_policydocs: List[str] = [], + marketable_auth_needed: List[str] = [], + marketable_auth_supported: List[str] = [], ) -> FBotInstallOutput: assert not ws_id.startswith("fx-"), "You can find workspace id in the browser address bar, when visiting for example the statistics page" http = await client.use_http() @@ -89,7 +91,7 @@ async def marketplace_upsert_dev_bot( experts_input.append(expert_dict) # NOTE: marketable_stage removed from mutation for staging API compatibility r = await h.execute( - gql.gql(f"""mutation InstallBot($ws: String!, $name: String!, $ver: String!, $title1: String!, $title2: String!, $author: String!, $accent_color: String!, $occupation: String!, $desc: String!, $typical_group: String!, $repo: String!, $run: String!, $setup: String!, $featured: [FFeaturedActionInput!]!, $intro: String!, $model: String!, $daily: Int!, $inbox: Int!, $experts: [FMarketplaceExpertInput!]!, $schedule: String!, $big: String!, $small: String!, $tags: [String!]!, $forms: String, $required_policydocs: [String!]!) {{ + gql.gql(f"""mutation InstallBot($ws: String!, $name: String!, $ver: String!, $title1: String!, $title2: String!, $author: String!, $accent_color: String!, $occupation: String!, $desc: String!, $typical_group: String!, $repo: String!, $run: String!, $setup: String!, $featured: [FFeaturedActionInput!]!, $intro: String!, $model: String!, $daily: Int!, $inbox: Int!, $experts: [FMarketplaceExpertInput!]!, $schedule: String!, $big: String!, $small: String!, $tags: [String!]!, $forms: String, $required_policydocs: [String!]!, $auth_needed: [String!]!, $auth_supported: [String!]!) {{ marketplace_upsert_dev_bot( ws_id: $ws, marketable_name: $name, @@ -115,7 +117,9 @@ async def marketplace_upsert_dev_bot( marketable_picture_small_b64: $small, marketable_tags: $tags, marketable_forms: $forms, - marketable_required_policydocs: $required_policydocs + marketable_required_policydocs: $required_policydocs, + marketable_auth_needed: $auth_needed, + marketable_auth_supported: $auth_supported ) {{ {gql_utils.gql_fields(FBotInstallOutput)} }} @@ -146,6 +150,8 @@ async def marketplace_upsert_dev_bot( "small": marketable_picture_small_b64, "forms": json.dumps(marketable_forms or {}), "required_policydocs": marketable_required_policydocs, + "auth_needed": marketable_auth_needed, + "auth_supported": marketable_auth_supported, }, ) return gql_utils.dataclass_from_dict(r["marketplace_upsert_dev_bot"], FBotInstallOutput) diff --git a/flexus_client_kit/ckit_bot_query.py b/flexus_client_kit/ckit_bot_query.py index 8803c6ca..2e1f695f 100644 --- a/flexus_client_kit/ckit_bot_query.py +++ b/flexus_client_kit/ckit_bot_query.py @@ -25,6 +25,8 @@ class FPersonaOutput: marketable_run_this: Optional[str] = None marketable_stage: Optional[str] = None marketable_radix: Optional[int] = None + marketable_auth_needed: Optional[List[str]] = None + marketable_auth_supported: Optional[List[str]] = None @dataclass @@ -48,6 +50,14 @@ class FExternalMessageOutput: emsg_created_ts: float ws_id: str +@dataclass +class FExternalAuth: + auth_id: str + auth_persona_id: str + auth_auth_type: str + auth_service_provider: str + auth_key2value: Any # {"TOKEN_ID": "TOKEN_VAL", ...} + ws_id: str @dataclass class FBotThreadsCallsTasks: @@ -62,7 +72,7 @@ class FBotThreadsCallsTasks: news_payload_erp_record_new: Optional[Dict[str, Any]] = None news_payload_erp_record_old: Optional[Dict[str, Any]] = None news_payload_emessage: Optional[FExternalMessageOutput] = None - + news_payload_auth: Optional[FExternalAuth] = None @dataclass class FThreadWithMessages: diff --git a/flexus_client_kit/ckit_external_auth.py b/flexus_client_kit/ckit_external_auth.py index 53451d81..9ac901f5 100644 --- a/flexus_client_kit/ckit_external_auth.py +++ b/flexus_client_kit/ckit_external_auth.py @@ -32,7 +32,6 @@ async def upsert_external_auth( fclient: ckit_client.FlexusClient, persona_id: str, auth_searchable: str, - auth_name: str, auth_service_provider: str, auth_json: dict, ) -> None: @@ -43,14 +42,12 @@ async def upsert_external_auth( mutation UpsertExternalAuth( $persona_id: String!, $auth_searchable: String!, - $auth_name: String!, $auth_service_provider: String!, $auth_json: String! ) {{ upsert_external_auth( persona_id: $persona_id, auth_searchable: $auth_searchable, - auth_name: $auth_name, auth_service_provider: $auth_service_provider, auth_json: $auth_json ) @@ -58,7 +55,6 @@ async def upsert_external_auth( variable_values={ "persona_id": persona_id, "auth_searchable": auth_searchable, - "auth_name": auth_name, "auth_service_provider": auth_service_provider, "auth_json": json.dumps(auth_json), }, diff --git a/flexus_simple_bots/frog/frog_bot.py b/flexus_simple_bots/frog/frog_bot.py index fcaf2902..48770f1c 100644 --- a/flexus_simple_bots/frog/frog_bot.py +++ b/flexus_simple_bots/frog/frog_bot.py @@ -98,6 +98,13 @@ async def frog_main_loop(fclient: ckit_client.FlexusClient, rcx: ckit_bot_exec.RobotContext) -> None: setup = ckit_bot_exec.official_setup_mixing_procedure(frog_install.frog_setup_schema, rcx.persona.persona_setup) + # Example of getting token from external_auth + elevenlabs_auth = rcx.external_auth.get("elevenlabs", {}) + elevenlabs_token = elevenlabs_auth.get("token") if elevenlabs_auth else None + if elevenlabs_token: + masked = elevenlabs_token[:8] + "..." + elevenlabs_token[-4:] if len(elevenlabs_token) > 12 else "***" + logger.info("Frog has ELEVENLABS token: %s (len=%d)", masked, len(elevenlabs_token)) + mongo_conn_str = await ckit_mongo.mongo_fetch_creds(fclient, rcx.persona.persona_id) mongo = AsyncMongoClient(mongo_conn_str) dbname = rcx.persona.persona_id + "_db" diff --git a/flexus_simple_bots/frog/frog_install.py b/flexus_simple_bots/frog/frog_install.py index c830bad7..f23a888a 100644 --- a/flexus_simple_bots/frog/frog_install.py +++ b/flexus_simple_bots/frog/frog_install.py @@ -104,6 +104,8 @@ async def install( marketable_github_repo="https://github.com/smallcloudai/flexus-client-kit.git", marketable_run_this="python -m flexus_simple_bots.frog.frog_bot", marketable_setup_default=frog_setup_schema, + marketable_auth_needed=[], + marketable_auth_supported=["elevenlabs"], marketable_featured_actions=[ {"feat_question": "Ribbit! Tell me something fun", "feat_expert": "default", "feat_depends_on_setup": []}, {"feat_question": "Give me a motivational boost", "feat_expert": "default", "feat_depends_on_setup": []},