playchat/backend/app/core/llm.py

import json
import logging
import time
import traceback
from typing import TYPE_CHECKING, AsyncIterator, AsyncGenerator
from contextlib import asynccontextmanager

import litellm
from app.core.config import DEFAULT_MODEL, DEBUG
from app.core.notes import NOTE_TOOLS, execute_note_tool
from app.core.search.websearch import WEB_SEARCH_TOOL, search_web

if TYPE_CHECKING:
    from sqlalchemy.orm import Session as DBSession

litellm.drop_params = True
if DEBUG:
    litellm._turn_on_debug()

logger = logging.getLogger(__name__)


@asynccontextmanager
async def _log_context(
    model: str,
    messages: list[dict],
    session_id: str | None = None,
) -> AsyncGenerator[dict, None]:
    """LLM呼び出しの前後でログデータを収集するコンテキストマネージャー。"""
    log: dict = {
        "session_id": session_id,
        "model": model,
        "request_json": json.dumps(messages, ensure_ascii=False),
        "response_text": None,
        "error": None,
        "latency_ms": None,
    }
    start = time.perf_counter()
    try:
        yield log
    except Exception as e:
        log["error"] = str(e)
        log["latency_ms"] = round((time.perf_counter() - start) * 1000, 1)
        logger.error("LLM call failed [model=%s]: %s\n%s", log["model"], e, traceback.format_exc())
        raise
    finally:
        if log["latency_ms"] is None:
            log["latency_ms"] = round((time.perf_counter() - start) * 1000, 1)


def _save_log(log: dict) -> None:
    from app.db.database import SessionLocal
    from app.db.models import ApiLog
    db = SessionLocal()
    try:
        db.add(ApiLog(**log))
        db.commit()
    finally:
        db.close()


async def stream_chat(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
    tools: list[dict] | None = None,
) -> AsyncIterator[str]:
    used_model = model or DEFAULT_MODEL
    async with _log_context(used_model, messages, session_id) as log:
        response = await litellm.acompletion(
            model=used_model,
            messages=messages,
            stream=True,
            **({"tools": tools} if tools else {}),
        )
        chunks: list[str] = []
        async for chunk in response:
            delta = chunk.choices[0].delta.content
            if delta:
                chunks.append(delta)
                yield delta
        log["response_text"] = "".join(chunks)
    _save_log(log)


async def complete(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
) -> str:
    used_model = model or DEFAULT_MODEL
    async with _log_context(used_model, messages, session_id) as log:
        response = await litellm.acompletion(
            model=used_model,
            messages=messages,
            stream=False,
        )
        result = response.choices[0].message.content or ""
        log["response_text"] = result
    _save_log(log)
    return result


async def embed(texts: list[str], model: str) -> list[list[float]]:
    """テキスト群を埋め込みベクトルに変換する。"""
    response = await litellm.aembedding(model=model, input=texts)
    return [item["embedding"] for item in response.data]


def _tool_display(name: str, args: dict) -> tuple[str, str]:
    """ツール呼び出しをGUI表示用の (ラベル, 詳細) に変換する。"""
    if name == "web_search":
        return "Web検索", args.get("query", "")
    if name == "read_note":
        return "ノートを読む", args.get("title", "")
    if name == "write_note":
        label = "ノートを上書き" if args.get("mode") == "replace" else "ノートに追記"
        return label, args.get("title", "")
    if name == "read_knowledge":
        return "蓄積知識を読む", args.get("query", "")
    return name, ""


def _first_line(text: str, limit: int = 80) -> str:
    """ツール結果の要点だけをGUIに出すため、1行目を切り出して長さを詰める。"""
    line = text.strip().splitlines()[0] if text.strip() else ""
    return line[:limit] + "…" if len(line) > limit else line


