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
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
from app.core.memory_kb import memory as mem
from app.core import llm
from app.core.memory_kb import knowledge
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


@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.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("/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 = await llm.resolve_web_search(
                context, model=used_model, session_id=session.id,
            )
            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 = "\n".join(r["content"][:300] for r in s["results"] if r["content"])
                    if summary:
                        knowledge.save_knowledge(db, s["query"], summary, s["results"])

            search_tools = [WEB_SEARCH_TOOL] 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=search_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",
        },
    )