Commit d9b6c146 authored by Kantz's avatar Kantz
Browse files

extractet tool_calls

parent 0c34d6f6
......@@ -12,7 +12,8 @@ from fastapi import APIRouter
from fastapi.responses import JSONResponse
import app.config as config
from app.deterministic_services import llm_client, llm_quota
from app.deterministic_services import llm_quota
from app.deterministic_services.llm_tool_client import resolve_mcp_base_url, with_mcp_session
router = APIRouter()
logger = logging.getLogger(__name__)
......@@ -161,7 +162,7 @@ async def _probe_mcp_server() -> dict:
mcp_settings = config.get_mcp_settings()
try:
url = llm_client._resolve_mcp_base_url(mcp_settings)
url = resolve_mcp_base_url(mcp_settings)
except ValueError as exc:
return {"status": "missing_config", "detail": str(exc)}
......@@ -173,7 +174,7 @@ async def _probe_mcp_server() -> dict:
}
try:
return await llm_client._with_mcp_session(mcp_settings, _handler)
return await with_mcp_session(mcp_settings, _handler)
except Exception as exc:
return {"status": "error", "url": url, "detail": str(exc)}
......
# LLM client with support for multiple providers and tool use, including quota tracking and tool call logging.
import asyncio
import inspect
import json
from datetime import date
from typing import Any, Callable
import httpx
import ollama
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mistralai.client import Mistral
from openai import AsyncOpenAI, OpenAI
from openai import OpenAI
from app import config
from app.deterministic_services import llm_quota, tool_log_context
AUTH_HEADER = "X-MCP-Shared-Secret"
from app.deterministic_services import llm_quota
from app.deterministic_services.llm_tool_client import (
append_tool_trace_log,
resolve_mcp_base_url,
run_tool_chat_sync,
with_mcp_session,
)
# ---
......@@ -39,14 +38,6 @@ def _extract_message(response) -> object:
return {}
def _extract_tool_calls(message: object) -> list[Any]:
if isinstance(message, dict):
return list(message.get("tool_calls") or [])
if hasattr(message, "tool_calls"):
return list(getattr(message, "tool_calls") or [])
return []
def _extract_total_tokens(response: object) -> int:
usage = None
if isinstance(response, dict):
......@@ -73,91 +64,6 @@ def _extract_total_tokens(response: object) -> int:
return int(prompt) + int(completion)
def _normalize_message(message: object) -> dict[str, Any]:
if isinstance(message, dict):
return dict(message)
normalized: dict[str, Any] = {}
for attr in ("role", "content", "tool_calls", "name"):
if hasattr(message, attr):
value = getattr(message, attr)
if value is not None:
normalized[attr] = value
return normalized
def _extract_tool_call_name(tool_call: object) -> str:
function = getattr(tool_call, "function", None)
if function is not None and hasattr(function, "name"):
return getattr(function, "name")
if isinstance(tool_call, dict):
function_dict = tool_call.get("function") or {}
return str(function_dict.get("name") or "")
return ""
def _extract_tool_call_arguments(tool_call: object) -> dict[str, Any]:
function = getattr(tool_call, "function", None)
raw_arguments: Any = None
if function is not None and hasattr(function, "arguments"):
raw_arguments = getattr(function, "arguments")
elif isinstance(tool_call, dict):
function_dict = tool_call.get("function") or {}
raw_arguments = function_dict.get("arguments")
if isinstance(raw_arguments, dict):
return raw_arguments
if isinstance(raw_arguments, str) and raw_arguments.strip():
return json.loads(raw_arguments)
return {}
def _extract_tool_call_id(tool_call: object) -> str:
if hasattr(tool_call, "id"):
return str(getattr(tool_call, "id") or "")
if isinstance(tool_call, dict):
return str(tool_call.get("id") or "")
return ""
def _extract_structured_tool_output(tool_result: Any) -> Any:
if hasattr(tool_result, "structured_content") and tool_result.structured_content is not None:
return tool_result.structured_content
if hasattr(tool_result, "structuredContent") and tool_result.structuredContent is not None:
return tool_result.structuredContent
blocks = []
for item in getattr(tool_result, "content", []) or []:
text = getattr(item, "text", None)
if text is not None:
blocks.append(text)
return "\n".join(blocks)
def _append_tool_trace_log(result: dict, provider: str) -> None:
tool_trace = result.get("tool_trace")
if not tool_trace:
return
tool_log = tool_log_context.get_active_tool_log()
if tool_log is None:
return
tool_log.append(
{
"name": "llm_tool_trace",
"arguments": {
"provider": provider,
"tool_call_count": len(tool_trace),
},
"response": {
"tool_trace": tool_trace,
"response_ids": result.get("response_ids", []),
},
}
)
# ---
# Quota
# ---
......@@ -286,257 +192,12 @@ def _chat_ollama(
return {"raw": response, "message": _extract_message(response)}
# ---
# MCP tool execution
# ---
def _resolve_mcp_base_url(mcp_settings: dict[str, Any]) -> str:
base_url = mcp_settings.get("base_url")
if not base_url:
raise ValueError("MCP_BASE_URL must be set")
return str(base_url)
async def _with_mcp_session(
mcp_settings: dict[str, Any],
handler,
http_client: httpx.AsyncClient | None = None,
):
headers: dict[str, str] = {}
shared_secret = mcp_settings.get("shared_secret")
if shared_secret:
headers[AUTH_HEADER] = str(shared_secret)
managed_http_client = http_client
if managed_http_client is None:
managed_http_client = httpx.AsyncClient(headers=headers, timeout=30.0)
elif headers:
managed_http_client.headers.update(headers)
should_close_http_client = http_client is None
try:
async with streamable_http_client(
_resolve_mcp_base_url(mcp_settings),
http_client=managed_http_client,
) as streams:
async with ClientSession(streams[0], streams[1]) as session:
await session.initialize()
listed_tools = await session.list_tools()
return await handler(session, listed_tools.tools)
finally:
if should_close_http_client:
await managed_http_client.aclose()
def _tool_attr(tool: Any, snake_name: str, camel_name: str) -> Any:
if hasattr(tool, snake_name):
return getattr(tool, snake_name)
return getattr(tool, camel_name)
def _tool_to_openai_chat_tool(tool: Any) -> dict[str, Any]:
return {
"type": "function",
"function": {
"name": _tool_attr(tool, "name", "name"),
"description": _tool_attr(tool, "description", "description") or "",
"parameters": _tool_attr(tool, "input_schema", "inputSchema"),
},
}
def _tool_to_ollama_chat_tool(tool: Any) -> dict[str, Any]:
return {
"type": "function",
"function": {
"name": _tool_attr(tool, "name", "name"),
"description": _tool_attr(tool, "description", "description") or "",
"parameters": _tool_attr(tool, "input_schema", "inputSchema"),
},
}
def _append_openai_tool_message(
messages: list[dict[str, Any]],
tool_call_id: str,
tool_name: str,
structured_output: Any,
) -> None:
messages.append(
{
"role": "tool",
"tool_call_id": tool_call_id,
"name": tool_name,
"content": json.dumps(structured_output),
}
)
def _append_ollama_tool_message(
messages: list[dict[str, Any]],
tool_name: str,
structured_output: Any,
) -> None:
messages.append(
{
"role": "tool",
"tool_name": tool_name,
"content": json.dumps(structured_output),
}
)
async def _run_openai_compatible_tool_loop(
messages: list[dict],
settings: config.OpenAIChatSettings,
http_client: httpx.AsyncClient | None = None,
openai_client: Any | None = None,
) -> dict:
mcp_settings = config.get_mcp_settings()
client_kwargs: dict[str, Any] = {"api_key": settings.api_key}
if settings.base_url:
client_kwargs["base_url"] = settings.base_url
if settings.timeout is not None:
client_kwargs["timeout"] = settings.timeout
client = openai_client or AsyncOpenAI(**client_kwargs)
async def _handler(session, tools) -> dict:
chat_tools = [_tool_to_openai_chat_tool(tool) for tool in tools]
request_messages = [dict(message) for message in messages]
response_ids: list[str] = []
tool_trace: list[dict[str, Any]] = []
while True:
request_kwargs: dict[str, Any] = {
"model": settings.model,
"messages": request_messages,
"tools": chat_tools,
}
if settings.temperature is not None:
request_kwargs["temperature"] = settings.temperature
response = await client.chat.completions.create(**request_kwargs)
response_id = getattr(response, "id", "")
if response_id:
response_ids.append(response_id)
message = response.choices[0].message if response.choices else {}
tool_calls = _extract_tool_calls(message)
if not tool_calls:
return {
"raw": response,
"message": message,
"tool_trace": tool_trace,
"response_ids": response_ids,
}
request_messages.append(_normalize_message(message))
for tool_call in tool_calls:
tool_name = _extract_tool_call_name(tool_call)
arguments = _extract_tool_call_arguments(tool_call)
tool_result = await session.call_tool(tool_name, arguments=arguments)
structured_output = _extract_structured_tool_output(tool_result)
tool_trace.append(
{
"name": tool_name,
"arguments": arguments,
"output": structured_output,
}
)
_append_openai_tool_message(
request_messages,
_extract_tool_call_id(tool_call),
tool_name,
structured_output,
)
return await _with_mcp_session(mcp_settings, _handler, http_client=http_client)
async def _run_ollama_tool_loop(
messages: list[dict],
http_client: httpx.AsyncClient | None = None,
ollama_client: Any | None = None,
) -> dict:
settings = config.get_ollama_settings()
mcp_settings = config.get_mcp_settings()
client_kwargs: dict[str, Any] = {"host": settings.base_url}
if settings.timeout is not None:
client_kwargs["timeout"] = settings.timeout
client = ollama_client or ollama.AsyncClient(**client_kwargs)
async def _handler(session, tools) -> dict:
chat_tools = [_tool_to_ollama_chat_tool(tool) for tool in tools]
request_messages = [dict(message) for message in messages]
response_ids: list[str] = []
tool_trace: list[dict[str, Any]] = []
while True:
request_kwargs: dict[str, Any] = {
"model": settings.model,
"messages": request_messages,
"tools": chat_tools,
}
if settings.keepalive is not None:
request_kwargs["keep_alive"] = settings.keepalive
if settings.temperature is not None:
request_kwargs["options"] = {"temperature": settings.temperature}
completion = await client.chat(**request_kwargs)
completion_id = getattr(completion, "id", "")
if completion_id:
response_ids.append(completion_id)
message = _extract_message(completion)
tool_calls = _extract_tool_calls(message)
if not tool_calls:
return {
"raw": completion,
"message": message,
"tool_trace": tool_trace,
"response_ids": response_ids,
}
request_messages.append(_normalize_message(message))
for tool_call in tool_calls:
tool_name = _extract_tool_call_name(tool_call)
arguments = _extract_tool_call_arguments(tool_call)
tool_result = await session.call_tool(tool_name, arguments=arguments)
structured_output = _extract_structured_tool_output(tool_result)
tool_trace.append(
{
"name": tool_name,
"arguments": arguments,
"output": structured_output,
}
)
_append_ollama_tool_message(request_messages, tool_name, structured_output)
return await _with_mcp_session(mcp_settings, _handler, http_client=http_client)
def _run_tool_chat_sync(
provider: str,
messages: list[dict],
settings: config.OpenAIChatSettings | None = None,
) -> dict:
if provider in {"openai", "gwdg"}:
assert settings is not None
return asyncio.run(_run_openai_compatible_tool_loop(messages, settings))
if provider == "ollama":
return asyncio.run(_run_ollama_tool_loop(messages))
raise ValueError(f"Unsupported tool-enabled LLM_PROVIDER: {provider}")
def _run_provider_chat(
provider: str,
chat_func: Callable[[], dict],
) -> dict:
result = _quota_tracked_chat(chat_func)
_append_tool_trace_log(result, provider)
append_tool_trace_log(result, provider)
return result
......@@ -564,7 +225,7 @@ def chat(
if tool_use_enabled:
return _run_provider_chat(
provider,
lambda: _run_tool_chat_sync(provider, messages, settings),
lambda: run_tool_chat_sync(provider, messages, settings),
)
return _run_provider_chat(provider, lambda: _chat_openai_compatible(messages, settings))
if provider == "gwdg":
......@@ -572,7 +233,7 @@ def chat(
if tool_use_enabled:
return _run_provider_chat(
provider,
lambda: _run_tool_chat_sync(provider, messages, settings),
lambda: run_tool_chat_sync(provider, messages, settings),
)
return _run_provider_chat(provider, lambda: _chat_openai_compatible(messages, settings))
if provider == "mistral":
......@@ -580,6 +241,6 @@ def chat(
return _run_provider_chat(provider, lambda: _chat_mistral(messages, settings))
if provider == "ollama":
if tool_use_enabled:
return _run_provider_chat(provider, lambda: _run_tool_chat_sync(provider, messages))
return _run_provider_chat(provider, lambda: run_tool_chat_sync(provider, messages))
return _run_provider_chat(provider, lambda: _chat_ollama(messages))
raise ValueError(f"Unsupported LLM_PROVIDER: {provider}")
import asyncio
import json
from typing import Any
import httpx
import ollama
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamable_http_client
from openai import AsyncOpenAI
from app import config
from app.deterministic_services import tool_log_context
AUTH_HEADER = "X-MCP-Shared-Secret"
def _extract_message(response) -> object:
if isinstance(response, dict) and "message" in response:
return response["message"]
if hasattr(response, "message"):
return response.message
return {}
def _extract_tool_calls(message: object) -> list[Any]:
if isinstance(message, dict):
return list(message.get("tool_calls") or [])
if hasattr(message, "tool_calls"):
return list(getattr(message, "tool_calls") or [])
return []
def _normalize_message(message: object) -> dict[str, Any]:
if isinstance(message, dict):
return dict(message)
normalized: dict[str, Any] = {}
for attr in ("role", "content", "tool_calls", "name"):
if hasattr(message, attr):
value = getattr(message, attr)
if value is not None:
normalized[attr] = value
return normalized
def _extract_tool_call_name(tool_call: object) -> str:
function = getattr(tool_call, "function", None)
if function is not None and hasattr(function, "name"):
return getattr(function, "name")
if isinstance(tool_call, dict):
function_dict = tool_call.get("function") or {}
return str(function_dict.get("name") or "")
return ""
def _extract_tool_call_arguments(tool_call: object) -> dict[str, Any]:
function = getattr(tool_call, "function", None)
raw_arguments: Any = None
if function is not None and hasattr(function, "arguments"):
raw_arguments = getattr(function, "arguments")
elif isinstance(tool_call, dict):
function_dict = tool_call.get("function") or {}
raw_arguments = function_dict.get("arguments")
if isinstance(raw_arguments, dict):
return raw_arguments
if isinstance(raw_arguments, str) and raw_arguments.strip():
return json.loads(raw_arguments)
return {}
def _extract_tool_call_id(tool_call: object) -> str:
if hasattr(tool_call, "id"):
return str(getattr(tool_call, "id") or "")
if isinstance(tool_call, dict):
return str(tool_call.get("id") or "")
return ""
def _extract_structured_tool_output(tool_result: Any) -> Any:
if hasattr(tool_result, "structured_content") and tool_result.structured_content is not None:
return tool_result.structured_content
if hasattr(tool_result, "structuredContent") and tool_result.structuredContent is not None:
return tool_result.structuredContent
blocks = []
for item in getattr(tool_result, "content", []) or []:
text = getattr(item, "text", None)
if text is not None:
blocks.append(text)
return "\n".join(blocks)
def append_tool_trace_log(result: dict, provider: str) -> None:
tool_trace = result.get("tool_trace")
if not tool_trace:
return
tool_log = tool_log_context.get_active_tool_log()
if tool_log is None:
return
tool_log.append(
{
"name": "llm_tool_trace",
"arguments": {
"provider": provider,
"tool_call_count": len(tool_trace),
},
"response": {
"tool_trace": tool_trace,
"response_ids": result.get("response_ids", []),
},
}
)
def resolve_mcp_base_url(mcp_settings: dict[str, Any]) -> str:
base_url = mcp_settings.get("base_url")
if not base_url:
raise ValueError("MCP_BASE_URL must be set")
return str(base_url)
async def with_mcp_session(
mcp_settings: dict[str, Any],
handler,
http_client: httpx.AsyncClient | None = None,
):
headers: dict[str, str] = {}
shared_secret = mcp_settings.get("shared_secret")
if shared_secret:
headers[AUTH_HEADER] = str(shared_secret)
managed_http_client = http_client
if managed_http_client is None:
managed_http_client = httpx.AsyncClient(headers=headers, timeout=30.0)
elif headers:
managed_http_client.headers.update(headers)
should_close_http_client = http_client is None
try:
async with streamable_http_client(
resolve_mcp_base_url(mcp_settings),
http_client=managed_http_client,
) as streams:
async with ClientSession(streams[0], streams[1]) as session:
await session.initialize()
listed_tools = await session.list_tools()
return await handler(session, listed_tools.tools)
finally:
if should_close_http_client:
await managed_http_client.aclose()
def _tool_attr(tool: Any, snake_name: str, camel_name: str) -> Any:
if hasattr(tool, snake_name):
return getattr(tool, snake_name)
return getattr(tool, camel_name)
def _tool_to_chat_tool(tool: Any) -> dict[str, Any]:
return {
"type": "function",
"function": {
"name": _tool_attr(tool, "name", "name"),
"description": _tool_attr(tool, "description", "description") or "",
"parameters": _tool_attr(tool, "input_schema", "inputSchema"),
},
}
def _append_openai_tool_message(
messages: list[dict[str, Any]],
tool_call_id: str,
tool_name: str,
structured_output: Any,
) -> None:
messages.append(
{
"role": "tool",
"tool_call_id": tool_call_id,
"name": tool_name,
"content": json.dumps(structured_output),
}
)
def _append_ollama_tool_message(
messages: list[dict[str, Any]],
tool_name: str,
structured_output: Any,
) -> None:
messages.append(
{
"role": "tool",
"tool_name": tool_name,
"content": json.dumps(structured_output),
}
)
async def _run_openai_compatible_tool_loop(
messages: list[dict],
settings: config.OpenAIChatSettings,
http_client: httpx.AsyncClient | None = None,
openai_client: Any | None = None,
) -> dict:
mcp_settings = config.get_mcp_settings()
client_kwargs: dict[str, Any] = {"api_key": settings.api_key}
if settings.base_url:
client_kwargs["base_url"] = settings.base_url
if settings.timeout is not None:
client_kwargs["timeout"] = settings.timeout
client = openai_client or AsyncOpenAI(**client_kwargs)
async def _handler(session, tools) -> dict:
chat_tools = [_tool_to_chat_tool(tool) for tool in tools]
request_messages = [dict(message) for message in messages]
response_ids: list[str] = []
tool_trace: list[dict[str, Any]] = []
while True:
request_kwargs: dict[str, Any] = {
"model": settings.model,
"messages": request_messages,
"tools": chat_tools,
}
if settings.temperature is not None:
request_kwargs["temperature"] = settings.temperature
response = await client.chat.completions.create(**request_kwargs)
response_id = getattr(response, "id", "")
if response_id:
response_ids.append(response_id)
message = response.choices[0].message if response.choices else {}
tool_calls = _extract_tool_calls(message)
if not tool_calls:
return {
"raw": response,
"message": message,
"tool_trace": tool_trace,
"response_ids": response_ids,
}
request_messages.append(_normalize_message(message))
for tool_call in tool_calls:
tool_name = _extract_tool_call_name(tool_call)
arguments = _extract_tool_call_arguments(tool_call)
tool_result = await session.call_tool(tool_name, arguments=arguments)
structured_output = _extract_structured_tool_output(tool_result)
tool_trace.append(
{
"name": tool_name,
"arguments": arguments,
"output": structured_output,
}
)
_append_openai_tool_message(
request_messages,
_extract_tool_call_id(tool_call),
tool_name,
structured_output,
)
return await with_mcp_session(mcp_settings, _handler, http_client=http_client)
async def _run_ollama_tool_loop(
messages: list[dict],
http_client: httpx.AsyncClient | None = None,
ollama_client: Any | None = None,
) -> dict:
settings = config.get_ollama_settings()
mcp_settings = config.get_mcp_settings()
client_kwargs: dict[str, Any] = {"host": settings.base_url}
if settings.timeout is not None:
client_kwargs["timeout"] = settings.timeout
client = ollama_client or ollama.AsyncClient(**client_kwargs)
async def _handler(session, tools) -> dict:
chat_tools = [_tool_to_chat_tool(tool) for tool in tools]
request_messages = [dict(message) for message in messages]
response_ids: list[str] = []
tool_trace: list[dict[str, Any]] = []
while True:
request_kwargs: dict[str, Any] = {
"model": settings.model,
"messages": request_messages,
"tools": chat_tools,
}
if settings.keepalive is not None:
request_kwargs["keep_alive"] = settings.keepalive
if settings.temperature is not None:
request_kwargs["options"] = {"temperature": settings.temperature}
completion = await client.chat(**request_kwargs)
completion_id = getattr(completion, "id", "")
if completion_id:
response_ids.append(completion_id)
message = _extract_message(completion)
tool_calls = _extract_tool_calls(message)
if not tool_calls:
return {
"raw": completion,
"message": message,
"tool_trace": tool_trace,
"response_ids": response_ids,
}
request_messages.append(_normalize_message(message))
for tool_call in tool_calls:
tool_name = _extract_tool_call_name(tool_call)
arguments = _extract_tool_call_arguments(tool_call)
tool_result = await session.call_tool(tool_name, arguments=arguments)
structured_output = _extract_structured_tool_output(tool_result)
tool_trace.append(
{
"name": tool_name,
"arguments": arguments,
"output": structured_output,
}
)
_append_ollama_tool_message(request_messages, tool_name, structured_output)
return await with_mcp_session(mcp_settings, _handler, http_client=http_client)
def run_tool_chat_sync(
provider: str,
messages: list[dict],
settings: config.OpenAIChatSettings | None = None,
) -> dict:
if provider in {"openai", "gwdg"}:
assert settings is not None
return asyncio.run(_run_openai_compatible_tool_loop(messages, settings))
if provider == "ollama":
return asyncio.run(_run_ollama_tool_loop(messages))
raise ValueError(f"Unsupported tool-enabled LLM_PROVIDER: {provider}")
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment