playchat/backend/app/core/llm.py

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

import litellm
from app.core.config import DEFAULT_MODEL, DEBUG
from app.core.websearch import WEB_SEARCH_TOOL, search_web

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 resolve_web_search(
    messages: list[dict],
    model: str | None = None,
    session_id: str | None = None,
    max_rounds: int = 2,
) -> tuple[list[dict], list[dict]]:
    """必要に応じてLLMにweb_searchツールを呼ばせ、検索結果を会話に組み込んだ上でメッセージ列を返す。

    モデルがtool callingに対応していない場合は何もせず元のmessagesをそのまま返す。
    戻り値: (拡張後のmessages, 実際に検索したクエリと結果のリスト)
    """
    used_model = model or DEFAULT_MODEL
    working_messages = list(messages)
    searches: list[dict] = []

    for _ in range(max_rounds):
        try:
            response = await litellm.acompletion(
                model=used_model,
                messages=working_messages,
                tools=[WEB_SEARCH_TOOL],
                tool_choice="auto",
            )
        except Exception:
            logger.info("Model %s does not support tool calling; skipping web search", used_model)
            return messages, searches

        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 = {}
            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})
            else:
                tool_text = "検索結果が見つかりませんでした。"

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

    return working_messages, searches