kairo-code 0.1.0__py3-none-any.whl

kairo/backend/services/conversation_service.py

@@ -0,0 +1,264 @@
+import logging
+
+from sqlalchemy import select, func, delete
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
+
+from backend.models.conversation import Conversation, Message
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationService:
+    def __init__(self, db: AsyncSession, user_id: str | None = None):
+        self.db = db
+        self.user_id = user_id
+
+    async def create(self, model: str = "nyx", title: str = "New Conversation", project_id: str | None = None) -> Conversation:
+        conv = Conversation(model=model, title=title, user_id=self.user_id, project_id=project_id)
+        self.db.add(conv)
+        await self.db.commit()
+        await self.db.refresh(conv)
+        return conv
+
+    async def list_all(self, limit: int = 50, offset: int = 0, project_id: str | None = None) -> tuple[list[dict], int]:
+        base = select(Conversation).where(Conversation.user_id == self.user_id) if self.user_id else select(Conversation)
+
+        # Count
+        count_stmt = select(func.count()).select_from(base.subquery())
+        total = (await self.db.execute(count_stmt)).scalar() or 0
+
+        # Data
+        stmt = (
+            select(
+                Conversation.id,
+                Conversation.title,
+                Conversation.model,
+                Conversation.updated_at,
+                Conversation.project_id,
+                func.count(Message.id).label("message_count"),
+            )
+            .outerjoin(Message)
+            .group_by(Conversation.id)
+            .order_by(Conversation.updated_at.desc())
+            .limit(limit)
+            .offset(offset)
+        )
+        if self.user_id:
+            stmt = stmt.where(Conversation.user_id == self.user_id)
+        if project_id is not None:
+            stmt = stmt.where(Conversation.project_id == project_id)
+
+        result = await self.db.execute(stmt)
+        rows = result.all()
+        items = [
+            {
+                "id": r.id,
+                "title": r.title,
+                "model": r.model,
+                "updated_at": r.updated_at,
+                "message_count": r.message_count,
+                "project_id": r.project_id,
+            }
+            for r in rows
+        ]
+        return items, total
+
+    async def search(self, query: str, limit: int = 50, offset: int = 0) -> tuple[list[dict], int]:
+        """Search conversations by title or message content."""
+        pattern = f"%{query}%"
+
+        # Find matching conversation IDs
+        matching_ids_stmt = (
+            select(Conversation.id)
+            .outerjoin(Message)
+            .where(
+                (Conversation.title.ilike(pattern)) | (Message.content.ilike(pattern))
+            )
+            .group_by(Conversation.id)
+        )
+        if self.user_id:
+            matching_ids_stmt = matching_ids_stmt.where(Conversation.user_id == self.user_id)
+
+        # Count
+        count_stmt = select(func.count()).select_from(matching_ids_stmt.subquery())
+        total = (await self.db.execute(count_stmt)).scalar() or 0
+
+        # Data
+        stmt = (
+            select(
+                Conversation.id,
+                Conversation.title,
+                Conversation.model,
+                Conversation.updated_at,
+                Conversation.project_id,
+                func.count(Message.id).label("message_count"),
+            )
+            .outerjoin(Message)
+            .where(Conversation.id.in_(matching_ids_stmt))
+            .group_by(Conversation.id)
+            .order_by(Conversation.updated_at.desc())
+            .limit(limit)
+            .offset(offset)
+        )
+        if self.user_id:
+            stmt = stmt.where(Conversation.user_id == self.user_id)
+
+        result = await self.db.execute(stmt)
+        rows = result.all()
+        items = [
+            {
+                "id": r.id,
+                "title": r.title,
+                "model": r.model,
+                "updated_at": r.updated_at,
+                "message_count": r.message_count,
+                "project_id": r.project_id,
+            }
+            for r in rows
+        ]
+        return items, total
+
+    async def get(self, conversation_id: str) -> Conversation | None:
+        stmt = (
+            select(Conversation)
+            .options(selectinload(Conversation.messages), selectinload(Conversation.project))
+            .where(Conversation.id == conversation_id)
+        )
+        if self.user_id:
+            stmt = stmt.where(Conversation.user_id == self.user_id)
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def get_fresh(self, conversation_id: str) -> Conversation | None:
+        """Get conversation with fresh data, bypassing the identity map cache."""
+        stmt = (
+            select(Conversation)
+            .options(selectinload(Conversation.messages), selectinload(Conversation.project))
+            .where(Conversation.id == conversation_id)
+            .execution_options(populate_existing=True)
+        )
+        if self.user_id:
+            stmt = stmt.where(Conversation.user_id == self.user_id)
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def delete(self, conversation_id: str) -> bool:
+        conv = await self.get(conversation_id)
+        if not conv:
+            return False
+        await self.db.delete(conv)
+        await self.db.commit()
+        return True
+
+    async def rename(self, conversation_id: str, title: str) -> Conversation | None:
+        conv = await self.get(conversation_id)
+        if not conv:
+            return None
+        conv.title = title
+        await self.db.commit()
+        await self.db.refresh(conv)
+        return conv
+
+    async def add_message(self, conversation_id: str, role: str, content: str, image_url: str | None = None) -> Message:
+        msg = Message(conversation_id=conversation_id, role=role, content=content, image_url=image_url)
+        self.db.add(msg)
+        await self.db.commit()
+        await self.db.refresh(msg)
+        return msg
+
+    async def get_message(self, message_id: str) -> Message | None:
+        stmt = select(Message).where(Message.id == message_id)
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def edit_message(self, message_id: str, new_content: str) -> Message | None:
+        """Edit a message and delete all subsequent messages in the conversation."""
+        msg = await self.get_message(message_id)
+        if not msg:
+            return None
+
+        # Verify user owns this conversation
+        conv = await self.get(msg.conversation_id)
+        if not conv:
+            return None
+
+        # Delete all messages after this one
+        await self.db.execute(
+            delete(Message).where(
+                Message.conversation_id == msg.conversation_id,
+                Message.created_at > msg.created_at,
+            )
+        )
+
+        # Update the message content
+        msg.content = new_content
+        await self.db.commit()
+        await self.db.refresh(msg)
+        return msg
+
+    async def search_messages(self, query: str, limit: int = 20, offset: int = 0) -> list[dict]:
+        """Search message content and return snippets with conversation context."""
+        pattern = f"%{query}%"
+        stmt = (
+            select(
+                Message.id,
+                Message.role,
+                Message.content,
+                Message.created_at,
+                Conversation.id.label("conversation_id"),
+                Conversation.title.label("conversation_title"),
+            )
+            .join(Conversation, Message.conversation_id == Conversation.id)
+            .where(Message.content.ilike(pattern))
+            .order_by(Message.created_at.desc())
+            .limit(limit)
+            .offset(offset)
+        )
+        if self.user_id:
+            stmt = stmt.where(Conversation.user_id == self.user_id)
+
+        result = await self.db.execute(stmt)
+        rows = result.all()
+        items = []
+        for r in rows:
+            # Build snippet: ~100 chars around first match
+            content = r.content
+            lower_content = content.lower()
+            idx = lower_content.find(query.lower())
+            if idx >= 0:
+                start = max(0, idx - 50)
+                end = min(len(content), idx + len(query) + 50)
+                snippet = ("..." if start > 0 else "") + content[start:end] + ("..." if end < len(content) else "")
+            else:
+                snippet = content[:100] + ("..." if len(content) > 100 else "")
+            items.append({
+                "message_id": r.id,
+                "role": r.role,
+                "snippet": snippet,
+                "created_at": r.created_at,
+                "conversation_id": r.conversation_id,
+                "conversation_title": r.conversation_title,
+            })
+        return items
+
+    async def delete_last_assistant_message(self, conversation_id: str) -> bool:
+        """Delete the last assistant message in a conversation for regeneration."""
+        conv = await self.get(conversation_id)
+        if not conv or not conv.messages:
+            return False
+
+        # Find last assistant message
+        messages = sorted(conv.messages, key=lambda m: m.created_at)
+        last_assistant = None
+        for m in reversed(messages):
+            if m.role == "assistant":
+                last_assistant = m
+                break
+
+        if not last_assistant:
+            return False
+
+        await self.db.delete(last_assistant)
+        await self.db.commit()
+        return True

