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 import notes
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
    # このセッションで本文をプロンプトに載せるノートのページID。
    # 新規セッションではまだ session_id が無く PUT できないため、送信時にまとめて反映する。
    note_page_ids: list[int] | None = None


class SessionNotesRequest(BaseModel):
    page_ids: list[int]


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
    # 要約に取り込まれ、LLMには渡らなくなったメッセージ。人間の履歴確認用に残してある。
    archived_at: datetime | None = None
    tags: list[str] = []

    class Config:
        from_attributes = True


class ImportMessage(BaseModel):
    role: str
    content: str
    created_at: datetime | None = None
    archived_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,
            archived_at=m.archived_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(),
                "archived_at": m.archived_at.isoformat() if m.archived_at else None,
                "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, archived_at=m.archived_at)
        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.get("/sessions/{session_id}/notes")
def get_session_notes(session_id: str, db: DBSession = Depends(get_db)):
    """このセッションで本文をプロンプトに載せるよう指定された共有ノートのページIDを返す。"""
    if not db.get(Session, session_id):
        raise HTTPException(status_code=404, detail="Session not found")
    return {"page_ids": notes.get_attached_page_ids(db, session_id)}


@router.put("/sessions/{session_id}/notes")
def put_session_notes(session_id: str, req: SessionNotesRequest, db: DBSession = Depends(get_db)):
    """このセッションに紐付けるノートのページIDを丸ごと置き換える。存在しないIDは無視する。"""
    if not db.get(Session, session_id):
        raise HTTPException(status_code=404, detail="Session not found")
    return {"page_ids": notes.set_attached_pages(db, session_id, req.page_ids)}


@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))

    # 紐付けるノートは送信のたびに送られてくる。新規セッションだとフロントは事前に
    # PUT できないため、コンテキストを組み立てる前にここで反映しておく。
    if req.note_page_ids is not None:
        notes.set_attached_pages(db, session.id, req.note_page_ids)

    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:
            # ツールの実行状況は tool イベントとしてそのままフロントへ流し、GUIに進行状況を出す
            resolved_context, searches, notes_changed = context, [], False
            async for event in llm.stream_tool_calls(
                context, model=used_model, session_id=session.id, db=db,
            ):
                if event["type"] == "resolved":
                    resolved_context = event["messages"]
                    searches = event["searches"]
                    notes_changed = event["notes_changed"]
                else:
                    yield sse({"type": "tool", "value": {k: v for k, v in event.items() if k != "type"}})

            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})
                yield sse({"type": "tool", "value": {
                    "id": "save_knowledge",
                    "name": "save_knowledge",
                    "label": "検索結果を要約して保存",
                    "detail": " / ".join(s["query"] for s in searches),
                    "status": "running",
                }})
                saved = 0
                for s in searches:
                    headline, summary = await knowledge.summarize_results(s["query"], s["results"])
                    if summary:
                        knowledge.save_knowledge(db, s["query"], summary, s["results"], headline=headline)
                        local_knowledge.save_search_result(s["query"], summary, s["results"])
                        saved += 1
                yield sse({"type": "tool", "value": {
                    "id": "save_knowledge",
                    "status": "done",
                    "result": f"{saved}件を蓄積知識に保存" if saved else "保存できる要約なし",
                }})

            # 履歴に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",
        },
    )