105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
"""Shared machine-readable fields for RAG and GraphRAG tool responses.
|
|
|
|
Completeness describes retrieval, not correctness or answerability. RAG is a
|
|
ranked selection and cannot promise exhaustive coverage of the wiki.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from typing import Any, Literal, Sequence
|
|
|
|
from neo4j.exceptions import Neo4jError, ServiceUnavailable, SessionExpired
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
RetrievalStatus = Literal["ok", "no_answer", "ambiguous", "error"]
|
|
|
|
|
|
class RetrievalSource(BaseModel):
|
|
model_config = ConfigDict(extra="allow")
|
|
|
|
source_id: str
|
|
document_path: str
|
|
source_url: str | None = None
|
|
title: str | None = None
|
|
|
|
|
|
class RetrievalResponse(BaseModel):
|
|
status: RetrievalStatus = "ok"
|
|
reason: str | None = None
|
|
sources: list[RetrievalSource] = Field(default_factory=list)
|
|
index_version: dict[str, Any] | None = None
|
|
complete: bool | None = Field(
|
|
default=None,
|
|
description=(
|
|
"Úplnosť vyhľadania v rozsahu nástroja: true=úplné, false=neúplné, "
|
|
"null=nemožno zaručiť. Nevyjadruje pravdivosť odpovede."
|
|
),
|
|
)
|
|
truncated: bool = False
|
|
truncation_reasons: list[str] = Field(default_factory=list)
|
|
|
|
|
|
class RetrievalErrorResponse(RetrievalResponse):
|
|
status: Literal["error"] = "error"
|
|
detail: Any = None
|
|
|
|
|
|
def source_references(rows: list[dict]) -> list[dict]:
|
|
"""Stable citation shape; never invent a URL or a document title."""
|
|
return [
|
|
{
|
|
"source_id": f"S{number}",
|
|
"document_path": row["document_path"],
|
|
"source_url": row.get("source_url"),
|
|
"title": row.get("title"),
|
|
}
|
|
for number, row in enumerate(rows, 1)
|
|
]
|
|
|
|
|
|
def result_fields(
|
|
*,
|
|
status: RetrievalStatus,
|
|
sources: list[dict],
|
|
reason: str | None = None,
|
|
complete: bool | None = None,
|
|
truncation_reasons: Sequence[str] = (),
|
|
index_version: dict | None = None,
|
|
) -> dict:
|
|
reasons = list(dict.fromkeys(truncation_reasons))
|
|
return {
|
|
"status": status,
|
|
"reason": reason,
|
|
"sources": sources,
|
|
"index_version": index_version,
|
|
"complete": False if reasons or status in {"error", "ambiguous"} else complete,
|
|
"truncated": bool(reasons),
|
|
"truncation_reasons": reasons,
|
|
}
|
|
|
|
|
|
def error_reason(error: Exception) -> str:
|
|
if isinstance(error, ValueError):
|
|
return "invalid_query"
|
|
if isinstance(
|
|
error,
|
|
(
|
|
FileNotFoundError,
|
|
sqlite3.Error,
|
|
RuntimeError,
|
|
Neo4jError,
|
|
ServiceUnavailable,
|
|
SessionExpired,
|
|
),
|
|
):
|
|
return "retrieval_unavailable"
|
|
return "internal_error"
|
|
|
|
|
|
def error_result(error: Exception | None = None, *, reason: str | None = None) -> dict:
|
|
# Raw exception messages can contain credentials, SQL or private paths.
|
|
return result_fields(
|
|
status="error", sources=[], complete=False, reason=reason or error_reason(error)
|
|
)
|