kairo/backend/services/device_auth_service.py

@@ -0,0 +1,193 @@
+import hashlib
+import logging
+import random
+import string
+import uuid
+from datetime import UTC, datetime, timedelta
+
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.api_key import ApiKey
+from backend.models.device_code import DeviceCode
+from backend.models.user import User
+
+logger = logging.getLogger(__name__)
+
+
+def _generate_user_code() -> str:
+    """Generate a human-readable code like ABCD-1234."""
+    letters = "".join(random.choices(string.ascii_uppercase, k=4))
+    digits = "".join(random.choices(string.digits, k=4))
+    return f"{letters}-{digits}"
+
+
+def _generate_device_code() -> str:
+    """Generate a 64-character hex device code."""
+    return uuid.uuid4().hex + uuid.uuid4().hex
+
+
+class DeviceAuthService:
+    EXPIRY_MINUTES = 10
+    POLL_INTERVAL = 5
+
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def create_device_code(self, client_name: str) -> DeviceCode:
+        """Create a new device code for the CLI auth flow."""
+        user_code = await self._unique_user_code()
+
+        device = DeviceCode(
+            device_code=_generate_device_code(),
+            user_code=user_code,
+            client_name=client_name,
+            expires_at=datetime.now(UTC) + timedelta(minutes=self.EXPIRY_MINUTES),
+            interval=self.POLL_INTERVAL,
+        )
+        self.db.add(device)
+        await self.db.commit()
+        await self.db.refresh(device)
+        logger.info("Device code created: user_code=%s client=%s", user_code, client_name)
+        return device
+
+    async def approve(self, user_code: str, user: User) -> tuple[str, DeviceCode] | None:
+        """Approve a pending device code.
+
+        Returns (raw_api_key, device_code) or None if not found / expired.
+        """
+        device = await self._find_pending_device(user_code)
+        if not device:
+            return None
+
+        raw_key = self._create_raw_cli_key()
+        api_key = self._build_api_key(raw_key, user.id, device.client_name)
+        self.db.add(api_key)
+        await self.db.flush()
+
+        self._mark_approved(device, user.id, api_key.id, raw_key)
+        await self.db.commit()
+
+        logger.info("Device code approved: user_code=%s user=%s", user_code, user.id)
+        return raw_key, device
+
+    async def poll_token(
+        self, device_code: str
+    ) -> tuple[str, dict | None]:
+        """Poll for token status.
+
+        Returns (status, user_info_or_none).
+        status: authorization_pending | expired_token | access_denied | approved
+        user_info: {"raw_key": ..., "plan": ..., "email": ...} when approved.
+        """
+        device = await self._find_device_by_code(device_code)
+        if not device:
+            return "expired_token", None
+
+        if self._is_expired_pending(device):
+            device.status = "expired"
+            await self.db.commit()
+            return "expired_token", None
+
+        if device.status in ("denied", "expired"):
+            status = "access_denied" if device.status == "denied" else "expired_token"
+            return status, None
+
+        if device.status == "pending":
+            return "authorization_pending", None
+
+        if device.status == "approved" and device.cli_token:
+            return await self._consume_approved_token(device)
+
+        return "authorization_pending", None
+
+    async def cleanup_expired(self) -> int:
+        """Delete device codes older than 1 hour."""
+        cutoff = datetime.now(UTC) - timedelta(hours=1)
+        stmt = delete(DeviceCode).where(DeviceCode.created_at < cutoff)
+        result = await self.db.execute(stmt)
+        await self.db.commit()
+        deleted = result.rowcount or 0
+        if deleted:
+            logger.info("Cleaned up %d expired device codes", deleted)
+        return deleted
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    async def _unique_user_code(self) -> str:
+        """Generate a user code that does not collide with pending codes."""
+        for _ in range(5):
+            user_code = _generate_user_code()
+            stmt = select(DeviceCode).where(
+                DeviceCode.user_code == user_code,
+                DeviceCode.status == "pending",
+            )
+            result = await self.db.execute(stmt)
+            if not result.scalar_one_or_none():
+                return user_code
+        # Extremely unlikely fallback
+        return _generate_user_code()
+
+    async def _find_pending_device(self, user_code: str) -> DeviceCode | None:
+        stmt = select(DeviceCode).where(
+            DeviceCode.user_code == user_code,
+            DeviceCode.status == "pending",
+            DeviceCode.expires_at > datetime.now(UTC),
+        )
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def _find_device_by_code(self, device_code: str) -> DeviceCode | None:
+        stmt = select(DeviceCode).where(DeviceCode.device_code == device_code)
+        result = await self.db.execute(stmt)
+        return result.scalar_one_or_none()
+
+    @staticmethod
+    def _create_raw_cli_key() -> str:
+        return f"sk-kairo-cli-{uuid.uuid4().hex}"
+
+    @staticmethod
+    def _build_api_key(raw_key: str, user_id: str, client_name: str | None) -> ApiKey:
+        prefix = raw_key[:24]
+        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
+        return ApiKey(
+            user_id=user_id,
+            name=f"Kairo Code CLI ({client_name or 'cli'})",
+            key_prefix=prefix,
+            key_hash=key_hash,
+            key_type="cli",
+        )
+
+    @staticmethod
+    def _mark_approved(
+        device: DeviceCode, user_id: str, api_key_id: str, raw_key: str
+    ) -> None:
+        device.status = "approved"
+        device.user_id = user_id
+        device.cli_api_key_id = api_key_id
+        device.cli_token = raw_key
+        device.approved_at = datetime.now(UTC)
+
+    @staticmethod
+    def _is_expired_pending(device: DeviceCode) -> bool:
+        return device.expires_at < datetime.now(UTC) and device.status == "pending"
+
+    async def _consume_approved_token(
+        self, device: DeviceCode
+    ) -> tuple[str, dict | None]:
+        """Return the token and clear it (single-use retrieval)."""
+        raw_key = device.cli_token
+        device.cli_token = None  # single-use: clear after first poll
+        await self.db.commit()
+
+        user = await self.db.get(User, device.user_id)
+        if not user:
+            return "access_denied", None
+
+        return "approved", {
+            "raw_key": raw_key,
+            "plan": user.plan,
+            "email": user.email,
+        }

