portainer-core/services/core-ai/main.py

import os
import logging
import json
import time
from aiohttp import web
from aiohttp_cors import setup as cors_setup, ResourceOptions
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
logger = logging.getLogger(__name__)

# Import the agent logic
from src.agents import (
    get_simple_litellm_agent,
    get_pydantic_agent,
    PYDANTIC_AI_AVAILABLE
)
from src.tools import get_all_tools
from src.utils import extract_user_id_from_request

async def chat_completions(request):
    """
    Handles OpenAI-compatible chat completion requests using PydanticAI.
    Default endpoint - uses PydanticAI with tools enabled.
    """
    if not PYDANTIC_AI_AVAILABLE:
        return web.json_response({
            "error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
        }, status=503)

    try:
        data = await request.json()
        logger.info(f"[DEFAULT/PYDANTIC_AI] Received chat request")

        # Extract relevant fields from the request
        messages = data.get("messages")
        model = data.get("model", "pydantic")
        stream = data.get("stream", False)
        conversation_id = data.get("conversation_id")
        enable_tools = data.get("enable_tools", True)

        # Extract user ID from request
        user_id = extract_user_id_from_request(data)

        if not messages:
            raise web.HTTPBadRequest(reason="'messages' field is required")

        # Get the agent instance
        agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)

        # For non-streaming requests, collect the full response
        if not stream:
            response_content = await agent.chat_completion(
                messages=messages,
                conversation_id=conversation_id
            )
            return web.json_response({
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response_content
                    },
                    "finish_reason": "stop"
                }],
                "model": model,
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                },
                "tools_enabled": enable_tools,
                "tools_count": len(agent.tools_dict) if enable_tools else 0
            })
        else:
            # Handle streaming response
            response = web.StreamResponse(
                status=200,
                reason='OK',
                headers={
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                }
            )
            await response.prepare(request)

            try:
                async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
                    if chunk["type"] == "content":
                        chunk_data = {
                            "choices": [{
                                "index": 0,
                                "delta": {"content": chunk["content"]},
                                "finish_reason": chunk.get("finish_reason")
                            }],
                            "model": model
                        }
                        await response.write(f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8'))

                await response.write(b"data: [DONE]\n\n")
            finally:
                await response.write_eof()

            return response

    except web.HTTPBadRequest:
        raise
    except Exception as e:
        logger.exception(f"Error in chat_completions: {e}")
        return web.json_response({
            "error": {"message": f"Internal server error: {str(e)}"}
        }, status=500)

async def chat_simple(request):
    """
    Handles chat requests using SimpleLiteLLMAgent (fallback, no tools).
    Endpoint: /v1/chat/simple
    """
    try:
        data = await request.json()
        logger.info(f"[SIMPLE/LITELLM] Received chat request")

        # Extract relevant fields from the request
        messages = data.get("messages")
        model = data.get("model", "simple")
        stream = data.get("stream", False)
        conversation_id = data.get("conversation_id")

        if not messages:
            raise web.HTTPBadRequest(reason="'messages' field is required")

        # Get simple agent instance (no tools)
        agent = get_simple_litellm_agent()

        # For non-streaming requests
        if not stream:
            response_content = await agent.chat_completion(
                messages=messages,
                conversation_id=conversation_id
            )
            return web.json_response({
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response_content
                    },
                    "finish_reason": "stop"
                }],
                "model": model,
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
            })
        else:
            # Streaming response
            response = web.StreamResponse(
                status=200,
                reason='OK',
                headers={
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                }
            )
            await response.prepare(request)

            try:
                async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
                    if chunk["type"] == "content":
                        chunk_data = {
                            "choices": [{
                                "index": 0,
                                "delta": {"content": chunk["content"]},
                                "finish_reason": chunk.get("finish_reason")
                            }],
                            "model": model
                        }
                        await response.write(f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8'))

                await response.write(b"data: [DONE]\n\n")
            finally:
                await response.write_eof()

            return response

    except web.HTTPBadRequest:
        raise
    except Exception as e:
        logger.exception(f"Error in chat_simple: {e}")
        return web.json_response({
            "error": {"message": f"Internal server error: {str(e)}"}
        }, status=500)

async def list_models(request):
    """List available models"""
    return web.json_response({
        "object": "list",
        "data": [
            {
                "id": "pydantic",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "core-ai"
            },
            {
                "id": "simple",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "core-ai"
            }
        ]
    })

async def list_tools(request):
    """List all available tools"""
    try:
        tools = get_all_tools()
        tool_list = []

        for name, func in tools.items():
            import inspect
            doc = inspect.getdoc(func) or "No description"
            tool_list.append({
                "name": name,
                "description": doc.split('\n')[0],
                "type": "local" if not name.startswith("core-api__") else "openapi"
            })

        return web.json_response({
            "tools": tool_list,
            "tools_count": len(tool_list)
        })

    except Exception as e:
        logger.exception(f"Error listing tools: {e}")
        return web.json_response({
            "error": {"message": str(e)}
        }, status=500)

async def health_check(request):
    """Health check endpoint"""
    return web.json_response({
        "status": "ok",
        "service": "core-ai",
        "agents": {
            "simple": True,
            "pydantic": PYDANTIC_AI_AVAILABLE
        },
        "default_agent": "pydantic" if PYDANTIC_AI_AVAILABLE else "simple",
        "tools_count": len(get_all_tools())
    })

async def setup_routes(app):
    # Chat endpoints
    app.router.add_post("/chat/completions", chat_completions)     # Alias without /v1
    app.router.add_post("/v1/chat/completions", chat_completions)  # Default: PydanticAI
    app.router.add_post("/v1/chat/simple", chat_simple)            # Fallback: Simple agent

    # Models endpoint
    app.router.add_get("/models", list_models)
    app.router.add_get("/v1/models", list_models)

    # Tools endpoint
    app.router.add_get("/tools", list_tools)
    app.router.add_get("/v1/tools", list_tools)

    # Health check
    app.router.add_get("/health", health_check)

    # Setup CORS
    cors = cors_setup(app, defaults={
        "*": ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*",
            allow_methods="*"
        )
    })

    # Apply CORS to all routes
    for route in list(app.router.routes()):
        cors.add(route)

def main():
    app = web.Application()

    # Set up routes
    import asyncio
    loop = asyncio.get_event_loop()
    loop.run_until_complete(setup_routes(app))

    # Run the application
    logger.info("Starting core-ai service on http://0.0.0.0:8086")
    web.run_app(app, host='0.0.0.0', port=8086)

if __name__ == '__main__':
    main()