playchat/backend/app/api/chat.py

import asyncio
import uuid
from datetime import datetime, timezone
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from sqlalchemy.orm import Session as DBSession

from app.db.database import get_db
from app.db.models import Session, Message, MessageTag
from app.core.memory_kb import memory as mem
from app.core import llm
from app.core.memory_kb import knowledge
from app.core.memory_kb import local_knowledge
from app.core.notes import NOTE_TOOLS
from app.core.search.websearch import WEB_SEARCH_TOOL
from app.core.config import DEFAULT_MODEL
from app.core.memory_kb.tagging import tag_message_by_id
from app.core.memory_kb.word_memory import record_message

router = APIRouter(prefix="/chat", tags=["chat"])


class ChatRequest(BaseModel):
    session_id: str | None = None
    message: str
    model: str | None = None


class SessionResponse(BaseModel):
    id: str
    title: str
    model: str
    updated_at: datetime

    class Config:
        from_attributes = True


class MessageResponse(BaseModel):
    id: int
    role: str
    content: str
    created_at: datetime
    tags: list[str] = []

    class Config:
        from_attributes = True


class ImportMessage(BaseModel):
    role: str
    content: str
    created_at: datetime | None = None
    tags: list[str] = []


class ImportSessionMeta(BaseModel):
    title: str | None = None
    model: str | None = None
    created_at: datetime | None = None
    updated_at: datetime | None = None


class ImportRequest(BaseModel):
    session: ImportSessionMeta = ImportSessionMeta()
    messages: list[ImportMessage]


@router.get("/sessions", response_model=list[SessionResponse])
def list_sessions(db: DBSession = Depends(get_db)):
    return db.query(Session).order_by(Session.updated_at.desc()).all()


@router.get("/sessions/{session_id}/messages", response_model=list[MessageResponse])
def get_messages(session_id: str, db: DBSession = Depends(get_db)):
    session = db.get(Session, session_id)
    if not session:
        raise HTTPException(status_code=404, detail="Session not found")
    return [
        MessageResponse(
            id=m.id,
            role=m.role,
            content=m.content,
            created_at=m.created_at,
            tags=[t.tag for t in m.tags],
        )
        for m in session.messages
    ]


@router.get("/sessions/{session_id}/export")
def export_session(session_id: str, db: DBSession = Depends(get_db)):
    """セッションのメタデータと全メッセージをJSONファイルとしてダウンロードする。"""
    session = db.get(Session, session_id)
    if not session:
        raise HTTPException(status_code=404, detail="Session not found")

    data = {
        "session": {
            "id": session.id,
            "title": session.title,
            "model": session.model,
            "created_at": session.created_at.isoformat(),
            "updated_at": session.updated_at.isoformat(),
        },
        "messages": [
            {
                "id": m.id,
                "role": m.role,
                "content": m.content,
                "created_at": m.created_at.isoformat(),
                "tags": [t.tag for t in m.tags],
            }
            for m in session.messages
        ],
    }
    return JSONResponse(
        content=data,
        headers={"Content-Disposition": f'attachment; filename="session_{session.id}.json"'},
    )


@router.post("/sessions/import", response_model=SessionResponse, status_code=201)
def import_session(req: ImportRequest, db: DBSession = Depends(get_db)):
    """export で出力したJSON（またはその形式に沿ったJSON）から新規セッションを作成する。"""
    if not req.messages:
        raise HTTPException(status_code=400, detail="messages is empty")
    for m in req.messages:
        if m.role not in ("user", "assistant", "system"):
            raise HTTPException(status_code=400, detail=f"invalid role: {m.role}")

    title = req.session.title or req.messages[0].content[:40]
    session = Session(
        id=str(uuid.uuid4()),
        title=title,
        model=req.session.model or DEFAULT_MODEL,
    )
    if req.session.created_at:
        session.created_at = req.session.created_at
    session.updated_at = req.session.updated_at or datetime.now(timezone.utc)
    db.add(session)
    db.flush()

    for m in req.messages:
        msg = Message(session_id=session.id, role=m.role, content=m.content)
        if m.created_at:
            msg.created_at = m.created_at
        db.add(msg)
        db.flush()
        for tag in m.tags:
            db.add(MessageTag(message_id=msg.id, tag=tag))

    db.commit()
    db.refresh(session)
    return session


