kairo-code 0.1.0__py3-none-any.whl

kairo/backend/api/openai_compat.py

@@ -0,0 +1,225 @@
+"""OpenAI-compatible API endpoint — mounted at /v1/ for SDK compatibility."""
+
+import json
+import logging
+import time
+import uuid
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.config import settings
+from backend.core.api_key_auth import get_api_key_user
+from backend.core.database import get_db
+from backend.core.rate_limit import rate_limit_api
+from backend.models.api_key import ApiKey
+from backend.models.user import User
+from backend.schemas.openai_compat import (
+    ChatCompletionChunk,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionChoice,
+    ChatMessage,
+    ChoiceDelta,
+    StreamChoice,
+    UsageInfo,
+)
+from backend.services.api_usage_service import ApiUsageService
+from backend.services.llm_service import LLMService
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/v1", tags=["OpenAI Compatible"])
+
+
+def _get_llm_service(request: Request) -> LLMService:
+    return request.app.state.llm_service
+
+
+@router.post("/chat/completions")
+async def chat_completions(
+    req: ChatCompletionRequest,
+    request: Request,
+    auth: tuple[User, ApiKey] = Depends(get_api_key_user),
+    db: AsyncSession = Depends(get_db),
+    llm_service: LLMService = Depends(_get_llm_service),
+    _rate: None = Depends(rate_limit_api),
+):
+    if not settings.FEATURE_KAIRO_API_ENABLED:
+        raise HTTPException(status_code=503, detail="Kairo API is coming soon.")
+
+    user, api_key = auth
+
+    # Check API usage limits
+    usage_svc = ApiUsageService(db)
+    allowed, reason = await usage_svc.check_limits(user.id)
+    if not allowed:
+        raise HTTPException(status_code=429, detail=reason)
+
+    # Resolve model
+    model_id = llm_service.resolve_model(req.model)
+    client, base_url, resolved_model, _is_fallback = llm_service._select_backend(req.model)
+
+    # Build vLLM payload — serialize full request, override model with resolved name
+    payload = req.model_dump(exclude_none=True)
+    payload["model"] = resolved_model
+
+    request_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
+    created = int(time.time())
+
+    if req.stream:
+        payload["stream_options"] = {"include_usage": True}
+        return StreamingResponse(
+            _stream_proxy(
+                client, base_url, payload, request_id, created, resolved_model,
+                api_key.id, user.id, db,
+            ),
+            media_type="text/event-stream",
+        )
+    else:
+        # Non-streaming: proxy full request
+        try:
+            resp = await client.post(f"{base_url}/v1/chat/completions", json=payload)
+            resp.raise_for_status()
+            data = resp.json()
+        except Exception as e:
+            logger.error("vLLM proxy error: %s", e)
+            raise HTTPException(status_code=502, detail="Inference backend unavailable")
+
+        # Record usage
+        usage = data.get("usage", {})
+        prompt_tokens = usage.get("prompt_tokens", 0)
+        completion_tokens = usage.get("completion_tokens", 0)
+        await usage_svc.record(
+            api_key_id=api_key.id, user_id=user.id, model=resolved_model,
+            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens,
+            endpoint="/v1/chat/completions",
+        )
+
+        # Reformat as our response — forward all message fields including tool_calls
+        choices = []
+        for c in data.get("choices", []):
+            msg = c.get("message", {})
+            choices.append(ChatCompletionChoice(
+                index=c.get("index", 0),
+                message=ChatMessage(
+                    role=msg.get("role", "assistant"),
+                    content=msg.get("content"),
+                    name=msg.get("name"),
+                    tool_calls=msg.get("tool_calls"),
+                    tool_call_id=msg.get("tool_call_id"),
+                ),
+                finish_reason=c.get("finish_reason"),
+            ))
+
+        return ChatCompletionResponse(
+            id=request_id, created=created, model=resolved_model,
+            choices=choices,
+            usage=UsageInfo(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+            ) if usage else None,
+        )
+
+
+async def _stream_proxy(
+    client, base_url, payload, request_id, created, model,
+    api_key_id, user_id, db,
+):
+    """Stream SSE from vLLM, relay in OpenAI format, record usage at end."""
+    prompt_tokens = 0
+    completion_tokens = 0
+
+    try:
+        async with client.stream("POST", f"{base_url}/v1/chat/completions", json=payload) as response:
+            response.raise_for_status()
+            async for line in response.aiter_lines():
+                if not line.startswith("data: "):
+                    continue
+                data_str = line[6:]
+                if data_str.strip() == "[DONE]":
+                    break
+                try:
+                    chunk = json.loads(data_str)
+                    # Capture usage
+                    if "usage" in chunk and chunk["usage"]:
+                        prompt_tokens = chunk["usage"].get("prompt_tokens", 0)
+                        completion_tokens = chunk["usage"].get("completion_tokens", 0)
+
+                    choices = chunk.get("choices", [])
+                    if not choices:
+                        # Emit usage-only final chunk if present
+                        if prompt_tokens or completion_tokens:
+                            usage_chunk = ChatCompletionChunk(
+                                id=request_id, created=created, model=model,
+                                choices=[],
+                                usage=UsageInfo(
+                                    prompt_tokens=prompt_tokens,
+                                    completion_tokens=completion_tokens,
+                                    total_tokens=prompt_tokens + completion_tokens,
+                                ),
+                            )
+                            yield f"data: {usage_chunk.model_dump_json()}\n\n"
+                        continue
+
+                    # Forward all choices, including tool_calls deltas
+                    stream_choices = []
+                    for c in choices:
+                        delta = c.get("delta", {})
+                        stream_choices.append(StreamChoice(
+                            index=c.get("index", 0),
+                            delta=ChoiceDelta(
+                                role=delta.get("role"),
+                                content=delta.get("content"),
+                                tool_calls=delta.get("tool_calls"),
+                            ),
+                            finish_reason=c.get("finish_reason"),
+                        ))
+                    out = ChatCompletionChunk(
+                        id=request_id, created=created, model=model,
+                        choices=stream_choices,
+                    )
+                    yield f"data: {out.model_dump_json()}\n\n"
+                except (json.JSONDecodeError, KeyError, IndexError):
+                    continue
+
+        yield "data: [DONE]\n\n"
+    except Exception as e:
+        logger.error("Stream proxy error: %s", e)
+        yield f"data: {json.dumps({'error': 'Inference backend error'})}\n\n"
+        yield "data: [DONE]\n\n"
+
+    # Record usage after stream completes
+    if prompt_tokens or completion_tokens:
+        try:
+            usage_svc = ApiUsageService(db)
+            await usage_svc.record(
+                api_key_id=api_key_id, user_id=user_id, model=model,
+                prompt_tokens=prompt_tokens, completion_tokens=completion_tokens,
+                endpoint="/v1/chat/completions",
+            )
+        except Exception as e:
+            logger.error("Failed to record API usage: %s", e)
+
+
+@router.get("/models")
+async def list_models(
+    llm_service: LLMService = Depends(_get_llm_service),
+):
+    """List available models — public endpoint, no auth required."""
+    models = await llm_service.list_available_models()
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": m["id"],
+                "object": "model",
+                "owned_by": "kairon-labs",
+                "permission": [],
+            }
+            for m in models
+            if m.get("available", True)
+        ],
+    }