async def stream_tool_calls(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
    db: "DBSession | None" = None,
    max_rounds: int = 2,
) -> AsyncIterator[dict]:
    """必要に応じてLLMにツールを呼ばせ、その進行状況をイベントとして逐次 yield する。

    扱うツールは web_search、共有ノート（read_note / write_note）、蓄積知識（read_knowledge）。
    DBを見るものは db が渡されたときだけ提示する。モデルがtool callingに対応していない場合は
    何もせず元のmessagesをそのまま返す。

    yield するイベント:
    - {"type": "tool", "id", "name", "label", "detail", "status": "running"} … 実行開始
    - {"type": "tool", "id", "status": "done", "result"} … 実行完了（GUI用の短い結果つき）
    - {"type": "resolved", "messages", "searches", "notes_changed"} … 最後に必ず1回

    最後の resolved イベントが従来の戻り値にあたる（resolve_tool_calls がそれを取り出す）。
    """
    # app.core.memory_kb.knowledge は app.core.llm を import しているため、
    # モジュールレベルで読むと循環 import になる（_save_log と同じ理由で関数内 import する）。
    from app.core.memory_kb.knowledge import KNOWLEDGE_TOOLS, execute_knowledge_tool

    used_model = model or DEFAULT_MODEL
    working_messages = list(messages)
    searches: list[dict] = []
    notes_changed = False

    tools = [WEB_SEARCH_TOOL] + (NOTE_TOOLS + KNOWLEDGE_TOOLS if db is not None else [])

    for _ in range(max_rounds):
        try:
            response = await litellm.acompletion(
                model=used_model,
                messages=working_messages,
                tools=tools,
                tool_choice="auto",
            )
        except Exception:
            logger.info("Model %s does not support tool calling; skipping tools", used_model)
            yield {"type": "resolved", "messages": messages, "searches": searches, "notes_changed": notes_changed}
            return

        choice_message = response.choices[0].message
        tool_calls = getattr(choice_message, "tool_calls", None)
        if not tool_calls:
            break

        working_messages.append({
            "role": "assistant",
            "content": choice_message.content or "",
            "tool_calls": [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                }
                for tc in tool_calls
            ],
        })

        for tc in tool_calls:
            try:
                args = json.loads(tc.function.arguments or "{}")
            except json.JSONDecodeError:
                args = {}

            name = tc.function.name
            label, detail = _tool_display(name, args)
            yield {
                "type": "tool",
                "id": tc.id,
                "name": name,
                "label": label,
                "detail": detail,
                "status": "running",
            }

            if name in ("read_note", "write_note") and db is not None:
                tool_text, changed = execute_note_tool(db, name, args)
                notes_changed = notes_changed or changed
                result = _first_line(tool_text)
            elif name == "read_knowledge" and db is not None:
                tool_text = execute_knowledge_tool(db, name, args)
                result = _first_line(tool_text)
            else:
                query = args.get("query", "")
                results = await search_web(query)
                if results:
                    tool_text = "\n\n".join(f"[{r['title']}]({r['url']})\n{r['content']}" for r in results)
                    searches.append({"query": query, "results": results})
                    result = f"{len(results)}件の結果"
                else:
                    tool_text = "検索結果が見つかりませんでした。"
                    result = "結果なし"

            yield {"type": "tool", "id": tc.id, "status": "done", "result": result}

            working_messages.append({
                "role": "tool",
                "tool_call_id": tc.id,
                "content": tool_text,
            })

    yield {
        "type": "resolved",
        "messages": working_messages,
        "searches": searches,
        "notes_changed": notes_changed,
    }


async def resolve_tool_calls(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
    db: "DBSession | None" = None,
    max_rounds: int = 2,
) -> tuple[list[dict], list[dict], bool]:
    """stream_tool_calls の進行状況イベントを捨て、最終結果だけを返すラッパー。

    戻り値: (拡張後のmessages, 実際に検索したクエリと結果のリスト, ノートを変更したか)
    """
    resolved = (messages, [], False)
    async for event in stream_tool_calls(messages, model, session_id, db, max_rounds):
        if event["type"] == "resolved":
            resolved = (event["messages"], event["searches"], event["notes_changed"])
    return resolved