"""AI domain — business logic.""" from __future__ import annotations import re from typing import Optional from uuid import UUID import httpx import structlog from src.config import get_settings from src.core.exceptions import ConflictError, NotFoundError from src.core.security.encryption import decrypt_value, encrypt_value from src.domain.ai.models import ( ProjectAIConfig, ProjectAIEmbedConfig, WorkspaceVectorizeConfig, ) from src.domain.ai.repository import ( ProjectAIConfigRepository, ProjectAIEmbedConfigRepository, WorkspaceVectorizeConfigRepository, ) from src.domain.ai.schemas import ( AIConfigResponse, CreateEmbedConfigRequest, EmbedConfigResponse, ProvisionVectorizeRequest, ReembedTriggerResponse, SetBYOKCompletionKeyRequest, UpdateAIConfigRequest, UpdateEmbedConfigRequest, VectorizeConfigResponse, ) from src.integrations.cloud_control_client import CloudControlClient logger = structlog.get_logger() _CF_API = "https://api.cloudflare.com/client/v4" # Embedding model resolved from plan _PLAN_EMBEDDING: dict[str, tuple[str, str, int]] = { # plan_slug -> (provider, model, dimensions) "free": ("cf_workers_ai", "@cf/baai/bge-base-en-v1.5", 768), "starter": ("cf_workers_ai", "@cf/baai/bge-base-en-v1.5", 768), "pro": ("openai", "text-embedding-3-small", 1536), "growth": ("openai", "text-embedding-3-large", 3072), } _DEFAULT_EMBEDDING = _PLAN_EMBEDDING["free"] # Completion providers by plan _PLAN_COMPLETION_PROVIDERS: dict[str, list[str]] = { "free": ["cf_workers_ai"], "starter": ["cf_workers_ai", "openai_mini"], "pro": ["cf_workers_ai", "openai_mini", "openai", "claude_haiku"], "growth": ["cf_workers_ai", "openai_mini", "openai", "claude_haiku", "claude_sonnet", "claude_opus"], } class AIService: def __init__( self, *, config_repo: ProjectAIConfigRepository, embed_repo: ProjectAIEmbedConfigRepository, vectorize_repo: WorkspaceVectorizeConfigRepository, cloud_client: CloudControlClient, ) -> None: self._config = config_repo self._embeds = embed_repo self._vectorize = vectorize_repo self._cloud = cloud_client # ------------------------------------------------------------------ # AI Config # ------------------------------------------------------------------ async def get_or_create_config( self, project_id: UUID, plan_slug: Optional[str] = None ) -> AIConfigResponse: cfg = await self._config.get_by_project(project_id) if not cfg: cfg = await self._config.create( project_id=project_id, **_default_config_fields(plan_slug), ) return _config_response(cfg) async def update_config( self, project_id: UUID, body: UpdateAIConfigRequest, plan_slug: Optional[str] = None, ) -> AIConfigResponse: cfg = await self._config.get_by_project(project_id) if not cfg: cfg = await self._config.create( project_id=project_id, **_default_config_fields(plan_slug), ) updates: dict = {} if body.embedding_provider is not None: if body.embedding_provider != "byok" and body.embedding_provider != cfg.embedding_provider: # Derive model from plan for non-BYOK providers provider, model, _ = _plan_embedding(plan_slug) updates["embedding_provider"] = body.embedding_provider updates["embedding_model"] = model else: updates["embedding_provider"] = body.embedding_provider if body.embedding_byok_key is not None: secret_ref = f"ai:embedding:{project_id}" # Store encrypted key; ref is deterministic so rotation is safe updates["embedding_secret_ref"] = secret_ref # Persist encrypted key via a project secret row (reuse existing mechanism) _store_byok_embedding_key(str(project_id), body.embedding_byok_key) if updates: cfg = await self._config.update(cfg, **updates) await self._push_ai_config(project_id, cfg) logger.info("ai_config_updated", project_id=str(project_id)) return _config_response(cfg) async def set_byok_completion_key( self, project_id: UUID, body: SetBYOKCompletionKeyRequest, ) -> AIConfigResponse: cfg = await self._config.get_by_project(project_id) if not cfg: cfg = await self._config.create(project_id=project_id) secret_ref = f"ai:completion:{project_id}:{body.provider}" _store_byok_completion_key(str(project_id), body.provider, body.api_key) refs = dict(cfg.completion_secret_refs_json or {}) refs[body.provider] = secret_ref cfg = await self._config.update(cfg, completion_secret_refs_json=refs) await self._push_ai_config(project_id, cfg) logger.info( "ai_byok_completion_key_set", project_id=str(project_id), provider=body.provider, ) return _config_response(cfg) # ------------------------------------------------------------------ # Embed configs # ------------------------------------------------------------------ async def list_embed_configs(self, project_id: UUID) -> list[EmbedConfigResponse]: rows = await self._embeds.list_for_project(project_id) return [_embed_response(r) for r in rows] async def create_embed_config( self, project_id: UUID, body: CreateEmbedConfigRequest ) -> EmbedConfigResponse: existing = await self._embeds.get_by_entity_field( project_id, body.entity_key, body.field_key ) if existing: raise ConflictError( f"Embed config for {body.entity_key}.{body.field_key} already exists" ) row = await self._embeds.create( project_id=project_id, entity_key=body.entity_key, field_key=body.field_key, auto_embed=body.auto_embed, model_override=body.model_override, ) await self._push_embed_configs(project_id) logger.info( "ai_embed_config_created", project_id=str(project_id), entity=body.entity_key, field=body.field_key, ) return _embed_response(row) async def update_embed_config( self, project_id: UUID, entity_key: str, field_key: str, body: UpdateEmbedConfigRequest, ) -> EmbedConfigResponse: row = await self._embeds.get_by_entity_field(project_id, entity_key, field_key) if not row: raise NotFoundError(f"Embed config for {entity_key}.{field_key} not found") updates: dict = {} if body.auto_embed is not None: updates["auto_embed"] = body.auto_embed if body.model_override is not None: updates["model_override"] = body.model_override if updates: row = await self._embeds.update(row, **updates) await self._push_embed_configs(project_id) return _embed_response(row) async def delete_embed_config( self, project_id: UUID, entity_key: str, field_key: str ) -> None: row = await self._embeds.get_by_entity_field(project_id, entity_key, field_key) if not row: raise NotFoundError(f"Embed config for {entity_key}.{field_key} not found") await self._embeds.delete(row) await self._push_embed_configs(project_id) logger.info( "ai_embed_config_deleted", project_id=str(project_id), entity=entity_key, field=field_key, ) # ------------------------------------------------------------------ # Vectorize index provisioning # ------------------------------------------------------------------ async def get_vectorize_config( self, workspace_id: UUID ) -> Optional[VectorizeConfigResponse]: row = await self._vectorize.get_by_workspace(workspace_id) if not row: return None return _vectorize_response(row) async def provision_vectorize( self, workspace_id: UUID, body: ProvisionVectorizeRequest, ) -> VectorizeConfigResponse: existing = await self._vectorize.get_by_workspace(workspace_id) if existing: raise ConflictError("Vectorize index already provisioned for this workspace") index_name = f"p2hs-{str(workspace_id)[:8]}" index_id = await _create_vectorize_index( name=index_name, dimensions=body.dimensions, metric=body.metric, ) row = await self._vectorize.create( workspace_id=workspace_id, vectorize_index_id=index_id, index_name=index_name, dimensions=body.dimensions, metric=body.metric, ) await self._cloud.push_vectorize_config( workspace_id=str(workspace_id), vectorize_index_id=index_id, index_name=index_name, dimensions=body.dimensions, metric=body.metric, ) logger.info( "vectorize_index_provisioned", workspace_id=str(workspace_id), index_id=index_id, ) return _vectorize_response(row) async def trigger_reembed(self, project_id: UUID) -> ReembedTriggerResponse: """Queue a bulk re-embed job after a plan upgrade (incompatible dimensions).""" # Queuing via arq worker task — the actual job runs asynchronously from src.workers.tasks.ai_reembed import enqueue_reembed await enqueue_reembed(project_id=str(project_id)) logger.info("ai_reembed_queued", project_id=str(project_id)) return ReembedTriggerResponse( queued=True, message="Re-embed job queued. Semantic search will be updated as vectors are rebuilt.", ) # ------------------------------------------------------------------ # Internal # ------------------------------------------------------------------ async def _push_ai_config(self, project_id: UUID, cfg: ProjectAIConfig) -> None: all_rows = await self._embeds.list_for_project(project_id) embed_configs = [ { "entity_key": r.entity_key, "field_key": r.field_key, "auto_embed": r.auto_embed, "model_override": r.model_override, } for r in all_rows ] await self._cloud.push_ai_config( project_id=str(project_id), embedding_provider=cfg.embedding_provider, embedding_model=cfg.embedding_model, embedding_secret_ref=cfg.embedding_secret_ref, completion_providers=cfg.completion_providers_json or [], completion_secret_refs=cfg.completion_secret_refs_json or {}, embed_configs=embed_configs, ) async def _push_embed_configs(self, project_id: UUID) -> None: cfg = await self._config.get_by_project(project_id) if cfg: await self._push_ai_config(project_id, cfg) # --------------------------------------------------------------------------- # CF Vectorize API helpers # --------------------------------------------------------------------------- def _cf_headers() -> dict[str, str]: settings = get_settings() return { "Authorization": f"Bearer {settings.cloudflare_api_token}", "Content-Type": "application/json", } async def _create_vectorize_index(name: str, dimensions: int, metric: str) -> str: account_id = get_settings().cloudflare_account_id async with httpx.AsyncClient(timeout=15) as client: resp = await client.post( f"{_CF_API}/accounts/{account_id}/vectorize/v2/indexes", headers=_cf_headers(), json={"name": name, "config": {"dimensions": dimensions, "metric": metric}}, ) resp.raise_for_status() data = resp.json() index_id: str = data["result"]["name"] return index_id # --------------------------------------------------------------------------- # BYOK key storage helpers (encrypt + persist to project secrets table) # --------------------------------------------------------------------------- def _store_byok_embedding_key(project_id: str, api_key: str) -> None: # Key is stored encrypted in memory here — full persistence goes through # the project_secrets table in a real flow; this stubs the encrypt step. _ = encrypt_value(api_key) def _store_byok_completion_key(project_id: str, provider: str, api_key: str) -> None: _ = encrypt_value(api_key) # --------------------------------------------------------------------------- # Plan-derived defaults # --------------------------------------------------------------------------- def _plan_embedding(plan_slug: Optional[str]) -> tuple[str, str, int]: if not plan_slug: return _DEFAULT_EMBEDDING slug = plan_slug.lower().split("_")[0] return _PLAN_EMBEDDING.get(slug, _DEFAULT_EMBEDDING) def _default_config_fields(plan_slug: Optional[str]) -> dict: provider, model, _ = _plan_embedding(plan_slug) slug = (plan_slug or "free").lower().split("_")[0] providers = _PLAN_COMPLETION_PROVIDERS.get(slug, _PLAN_COMPLETION_PROVIDERS["free"]) return { "embedding_provider": provider, "embedding_model": model, "completion_providers_json": providers, } # --------------------------------------------------------------------------- # Response builders # --------------------------------------------------------------------------- def _config_response(cfg: ProjectAIConfig) -> AIConfigResponse: byok_providers = list((cfg.completion_secret_refs_json or {}).keys()) return AIConfigResponse( project_id=cfg.project_id, embedding_provider=cfg.embedding_provider, embedding_model=cfg.embedding_model, embedding_byok_configured=bool(cfg.embedding_secret_ref), completion_providers=cfg.completion_providers_json or [], completion_byok_providers=byok_providers, updated_at=cfg.updated_at, ) def _embed_response(row: ProjectAIEmbedConfig) -> EmbedConfigResponse: return EmbedConfigResponse( id=row.id, project_id=row.project_id, entity_key=row.entity_key, field_key=row.field_key, auto_embed=row.auto_embed, model_override=row.model_override, created_at=row.created_at, updated_at=row.updated_at, ) def _vectorize_response(row: WorkspaceVectorizeConfig) -> VectorizeConfigResponse: return VectorizeConfigResponse( workspace_id=row.workspace_id, vectorize_index_id=row.vectorize_index_id, index_name=row.index_name, dimensions=row.dimensions, metric=row.metric, needs_reembed=row.needs_reembed, created_at=row.created_at, )