kairo/backend/api/projects.py

@@ -0,0 +1,102 @@
+from fastapi import APIRouter, Depends, HTTPException
+
+from backend.core.dependencies import get_project_service
+from backend.schemas.project import ProjectCreate, ProjectUpdate, ProjectResponse
+from backend.services.project_service import ProjectService
+
+router = APIRouter()
+
+
+@router.post("/", response_model=ProjectResponse)
+async def create_project(
+    body: ProjectCreate,
+    service: ProjectService = Depends(get_project_service),
+):
+    project = await service.create(body.name)
+    return {
+        "id": project.id,
+        "name": project.name,
+        "instructions": project.instructions,
+        "conversation_count": 0,
+        "created_at": project.created_at,
+        "updated_at": project.updated_at,
+    }
+
+
+@router.get("/", response_model=list[ProjectResponse])
+async def list_projects(
+    service: ProjectService = Depends(get_project_service),
+):
+    return await service.list_all()
+
+
+@router.get("/{project_id}", response_model=ProjectResponse)
+async def get_project(
+    project_id: str,
+    service: ProjectService = Depends(get_project_service),
+):
+    project = await service.get(project_id)
+    if not project:
+        raise HTTPException(status_code=404, detail="Project not found")
+    return {
+        "id": project.id,
+        "name": project.name,
+        "instructions": project.instructions,
+        "conversation_count": len(project.conversations),
+        "created_at": project.created_at,
+        "updated_at": project.updated_at,
+    }
+
+
+@router.patch("/{project_id}", response_model=ProjectResponse)
+async def update_project(
+    project_id: str,
+    body: ProjectUpdate,
+    service: ProjectService = Depends(get_project_service),
+):
+    updates = body.model_dump(exclude_unset=True)
+    if not updates:
+        raise HTTPException(status_code=400, detail="No fields to update")
+    project = await service.update(project_id, **updates)
+    if not project:
+        raise HTTPException(status_code=404, detail="Project not found")
+    return {
+        "id": project.id,
+        "name": project.name,
+        "instructions": project.instructions,
+        "conversation_count": len(project.conversations),
+        "created_at": project.created_at,
+        "updated_at": project.updated_at,
+    }
+
+
+@router.delete("/{project_id}", status_code=204)
+async def delete_project(
+    project_id: str,
+    service: ProjectService = Depends(get_project_service),
+):
+    deleted = await service.delete(project_id)
+    if not deleted:
+        raise HTTPException(status_code=404, detail="Project not found")
+
+
+@router.post("/{project_id}/conversations/{conversation_id}", status_code=204)
+async def add_conversation_to_project(
+    project_id: str,
+    conversation_id: str,
+    service: ProjectService = Depends(get_project_service),
+):
+    added = await service.add_conversation(project_id, conversation_id)
+    if not added:
+        raise HTTPException(status_code=404, detail="Project or conversation not found")
+
+
+@router.delete("/{project_id}/conversations/{conversation_id}", status_code=204)
+async def remove_conversation_from_project(
+    project_id: str,
+    conversation_id: str,
+    service: ProjectService = Depends(get_project_service),
+):
+    removed = await service.remove_conversation(conversation_id)
+    if not removed:
+        raise HTTPException(status_code=404, detail="Conversation not found")

