30c7bda389
- Removed multiple failed document entries from `documents.json`. - Added a new document entry with updated metadata and changed the index name to `regulations_dense_1024_v2`. - Updated architecture documentation to reflect changes in the Milvus collection name. - Adjusted requirements by removing the sqlalchemy dependency. - Modified test cases to align with new document structure and naming conventions. - Introduced a new test file for Milvus vector index runtime recovery and error handling. - Updated assertions in various test files to ensure compatibility with the new schema.
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
"""Define domain models for retrieval."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
# Keep module behavior explicit so the backend flow stays easy to audit.
|
|
|
|
|
|
|
|
@dataclass
|
|
class RetrievalQuery:
|
|
"""Represent the Retrieval Query type."""
|
|
query: str
|
|
top_k: int
|
|
filters: str | None = None
|
|
|
|
|
|
@dataclass(init=False)
|
|
class RetrievedChunk:
|
|
"""Represent the retrieved chunk payload with legacy aliases."""
|
|
|
|
chunk_id: str
|
|
doc_id: str
|
|
doc_title: str
|
|
text: str
|
|
score: float
|
|
chunk_type: str = ""
|
|
section_title: str = ""
|
|
page_start: int = 0
|
|
page_end: int = 0
|
|
section_level: int = 0
|
|
chunk_index: int = 0
|
|
piece_index: int = 0
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
chunk_id: str,
|
|
doc_id: str,
|
|
doc_title: str | None = None,
|
|
text: str | None = None,
|
|
score: float = 0.0,
|
|
chunk_type: str = "",
|
|
section_title: str = "",
|
|
page_start: int = 0,
|
|
page_end: int = 0,
|
|
section_level: int = 0,
|
|
chunk_index: int = 0,
|
|
piece_index: int = 0,
|
|
metadata: dict[str, Any] | None = None,
|
|
doc_name: str | None = None,
|
|
content: str | None = None,
|
|
page_number: int | None = None,
|
|
block_type: str | None = None,
|
|
**_: Any,
|
|
) -> None:
|
|
"""Initialize the retrieved chunk while accepting legacy field names."""
|
|
self.chunk_id = chunk_id
|
|
self.doc_id = doc_id
|
|
self.doc_title = doc_title if doc_title is not None else (doc_name or "")
|
|
self.text = text if text is not None else (content or "")
|
|
self.score = float(score)
|
|
self.chunk_type = chunk_type or (block_type or "")
|
|
self.section_title = section_title
|
|
self.page_start = int(page_start or page_number or 0)
|
|
self.page_end = int(page_end or self.page_start)
|
|
self.section_level = int(section_level or 0)
|
|
self.chunk_index = int(chunk_index or 0)
|
|
self.piece_index = int(piece_index or 0)
|
|
self.metadata = dict(metadata or {})
|
|
|
|
@property
|
|
def doc_name(self) -> str:
|
|
"""Return the legacy document name alias."""
|
|
return self.doc_title
|
|
|
|
@doc_name.setter
|
|
def doc_name(self, value: str) -> None:
|
|
"""Update the legacy document name alias."""
|
|
self.doc_title = value
|
|
|
|
@property
|
|
def content(self) -> str:
|
|
"""Return the legacy content alias."""
|
|
return self.text
|
|
|
|
@content.setter
|
|
def content(self, value: str) -> None:
|
|
"""Update the legacy content alias."""
|
|
self.text = value
|
|
|
|
@property
|
|
def page_number(self) -> int:
|
|
"""Return the legacy page number alias."""
|
|
return self.page_start
|
|
|
|
@page_number.setter
|
|
def page_number(self, value: int) -> None:
|
|
"""Update the legacy page number alias."""
|
|
self.page_start = value
|
|
self.page_end = max(self.page_end, value)
|