api.py

"""
FastAPI应用 - OpenAI兼容的API接口
"""
import asyncio
import json
import logging
import time
from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException, Depends, Security, status
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

from config import settings
from models import (
    ChatCompletionRequest, ChatCompletionResponse, ChatCompletionChoice,
    ChatCompletionStreamResponse, ChatCompletionStreamChoice,
    ChatCompletionUsage, ChatMessage, ErrorResponse, ModelsResponse, ModelInfo
)
from simple_rovo_client import rovo_client
from session_manager import SessionManager
from auth import verify_api_key
from exceptions import (
    setup_exception_handlers, create_error_response, RovoAPIException,
    SessionNotFoundError, RovoConnectionError, ValidationError
)
from middleware import setup_middleware, performance_middleware, health_check


# 配置日志
logging.basicConfig(
    level=getattr(logging, settings.log_level.upper()),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    filename=settings.log_file
)
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(
    title="Rovo Dev OpenAI API",
    description="OpenAI兼容的Rovo Dev API接口",
    version="1.0.0",
    docs_url="/docs" if settings.debug else None,
    redoc_url="/redoc" if settings.debug else None
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 设置异常处理器
setup_exception_handlers(app)

# 设置中间件
setup_middleware(app)

# 创建会话管理器
session_manager = SessionManager(rovo_client)


@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    await rovo_client.start()
    await session_manager.start()
    logger.info("Rovo Dev OpenAI API服务已启动")


@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭事件"""
    await session_manager.stop()
    await rovo_client.stop()
    logger.info("Rovo Dev OpenAI API服务已关闭")


@app.get("/")
async def root():
    """根路径"""
    return {"message": "Rovo Dev OpenAI API", "version": "1.0.0"}


@app.get("/health")
async def health_check_endpoint():
    """健康检查"""
    return health_check.get_health_status()


@app.get("/stats")
async def get_stats():
    """获取性能统计"""
    if performance_middleware:
        return performance_middleware.get_stats()
    return {"message": "Performance middleware not available"}


@app.get("/sessions")
async def list_sessions(auth: bool = Depends(verify_api_key)):
    """列出所有会话"""
    return {
        "sessions": session_manager.list_sessions(),
        "total": len(session_manager.sessions)
    }


@app.get("/sessions/{session_id}")
async def get_session(session_id: str, auth: bool = Depends(verify_api_key)):
    """获取指定会话信息"""
    session_info = session_manager.get_session_info(session_id)
    if not session_info:
        raise HTTPException(status_code=404, detail="Session not found")
    return session_info


@app.delete("/sessions/{session_id}")
async def close_session_endpoint(session_id: str, auth: bool = Depends(verify_api_key)):
    """关闭指定会话"""
    session_info = session_manager.get_session_info(session_id)
    if not session_info:
        raise HTTPException(status_code=404, detail="Session not found")

    await session_manager.close_session(session_id)
    return {"message": f"Session {session_id} closed successfully"}


@app.get("/v1/models", response_model=ModelsResponse)
async def list_models(auth: bool = Depends(verify_api_key)):
    """列出可用模型"""
    return ModelsResponse(
        data=[
            ModelInfo(
                id="rovo-dev",
                owned_by="atlassian"
            )
        ]
    )


@app.post("/v1/chat/completions")
async def create_chat_completion(
    request: ChatCompletionRequest,
    auth: bool = Depends(verify_api_key)
):
    """创建聊天完成"""
    try:
        # 验证请求
        if not request.messages:
            raise ValidationError("Messages cannot be empty")

        # 提取用户信息
        user_id = request.user
        conversation_id = None

        # 尝试从消息中提取conversation_id（如果有的话）
        if hasattr(request, 'conversation_id'):
            conversation_id = request.conversation_id

        # 获取或创建会话
        try:
            # 将Pydantic模型转换为字典
            messages_dict = [msg.dict() for msg in request.messages]
            session_id = await session_manager.get_or_create_session(
                messages=messages_dict,
                user_id=user_id,
                conversation_id=conversation_id
            )
        except Exception as e:
            raise RovoConnectionError(f"Failed to get or create session: {str(e)}")

        try:
            
            if request.stream:
                # 流式响应
                return StreamingResponse(
                    _stream_chat_completion(session_id, request),
                    media_type="text/plain",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "Content-Type": "text/event-stream"
                    }
                )
            else:
                # 非流式响应
                response_content = ""
                async for chunk in session_manager.send_message(
                    session_id,
                    messages_dict,
                    user_id
                ):
                    response_content += chunk
                
                # 创建响应
                message_content = _build_message_content(request.messages)
                response = ChatCompletionResponse(
                    model=request.model,
                    choices=[
                        ChatCompletionChoice(
                            index=0,
                            message=ChatMessage(
                                role="assistant",
                                content=response_content.strip()
                            ),
                            finish_reason="stop"
                        )
                    ],
                    usage=ChatCompletionUsage(
                        prompt_tokens=_estimate_tokens(message_content),
                        completion_tokens=_estimate_tokens(response_content),
                        total_tokens=_estimate_tokens(message_content + response_content)
                    )
                )

                return response

        except Exception as e:
            # 如果出错，记录但不关闭会话（会话管理器会处理）
            logger.error(f"处理消息时出错: {e}")
            raise
            
    except RovoAPIException:
        raise
    except Exception as e:
        logger.error(f"处理聊天完成请求时出错: {e}")
        raise RovoAPIException(f"Internal server error: {str(e)}")


def _build_message_content(messages: List[ChatMessage]) -> str:
    """构建发送给Rovo Dev的消息内容"""
    # 简单地使用最后一条用户消息
    user_messages = [msg for msg in messages if msg.role == "user"]
    if user_messages:
        return user_messages[-1].content
    
    # 如果没有用户消息，使用最后一条消息
    return messages[-1].content if messages else ""


async def _stream_chat_completion(session_id: str, request: ChatCompletionRequest):
    """生成流式聊天完成响应"""
    try:
        completion_id = f"chatcmpl-{int(time.time())}"

        # 发送初始响应
        initial_response = ChatCompletionStreamResponse(
            id=completion_id,
            model=request.model,
            choices=[
                ChatCompletionStreamChoice(
                    index=0,
                    delta={"role": "assistant", "content": ""},
                    finish_reason=None
                )
            ]
        )
        yield f"data: {initial_response.json()}\n\n"

        # 流式发送内容
        content_buffer = ""
        user_id = request.user
        messages_dict = [msg.dict() for msg in request.messages]
        async for chunk in session_manager.send_message(
            session_id,
            messages_dict,
            user_id
        ):
            if chunk.strip():
                content_buffer += chunk
                stream_response = ChatCompletionStreamResponse(
                    id=completion_id,
                    model=request.model,
                    choices=[
                        ChatCompletionStreamChoice(
                            index=0,
                            delta={"content": chunk},
                            finish_reason=None
                        )
                    ]
                )
                yield f"data: {stream_response.json()}\n\n"

                # 添加小延迟以避免过快的流式输出
                await asyncio.sleep(0.01)

        # 发送结束标记
        final_response = ChatCompletionStreamResponse(
            id=completion_id,
            model=request.model,
            choices=[
                ChatCompletionStreamChoice(
                    index=0,
                    delta={},
                    finish_reason="stop"
                )
            ]
        )
        yield f"data: {final_response.json()}\n\n"
        yield "data: [DONE]\n\n"

    except Exception as e:
        logger.error(f"流式响应时出错: {e}")
        error_response = {
            "error": {
                "message": str(e),
                "type": "server_error"
            }
        }
        yield f"data: {json.dumps(error_response)}\n\n"
        yield "data: [DONE]\n\n"


def _estimate_tokens(text: str) -> int:
    """估算token数量（简单实现）"""
    # 简单的token估算：大约4个字符=1个token
    return max(1, len(text) // 4)


if __name__ == "__main__":
    uvicorn.run(
        "api:app",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
        log_level=settings.log_level.lower()
    )