@router.delete("/sessions/{session_id}", status_code=204)
def delete_session(session_id: str, db: DBSession = Depends(get_db)):
    session = db.get(Session, session_id)
    if not session:
        raise HTTPException(status_code=404, detail="Session not found")
    db.delete(session)
    db.commit()


@router.post("/knowledge/regenerate")
async def regenerate_knowledge(db: DBSession = Depends(get_db)):
    """既存のKnowledgeEntryを出典URLから取得し直し、構造化要約に再生成する（一括実行用）。"""
    count = await knowledge.regenerate_from_sources(db)
    return {"regenerated": count}


@router.post("/stream")
async def stream_chat(req: ChatRequest, db: DBSession = Depends(get_db)):
    # セッション取得 or 新規作成
    if req.session_id:
        session = db.get(Session, req.session_id)
        if not session:
            raise HTTPException(status_code=404, detail="Session not found")
    else:
        session = Session(
            id=str(uuid.uuid4()),
            title=req.message[:40],
            model=req.model or DEFAULT_MODEL,
        )
        db.add(session)
        db.commit()

    # ユーザーメッセージを保存
    user_msg = Message(session_id=session.id, role="user", content=req.message)
    db.add(user_msg)
    db.commit()
    db.refresh(user_msg)

    record_message(req.message)
    asyncio.create_task(tag_message_by_id(user_msg.id))

    context = await mem.get_context_messages(db, session.id)

    async def generate():
        import json as _json

        def sse(payload: dict) -> str:
            return f"data: {_json.dumps(payload, ensure_ascii=False)}\n\n"

        # セッションIDをまず返す（新規セッション時にフロントへ通知）
        yield sse({"type": "session_id", "value": session.id})

        used_model = req.model or session.model
        full_response = []
        try:
            resolved_context, searches, notes_changed = await llm.resolve_tool_calls(
                context, model=used_model, session_id=session.id, db=db,
            )
            if notes_changed:
                # AIがノートを書き換えたので、フロントのノートパネルを再読み込みさせる
                yield sse({"type": "notes_updated"})
            if searches:
                sources = [
                    {"title": r["title"], "url": r["url"]}
                    for s in searches for r in s["results"]
                ]
                yield sse({"type": "sources", "value": sources})
                for s in searches:
                    summary = await knowledge.summarize_results(s["query"], s["results"])
                    if summary:
                        knowledge.save_knowledge(db, s["query"], summary, s["results"])
                        local_knowledge.save_search_result(s["query"], summary, s["results"])

            # 履歴にtool_callsが含まれる場合、最終呼び出しにも同じツール定義を渡す必要がある
            used_tools = [WEB_SEARCH_TOOL, *NOTE_TOOLS] if len(resolved_context) > len(context) else None
            async for chunk in llm.stream_chat(
                resolved_context, model=used_model, session_id=session.id, tools=used_tools,
            ):
                full_response.append(chunk)
                yield sse({"type": "chunk", "value": chunk})

            # アシスタントメッセージを保存
            assistant_content = "".join(full_response)
            assistant_msg = Message(session_id=session.id, role="assistant", content=assistant_content)
            db.add(assistant_msg)
            session.updated_at = datetime.now(timezone.utc)
            db.commit()
            db.refresh(assistant_msg)

            record_message(assistant_content)
            asyncio.create_task(tag_message_by_id(assistant_msg.id))

            # メモリ要約チェック
            await mem.maybe_summarize(db, session)

            yield sse({"type": "done"})

        except Exception as e:
            yield sse({"type": "error", "value": str(e)})

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        },
    )