kairo/backend/api/usage.py

@@ -0,0 +1,32 @@
+from fastapi import APIRouter, Depends, Query
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.core.database import get_db
+from backend.core.dependencies import get_current_user
+from backend.models.user import User
+from backend.schemas.usage import UsageSummaryResponse, UsageHistoryResponse, UsageHistoryEntry
+from backend.services.usage_service import UsageService
+
+router = APIRouter(prefix="/usage", tags=["usage"])
+
+
+@router.get("/summary", response_model=UsageSummaryResponse)
+async def usage_summary(
+    user: User = Depends(get_current_user),
+    db: AsyncSession = Depends(get_db),
+):
+    service = UsageService(db)
+    return await service.get_usage_summary(user.id)
+
+
+@router.get("/history", response_model=UsageHistoryResponse)
+async def usage_history(
+    days: int = Query(default=30, ge=1, le=365),
+    user: User = Depends(get_current_user),
+    db: AsyncSession = Depends(get_db),
+):
+    service = UsageService(db)
+    entries = await service.get_usage_history(user.id, days=days)
+    return UsageHistoryResponse(
+        entries=[UsageHistoryEntry(**e) for e in entries]
+    )

kairo/backend/api/webhooks.py

@@ -0,0 +1,79 @@
+import logging
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.core.database import get_db
+from backend.models.user import PlanType, User
+from backend.services.stripe_service import StripeService
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/webhooks", tags=["webhooks"])
+
+
+@router.post("/stripe")
+async def stripe_webhook(request: Request, db: AsyncSession = Depends(get_db)):
+    payload = await request.body()
+    sig = request.headers.get("stripe-signature", "")
+
+    stripe_service = StripeService()
+    try:
+        result = stripe_service.handle_webhook_event(payload, sig)
+    except Exception as e:
+        logger.warning("Webhook signature verification failed: %s", e)
+        raise HTTPException(status_code=400, detail="Invalid signature")
+
+    action = result.get("action")
+
+    if action == "upgrade":
+        customer_id = result["customer_id"]
+        subscription_id = result["subscription_id"]
+        user_id = result.get("user_id")
+
+        user = None
+        if user_id:
+            user = await db.get(User, user_id)
+        if not user:
+            stmt = select(User).where(User.stripe_customer_id == customer_id)
+            res = await db.execute(stmt)
+            user = res.scalar_one_or_none()
+        if user:
+            user.plan = PlanType.PRO.value
+            user.stripe_customer_id = customer_id
+            user.stripe_subscription_id = subscription_id
+            await db.commit()
+            logger.info("User %s upgraded to Pro", user.id)
+        else:
+            logger.error("Webhook upgrade: no user found for customer=%s user_id=%s", customer_id, user_id)
+
+    elif action == "downgrade":
+        customer_id = result["customer_id"]
+        stmt = select(User).where(User.stripe_customer_id == customer_id)
+        res = await db.execute(stmt)
+        user = res.scalar_one_or_none()
+        if user:
+            user.plan = PlanType.FREE.value
+            user.stripe_subscription_id = None
+            await db.commit()
+            logger.info("User %s downgraded to Free", user.id)
+        else:
+            logger.error("Webhook downgrade: no user found for customer=%s", customer_id)
+
+    elif action == "sync_active":
+        customer_id = result["customer_id"]
+        subscription_id = result["subscription_id"]
+        stmt = select(User).where(User.stripe_customer_id == customer_id)
+        res = await db.execute(stmt)
+        user = res.scalar_one_or_none()
+        if user:
+            user.plan = PlanType.PRO.value
+            user.stripe_subscription_id = subscription_id
+            await db.commit()
+
+    elif action == "payment_failed":
+        customer_id = result["customer_id"]
+        logger.warning("Payment failed for customer %s", customer_id)
+
+    return {"received": True}