kairo/backend/services/email_service.py

@@ -0,0 +1,55 @@
+import logging
+import smtplib
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+
+from backend.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class EmailService:
+    def send_verification_email(self, to: str, token: str) -> None:
+        link = f"{settings.APP_BASE_URL}/verify-email?token={token}"
+        subject = "Verify your Kairo email"
+        html = (
+            f"<h2>Welcome to Kairo</h2>"
+            f"<p>Click the link below to verify your email address:</p>"
+            f'<p><a href="{link}">Verify Email</a></p>'
+            f"<p>Or copy this URL: {link}</p>"
+            f"<p>If you didn't create a Kairo account, you can ignore this email.</p>"
+        )
+        self._send(to, subject, html)
+
+    def send_password_reset_email(self, to: str, token: str) -> None:
+        link = f"{settings.APP_BASE_URL}/reset-password?token={token}"
+        subject = "Reset your Kairo password"
+        html = (
+            f"<h2>Password Reset</h2>"
+            f"<p>Click the link below to reset your password. This link expires in 1 hour.</p>"
+            f'<p><a href="{link}">Reset Password</a></p>'
+            f"<p>Or copy this URL: {link}</p>"
+            f"<p>If you didn't request this, you can ignore this email.</p>"
+        )
+        self._send(to, subject, html)
+
+    def _send(self, to: str, subject: str, html_body: str) -> None:
+        if not settings.SMTP_USERNAME:
+            logger.warning("SMTP not configured, skipping email to %s", to)
+            return
+
+        msg = MIMEMultipart("alternative")
+        msg["Subject"] = subject
+        msg["From"] = f"{settings.SMTP_FROM_NAME} <{settings.SMTP_FROM_EMAIL}>"
+        msg["To"] = to
+        msg.attach(MIMEText(html_body, "html"))
+
+        try:
+            with smtplib.SMTP(settings.SMTP_HOST, settings.SMTP_PORT) as server:
+                server.starttls()
+                server.login(settings.SMTP_USERNAME, settings.SMTP_PASSWORD)
+                server.sendmail(settings.SMTP_FROM_EMAIL, to, msg.as_string())
+            logger.info("Email sent to %s: %s", to, subject)
+        except Exception:
+            logger.exception("Failed to send email to %s", to)
+            raise

