dp-zp-agent/evaluation/rag_runner.py

969 lines
18 KiB
Python

from __future__ import annotations
import json
import os
import re
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any
PROJECT_ROOT = Path(
__file__
).resolve().parents[1]
OPENWEBUI_URL = (
"https://ui.tukekemt.xyz/api/chat/completions"
)
LOCAL_OPENAPI_URL = (
"http://localhost:8000/openapi.json"
)
LOCAL_RAG_URL = (
"http://localhost:8000/rag"
)
DEFAULT_MODEL = (
"model120-fast"
)
DEFAULT_TIMEOUT = 180
MARKDOWN_URL_RE = re.compile(
r"\[[^\]]*\]\((https?://[^)]+)\)"
)
def load_env_value(
key: str,
env_path: Path | None = None,
) -> str:
value = os.environ.get(
key
)
if value:
return value
if env_path is None:
env_path = (
PROJECT_ROOT
/ ".env"
)
if not env_path.exists():
raise RuntimeError(
f"{key} nie je v environment "
f"a {env_path} neexistuje."
)
for raw_line in env_path.read_text(
encoding="utf-8"
).splitlines():
line = raw_line.strip()
if (
not line
or line.startswith("#")
or "=" not in line
):
continue
name, value = line.split(
"=",
1,
)
if name.strip() != key:
continue
value = value.strip()
if (
len(value) >= 2
and value[0] == value[-1]
and value[0] in {
"'",
'"',
}
):
value = value[
1:-1
]
if value:
return value
raise RuntimeError(
f"{key} sa nepodarilo nájsť."
)
def request_json(
url: str,
*,
method: str = "GET",
headers: dict[str, str] | None = None,
payload: dict[str, Any] | None = None,
timeout: int = DEFAULT_TIMEOUT,
) -> dict[str, Any]:
data = None
if payload is not None:
data = json.dumps(
payload,
ensure_ascii=False,
).encode(
"utf-8"
)
request = urllib.request.Request(
url,
data=data,
method=method,
headers=(
headers
or {}
),
)
try:
with urllib.request.urlopen(
request,
timeout=timeout,
) as response:
raw = (
response
.read()
.decode(
"utf-8"
)
)
except urllib.error.HTTPError as exc:
body = (
exc.read()
.decode(
"utf-8",
errors="replace",
)
)
raise RuntimeError(
f"HTTP {exc.code} pre {url}: "
f"{body[:1500]}"
) from exc
except urllib.error.URLError as exc:
raise RuntimeError(
f"Sieťová chyba pre {url}: "
f"{exc}"
) from exc
if not raw.strip():
raise RuntimeError(
f"Prázdna odpoveď z {url}."
)
try:
parsed = json.loads(
raw
)
except json.JSONDecodeError as exc:
raise RuntimeError(
f"Neplatný JSON z {url}: "
f"{raw[:1000]}"
) from exc
if not isinstance(
parsed,
dict,
):
raise RuntimeError(
"Očakávaný JSON objekt "
f"z {url}, dostal som "
f"{type(parsed).__name__}."
)
return parsed
def resolve_refs(
value: Any,
document: dict[str, Any],
) -> Any:
if isinstance(
value,
list,
):
return [
resolve_refs(
item,
document,
)
for item in value
]
if not isinstance(
value,
dict,
):
return value
ref = value.get(
"$ref"
)
if ref:
if not ref.startswith(
"#/"
):
raise RuntimeError(
"Nepodporovaný "
f"OpenAPI $ref: {ref}"
)
current: Any = document
for part in ref[
2:
].split("/"):
current = current[
part
]
resolved = resolve_refs(
current,
document,
)
extra = {
key: item
for key, item
in value.items()
if key != "$ref"
}
if (
extra
and isinstance(
resolved,
dict,
)
):
resolved = {
**resolved,
**resolve_refs(
extra,
document,
),
}
return resolved
return {
key: resolve_refs(
item,
document,
)
for key, item
in value.items()
}
def build_rag_tool(
openapi: dict[str, Any],
) -> tuple[
str,
dict[str, Any],
]:
try:
operation = (
openapi[
"paths"
][
"/rag"
][
"post"
]
)
except KeyError as exc:
raise RuntimeError(
"OpenAPI schéma neobsahuje "
"POST /rag."
) from exc
operation_id = operation.get(
"operationId"
)
if not operation_id:
raise RuntimeError(
"POST /rag nemá operationId."
)
try:
schema = (
operation[
"requestBody"
][
"content"
][
"application/json"
][
"schema"
]
)
except KeyError as exc:
raise RuntimeError(
"POST /rag nemá "
"request JSON schema."
) from exc
parameters = resolve_refs(
schema,
openapi,
)
description = (
operation.get(
"description"
)
or operation.get(
"summary"
)
or (
"Vyhľadá relevantný kontext "
"v dokumentoch ZP Wiki."
)
)
tool = {
"type": "function",
"function": {
"name": operation_id,
"description": (
description
),
"parameters": (
parameters
),
},
}
return (
operation_id,
tool,
)
def normalize_url(
value: str,
) -> str:
value = value.strip()
match = (
MARKDOWN_URL_RE.search(
value
)
)
if match:
value = match.group(
1
)
return value.rstrip(
"/"
)
def extract_urls_from_object(
value: Any,
) -> list[str]:
result: list[
str
] = []
if isinstance(
value,
dict,
):
for key, item in value.items():
if (
key == "source_url"
and isinstance(
item,
str,
)
):
result.append(
normalize_url(
item
)
)
result.extend(
extract_urls_from_object(
item
)
)
elif isinstance(
value,
list,
):
for item in value:
result.extend(
extract_urls_from_object(
item
)
)
return list(
dict.fromkeys(
result
)
)
def get_usage(
response: dict[str, Any],
) -> dict[str, int]:
usage = response.get(
"usage"
)
if not isinstance(
usage,
dict,
):
return {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
}
return {
"prompt_tokens": int(
usage.get(
"prompt_tokens",
0,
)
or 0
),
"completion_tokens": int(
usage.get(
"completion_tokens",
0,
)
or 0
),
"total_tokens": int(
usage.get(
"total_tokens",
0,
)
or 0
),
}
def add_usage(
total: dict[str, int],
current: dict[str, int],
) -> None:
for key in (
"prompt_tokens",
"completion_tokens",
"total_tokens",
):
total[
key
] += current[
key
]
def get_first_message(
response: dict[str, Any],
) -> dict[str, Any]:
choices = response.get(
"choices"
)
if (
not isinstance(
choices,
list,
)
or not choices
):
raise RuntimeError(
"Model nevrátil choices."
)
choice = choices[
0
]
if not isinstance(
choice,
dict,
):
raise RuntimeError(
"Neplatný choices[0]."
)
message = choice.get(
"message"
)
if not isinstance(
message,
dict,
):
raise RuntimeError(
"Model nevrátil message."
)
return message
def parse_tool_arguments(
raw_arguments: Any,
) -> dict[str, Any]:
if isinstance(
raw_arguments,
dict,
):
return raw_arguments
if not isinstance(
raw_arguments,
str,
):
raise RuntimeError(
"Neplatný formát "
"tool arguments."
)
try:
parsed = json.loads(
raw_arguments
)
except json.JSONDecodeError as exc:
raise RuntimeError(
"Model vrátil neplatné JSON "
"argumenty toolu: "
f"{raw_arguments[:1000]}"
) from exc
if not isinstance(
parsed,
dict,
):
raise RuntimeError(
"Tool arguments nie sú "
"JSON objekt."
)
return parsed
def run_question(
question: str,
*,
model: str,
operation_id: str,
rag_tool: dict[str, Any],
openwebui_api_key: str,
search_api_key: str,
timeout: int,
) -> dict[str, Any]:
question = question.strip()
if not question:
raise RuntimeError(
"Otázka je prázdna."
)
started = (
time.perf_counter()
)
usage_total = {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
}
tool_calls_record: list[
dict[str, Any]
] = []
rag_source_urls: list[
str
] = []
messages: list[
dict[str, Any]
] = [
{
"role": "user",
"content": question,
}
]
first_started = (
time.perf_counter()
)
first_response = request_json(
OPENWEBUI_URL,
method="POST",
headers={
"Authorization": (
"Bearer "
f"{openwebui_api_key}"
),
"Content-Type": (
"application/json"
),
},
payload={
"model": model,
"messages": messages,
"tools": [
rag_tool
],
"tool_choice": "auto",
"stream": False,
},
timeout=timeout,
)
first_latency = (
time.perf_counter()
- first_started
)
add_usage(
usage_total,
get_usage(
first_response
),
)
first_message = (
get_first_message(
first_response
)
)
tool_calls = (
first_message.get(
"tool_calls"
)
or []
)
if not isinstance(
tool_calls,
list,
):
tool_calls = []
if not tool_calls:
answer = str(
first_message.get(
"content"
)
or ""
)
total_latency = (
time.perf_counter()
- started
)
return {
"answer": answer,
"tool_called": False,
"tool_call_count": 0,
"tool_calls": [],
"rag_source_urls": [],
"first_model_latency_seconds": (
round(
first_latency,
6,
)
),
"tool_latency_seconds": 0.0,
"final_model_latency_seconds": 0.0,
"total_latency_seconds": (
round(
total_latency,
6,
)
),
"usage": (
usage_total
),
"response_model": (
first_response.get(
"model"
)
),
}
messages.append(
{
"role": "assistant",
"content": (
first_message.get(
"content"
)
or ""
),
"tool_calls": (
tool_calls
),
}
)
tool_latency_total = 0.0
for tool_call in tool_calls:
if not isinstance(
tool_call,
dict,
):
raise RuntimeError(
"Neplatný tool_call objekt."
)
function = tool_call.get(
"function"
)
if not isinstance(
function,
dict,
):
raise RuntimeError(
"tool_call nemá function."
)
name = function.get(
"name"
)
if name != operation_id:
raise RuntimeError(
"Model zavolal "
"neočakávaný tool: "
f"{name!r}"
)
arguments = (
parse_tool_arguments(
function.get(
"arguments"
)
)
)
tool_started = (
time.perf_counter()
)
rag_result = request_json(
LOCAL_RAG_URL,
method="POST",
headers={
"X-API-Key": (
search_api_key
),
"Content-Type": (
"application/json"
),
},
payload=arguments,
timeout=timeout,
)
tool_latency = (
time.perf_counter()
- tool_started
)
tool_latency_total += (
tool_latency
)
current_urls = (
extract_urls_from_object(
rag_result
)
)
for url in current_urls:
if (
url
not in rag_source_urls
):
rag_source_urls.append(
url
)
tool_calls_record.append(
{
"name": name,
"arguments": arguments,
"latency_seconds": (
round(
tool_latency,
6,
)
),
"source_urls": (
current_urls
),
}
)
tool_call_id = (
tool_call.get(
"id"
)
)
if not tool_call_id:
raise RuntimeError(
"tool_call nemá id."
)
messages.append(
{
"role": "tool",
"tool_call_id": (
tool_call_id
),
"name": name,
"content": json.dumps(
rag_result,
ensure_ascii=False,
),
}
)
final_started = (
time.perf_counter()
)
final_response = request_json(
OPENWEBUI_URL,
method="POST",
headers={
"Authorization": (
"Bearer "
f"{openwebui_api_key}"
),
"Content-Type": (
"application/json"
),
},
payload={
"model": model,
"messages": messages,
"stream": False,
},
timeout=timeout,
)
final_latency = (
time.perf_counter()
- final_started
)
add_usage(
usage_total,
get_usage(
final_response
),
)
final_message = (
get_first_message(
final_response
)
)
answer = str(
final_message.get(
"content"
)
or ""
)
total_latency = (
time.perf_counter()
- started
)
return {
"answer": answer,
"tool_called": True,
"tool_call_count": len(
tool_calls_record
),
"tool_calls": (
tool_calls_record
),
"rag_source_urls": (
rag_source_urls
),
"first_model_latency_seconds": (
round(
first_latency,
6,
)
),
"tool_latency_seconds": (
round(
tool_latency_total,
6,
)
),
"final_model_latency_seconds": (
round(
final_latency,
6,
)
),
"total_latency_seconds": (
round(
total_latency,
6,
)
),
"usage": (
usage_total
),
"response_model": (
final_response.get(
"model"
)
),
}