ckit_integrations_db.py

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Awaitable, Callable

from flexus_client_kit import ckit_ask_model, ckit_bot_exec, ckit_cloudtool


@dataclass
class IntegrationRecord:
    integr_name: str
    integr_tools: list[ckit_cloudtool.CloudTool]
    integr_init: Callable[..., Awaitable[Any]]
    integr_setup_handlers: Callable
    integr_provider: str = ""
    integr_scopes: list[str] = field(default_factory=list)
    integr_prompt: str = ""
    integr_is_messenger: bool = False
    integr_need_mongo: bool = False


@dataclass
class _MongoStoreState:
    collection: Any = None


def static_integrations_load(bot_dir: Path, allowlist: list[str], builtin_skills: list[str]) -> list[IntegrationRecord]:
    # static means designed to save into constant on top level of a bot file
    # logger is not yet initilized here, no logs possible
    result = []
    for name in allowlist:
        if name == "skills":
            if len(builtin_skills) > 0:
                from flexus_client_kit import ckit_skills
                async def _init_skills(rcx, setup):
                    return None
                result.append(IntegrationRecord(
                    integr_name=name,
                    integr_tools=[ckit_skills.FETCH_SKILL_TOOL],
                    integr_init=_init_skills,
                    integr_setup_handlers=lambda obj, rcx, _d=bot_dir, _s=builtin_skills: [
                        rcx.on_tool_call("flexus_fetch_skill")(lambda tc, args: ckit_skills.called_by_model(tc, args, _d, _s))
                    ],
                ))

        elif name == "flexus_policy_document":
            from flexus_client_kit.integrations import fi_pdoc
            async def _init_pdoc(rcx, setup):
                return fi_pdoc.IntegrationPdoc(rcx, rcx.persona.ws_root_group_id)
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_pdoc.POLICY_DOCUMENT_TOOL],
                integr_init=_init_pdoc,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("flexus_policy_document")(obj.called_by_model)],
                integr_prompt=fi_pdoc.POLICY_DOCUMENT_PROMPT,
            ))

        elif name == "print_widget":
            from flexus_client_kit.integrations import fi_widget
            async def _init_widget(rcx, setup):
                return None
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_widget.PRINT_WIDGET_TOOL],
                integr_init=_init_widget,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("print_widget")(fi_widget.handle_print_widget)],
                integr_prompt=fi_widget.PRINT_WIDGET_PROMPT,
            ))

        elif name == "gmail":
            from flexus_client_kit.integrations import fi_gmail
            async def _init_gmail(rcx, setup):
                return fi_gmail.IntegrationGmail(rcx.fclient, rcx)
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_gmail.GMAIL_TOOL],
                integr_init=_init_gmail,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("gmail")(obj.called_by_model)],
                integr_provider="google",
                integr_scopes=fi_gmail.GMAIL_SCOPES,
                integr_prompt=fi_gmail.GMAIL_PROMPT,
            ))

        elif name == "google_calendar":
            from flexus_client_kit.integrations import fi_google_calendar
            async def _init_gcal(rcx, setup):
                return fi_google_calendar.IntegrationGoogleCalendar(rcx.fclient, rcx)
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_google_calendar.GOOGLE_CALENDAR_TOOL],
                integr_init=_init_gcal,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("google_calendar")(obj.called_by_model)],
                integr_provider="google",
                integr_scopes=fi_google_calendar.REQUIRED_SCOPES,
                integr_prompt="",
            ))

        elif name == "jira":
            from flexus_client_kit.integrations import fi_jira
            async def _init_jira(rcx, setup):
                url = (setup or {}).get("jira_instance_url", "")
                return fi_jira.IntegrationJira(rcx.fclient, rcx, jira_instance_url=url)
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_jira.JIRA_TOOL],
                integr_init=_init_jira,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("jira")(obj.called_by_model)],
                integr_provider="atlassian",
                integr_scopes=fi_jira.REQUIRED_SCOPES,
                integr_prompt="",
            ))

        elif name.startswith("facebook"):   # "facebook[account, adset]"
            # In manifest.json, it's not possible to write [] brackets (according to the json schema), but in a
            # bot defined by code, you can
            from flexus_client_kit.integrations import fi_facebook2
            fb_bunch = fi_facebook2.make_facebook_bunch(_parse_bracket_list(name))
            fb_tool = fb_bunch.make_tool()
            async def _init_facebook(rcx, setup, _bunch=fb_bunch):
                return fi_facebook2.IntegrationFacebook2(rcx.fclient, rcx, _bunch)
            result.append(IntegrationRecord(
                integr_name="facebook",
                integr_tools=[fb_tool],
                integr_init=_init_facebook,
                integr_setup_handlers=lambda obj, rcx, _t=fb_tool: [rcx.on_tool_call(_t.name)(obj.called_by_model)],
                integr_provider="facebook",
                integr_prompt="",
            ))

        elif name == "linkedin":
            from flexus_client_kit.integrations import fi_linkedin
            async def _init_linkedin(rcx, setup):
                ad_account_id = (setup or {}).get("ad_account_id", "")
                return fi_linkedin.IntegrationLinkedIn(
                    rcx.fclient, rcx, ad_account_id=ad_account_id
                )
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_linkedin.LINKEDIN_TOOL],
                integr_init=_init_linkedin,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("linkedin")(obj.called_by_model)],
                integr_provider="linkedin",
                integr_scopes=[
                    "r_profile_basicinfo",
                    "email",
                    "w_member_social",
                ],
                integr_prompt="",
            ))

        elif name == "github":
            from flexus_client_kit.integrations import fi_github
            async def _init_github(rcx, setup):
                return fi_github.IntegrationGitHub(rcx.fclient, rcx)
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_github.GITHUB_TOOL],
                integr_init=_init_github,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("github")(obj.called_by_model)],
                integr_provider="github",
                integr_prompt="",
            ))

        elif name == "slack":
            from flexus_client_kit.integrations import fi_slack, fi_messenger
            async def _init_slack(rcx, setup):
                bot_name = (setup or {}).get("slack_bot_name", "") or rcx.persona.persona_name
                bot_icon_url = (setup or {}).get("slack_bot_icon_url", "")
                obj = fi_slack.IntegrationSlack(rcx.fclient, rcx, bot_name=bot_name, bot_icon_url=bot_icon_url)
                await obj.load_workspace_maps()
                return obj
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_slack.SLACK_TOOL],
                integr_init=_init_slack,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("slack")(obj.called_by_model)],
                integr_provider="slack",
                integr_is_messenger=True,
                integr_scopes=[
                    "channels:read",
                    "chat:write",
                    "chat:write.customize",
                    "files:read",
                    "users:read",
                    "im:read",
                ],
                integr_prompt=fi_messenger.MESSENGER_PROMPT,
            ))

        elif name == "telegram":
            from flexus_client_kit.integrations import fi_telegram, fi_messenger
            async def _init_telegram(rcx, setup):
                obj = fi_telegram.IntegrationTelegram(rcx.fclient, rcx)
                await obj.initialize()
                return obj
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_telegram.TELEGRAM_TOOL],
                integr_init=_init_telegram,
                integr_setup_handlers=lambda obj, rcx: [rcx.on_tool_call("telegram")(obj.called_by_model)],
                integr_provider="telegram",
                integr_is_messenger=True,
                integr_prompt=fi_messenger.MESSENGER_PROMPT,
            ))

        elif name.startswith("erp"):   # "erp[meta, data]" or "erp[meta, data, crud, csv_import]"
            from flexus_client_kit.integrations import fi_erp
            subset = _parse_bracket_list(name)
            tool_map = {
                "meta": (fi_erp.ERP_TABLE_META_TOOL, "handle_erp_meta"),
                "data": (fi_erp.ERP_TABLE_DATA_TOOL, "handle_erp_data"),
                "crud": (fi_erp.ERP_TABLE_CRUD_TOOL, "handle_erp_crud"),
                "csv_import": (fi_erp.ERP_CSV_IMPORT_TOOL, "handle_csv_import"),
            }
            if subset is None:
                subset = list(tool_map.keys())
            tools_and_methods = [(tool_map[s][0], tool_map[s][1]) for s in subset]
            async def _init_erp(rcx, setup):
                return fi_erp.IntegrationErp(rcx.fclient, rcx.persona.ws_id, rcx.personal_mongo)
            def _setup_erp(obj, rcx, _tam=tools_and_methods):
                for tool, method_name in _tam:
                    rcx.on_tool_call(tool.name)(getattr(obj, method_name))
            result.append(IntegrationRecord(
                integr_name="erp",
                integr_tools=[t for t, _ in tools_and_methods],
                integr_init=_init_erp,
                integr_setup_handlers=_setup_erp,
                integr_need_mongo=True,
            ))

        elif name.startswith("crm"):   # "crm[manage_contact, manage_deal, log_activity]"
            from flexus_client_kit.integrations import fi_crm
            subset = _parse_bracket_list(name)
            tool_map = {
                "manage_contact": (fi_crm.MANAGE_CRM_CONTACT_TOOL, "handle_manage_crm_contact"),
                "manage_deal": (fi_crm.MANAGE_CRM_DEAL_TOOL, "handle_manage_crm_deal"),
                "log_activity": (fi_crm.LOG_CRM_ACTIVITY_TOOL, "handle_log_crm_activity"),
            }
            if subset is None:
                subset = list(tool_map.keys())
            tools_and_methods = [(tool_map[s][0], tool_map[s][1]) for s in subset]
            async def _init_crm(rcx, setup):
                return fi_crm.IntegrationCrm(rcx.fclient, rcx.persona.ws_id)
            def _setup_crm(obj, rcx, _tam=tools_and_methods):
                for tool, method_name in _tam:
                    rcx.on_tool_call(tool.name)(getattr(obj, method_name))
            result.append(IntegrationRecord(
                integr_name="crm",
                integr_tools=[t for t, _ in tools_and_methods],
                integr_init=_init_crm,
                integr_setup_handlers=_setup_crm,
                integr_prompt=fi_crm.LOG_CRM_ACTIVITIES_PROMPT if "log_activity" in subset else "",
            ))

        elif name == "mongo_store":
            from flexus_client_kit.integrations import fi_mongo_store
            from flexus_client_kit import ckit_mongo
            state = _MongoStoreState()
            async def _init_mongo_store(rcx, setup, _state=state):
                from pymongo import AsyncMongoClient
                mongo_conn_str = await ckit_mongo.mongo_fetch_creds(rcx.fclient, rcx.persona.persona_id)
                mongo = AsyncMongoClient(mongo_conn_str)
                dbname = rcx.persona.persona_id + "_db"
                _state.collection = mongo[dbname]["personal_mongo"]
                return _state
            result.append(IntegrationRecord(
                integr_name=name,
                integr_tools=[fi_mongo_store.MONGO_STORE_TOOL],
                integr_init=_init_mongo_store,
                integr_setup_handlers=lambda obj, rcx, _s=state: [
                    rcx.on_tool_call("mongo_store")(
                        lambda tc, args: fi_mongo_store.handle_mongo_store(rcx.workdir, _s.collection, tc, args)
                    )
                ],
                integr_prompt="",
            ))

        else:
            raise ValueError(f"Unknown integration {name!r}")
    return result


def _parse_bracket_list(name: str) -> list[str] | None:
    if "[" not in name:
        return None
    return [g.strip() for g in name.split("[", 1)[1].rstrip("]").split(",")]


async def main_loop_integrations_init(records: list[IntegrationRecord], rcx: ckit_bot_exec.RobotContext, setup: dict) -> dict[str, Any]:
    from flexus_client_kit.integrations import fi_messenger
    if any(rec.integr_need_mongo for rec in records) and rcx.personal_mongo is None:
        from pymongo import AsyncMongoClient
        from flexus_client_kit import ckit_mongo
        mongo_conn_str = await ckit_mongo.mongo_fetch_creds(rcx.fclient, rcx.persona.persona_id)
        mongo = AsyncMongoClient(mongo_conn_str)
        rcx.personal_mongo = mongo[rcx.persona.persona_id + "_db"]["personal_mongo"]
    result = {}
    for rec in records:
        obj = await rec.integr_init(rcx, setup)
        rec.integr_setup_handlers(obj, rcx)
        result[rec.integr_name] = obj
        if rec.integr_is_messenger:
            assert isinstance(obj, fi_messenger.FlexusMessenger), f"{rec.integr_name} has integr_is_messenger=True but {type(obj).__name__} doesn't inherit FlexusMessenger"
            rcx.messengers.append(obj)
            @rcx.on_emessage(obj.emessage_type)
            async def _emessage_handler(emsg, _m=obj):
                await _m.handle_emessage(emsg)

    if rcx.messengers:
        @rcx.on_updated_message
        async def _messenger_updated_message(msg: ckit_ask_model.FThreadMessageOutput):
            # Don't worry, you can override it. The default reaction to assistant messages is to get it past messengers:
            for m in rcx.messengers:
                await m.look_assistant_might_have_posted_something(msg)

    return result