kairo/backend/services/image_service.py

@@ -0,0 +1,181 @@
+import logging
+import uuid
+from datetime import datetime, timezone
+
+import boto3
+import httpx
+from sqlalchemy import select, func
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.config import settings
+from backend.models.image_generation import ImageGeneration
+from backend.models.user import PlanType, User
+from backend.services.conversation_service import ConversationService
+
+logger = logging.getLogger(__name__)
+
+IMAGE_LIMITS: dict[str, dict[str, int]] = {
+    PlanType.PRO.value: {
+        "daily": settings.PRO_DAILY_IMAGE_LIMIT,
+        "monthly": settings.PRO_MONTHLY_IMAGE_LIMIT,
+    },
+    PlanType.MAX.value: {
+        "daily": settings.MAX_DAILY_IMAGE_LIMIT,
+        "monthly": settings.MAX_MONTHLY_IMAGE_LIMIT,
+    },
+}
+
+
+class ImageService:
+    def __init__(self, db: AsyncSession, user_id: str):
+        self.db = db
+        self.user_id = user_id
+        self.conv_service = ConversationService(db, user_id=user_id)
+
+    async def check_access(self) -> tuple[bool, str]:
+        """Check plan tier and image usage limits."""
+        user = await self.db.get(User, self.user_id)
+        if not user:
+            return False, "User not found."
+
+        if user.plan not in (PlanType.PRO.value, PlanType.MAX.value):
+            return False, "Image generation requires a Pro or Max plan."
+
+        limits = IMAGE_LIMITS.get(user.plan)
+        if not limits:
+            return False, "No image limits configured for your plan."
+
+        now = datetime.now(timezone.utc)
+        start_of_day = now.replace(hour=0, minute=0, second=0, microsecond=0)
+        start_of_month = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+
+        daily_count = (await self.db.execute(
+            select(func.count())
+            .where(ImageGeneration.user_id == self.user_id)
+            .where(ImageGeneration.created_at >= start_of_day)
+        )).scalar() or 0
+
+        if daily_count >= limits["daily"]:
+            return False, "Daily image generation limit reached."
+
+        monthly_count = (await self.db.execute(
+            select(func.count())
+            .where(ImageGeneration.user_id == self.user_id)
+            .where(ImageGeneration.created_at >= start_of_month)
+        )).scalar() or 0
+
+        if monthly_count >= limits["monthly"]:
+            return False, "Monthly image generation limit reached."
+
+        return True, ""
+
+    async def _save_result(
+        self,
+        image_bytes: bytes,
+        prompt: str,
+        conversation_id: str | None,
+        width: int = 1024,
+        height: int = 1024,
+    ) -> dict:
+        """Upload image to S3, save to conversation, record usage."""
+        # Upload to S3
+        image_id = str(uuid.uuid4())
+        s3_key = f"generated/{image_id}.png"
+        s3 = boto3.client("s3", region_name=settings.S3_IMAGES_REGION)
+        s3.put_object(
+            Bucket=settings.S3_IMAGES_BUCKET,
+            Key=s3_key,
+            Body=image_bytes,
+            ContentType="image/png",
+        )
+        image_url = f"https://{settings.S3_IMAGES_BUCKET}.s3.{settings.S3_IMAGES_REGION}.amazonaws.com/{s3_key}"
+
+        # Get or create conversation
+        conv = None
+        if conversation_id:
+            conv = await self.conv_service.get(conversation_id)
+        if not conv:
+            conv = await self.conv_service.create(model="flux-dev")
+
+        # Save user message (the prompt)
+        await self.conv_service.add_message(conv.id, "user", prompt)
+
+        # Save assistant message with image_url
+        msg = await self.conv_service.add_message(
+            conv.id, "assistant", f"![Generated image]({image_url})",
+            image_url=image_url,
+        )
+        msg_id = msg.id
+
+        # Record in image_generations table
+        record = ImageGeneration(
+            user_id=self.user_id,
+            conversation_id=conv.id,
+            prompt=prompt,
+            image_url=image_url,
+            width=width,
+            height=height,
+        )
+        self.db.add(record)
+        await self.db.commit()
+
+        # Auto-title if new conversation
+        if conv.title == "New Conversation":
+            title = prompt[:47] + "..." if len(prompt) > 50 else prompt
+            await self.conv_service.rename(conv.id, title)
+
+        return {
+            "image_url": image_url,
+            "message_id": msg_id,
+            "conversation_id": conv.id,
+            "prompt": prompt,
+        }
+
+    async def generate(
+        self,
+        prompt: str,
+        conversation_id: str | None,
+        width: int = 1024,
+        height: int = 1024,
+    ) -> dict:
+        """Call FLUX text2img service, upload to S3, save to conversation."""
+        base_url = settings.FLUX_BASE_URL.rstrip("/")
+        headers = {}
+        if settings.FLUX_API_KEY:
+            headers["Authorization"] = f"Bearer {settings.FLUX_API_KEY}"
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=10.0)) as client:
+            resp = await client.post(
+                f"{base_url}/generate",
+                json={"prompt": prompt, "width": width, "height": height},
+                headers=headers,
+            )
+            resp.raise_for_status()
+            image_bytes = resp.content
+
+        return await self._save_result(image_bytes, prompt, conversation_id, width, height)
+
+    async def generate_img2img(
+        self,
+        image_bytes: bytes,
+        prompt: str,
+        conversation_id: str | None,
+        strength: float = 0.75,
+    ) -> dict:
+        """Call FLUX img2img service, upload to S3, save to conversation."""
+        base_url = settings.FLUX_BASE_URL.rstrip("/")
+        headers = {}
+        if settings.FLUX_API_KEY:
+            headers["Authorization"] = f"Bearer {settings.FLUX_API_KEY}"
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=10.0)) as client:
+            resp = await client.post(
+                f"{base_url}/img2img",
+                files={"image": ("input.png", image_bytes, "image/png")},
+                data={"prompt": prompt, "strength": str(strength)},
+                headers=headers,
+            )
+            resp.raise_for_status()
+            result_bytes = resp.content
+
+        return await self._save_result(result_bytes, prompt, conversation_id)