77 lines
2.8 KiB
Python
77 lines
2.8 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))
|