| from __future__ import annotations |
|
|
| import json |
| import uuid |
| from dataclasses import dataclass, field |
| from datetime import UTC, datetime |
| from enum import StrEnum |
| from typing import Any |
|
|
|
|
| class Mode(StrEnum): |
| LIST = "list" |
| SUMMARIZE = "summarize" |
| GENERATE = "generate" |
|
|
|
|
| |
| |
| |
| DEFAULT_MODE = Mode.SUMMARIZE |
|
|
|
|
| @dataclass(slots=True) |
| class SearchRequest: |
| query: str |
| site: str | None = None |
| mode: Mode = Mode.LIST |
| previous_queries: tuple[str, ...] = () |
| max_results: int = 10 |
| min_score: int = 70 |
| canonical_query: str | None = None |
| retrieval: str = "compare" |
| ranking_model: str | None = None |
|
|
| @property |
| def effective_query(self) -> str: |
| return self.canonical_query or self.query |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Document: |
| id: str |
| url: str |
| name: str |
| site: str |
| text: str |
| schema_object: dict[str, Any] |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Candidate: |
| document: Document |
| vector_score: float | None |
| bm25_score: float | None |
| retrieval_score: float |
| retrieval_source: str |
|
|
| @property |
| def similarity(self) -> float: |
| return self.vector_score or 0.0 |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class Result: |
| document: Document |
| score: int |
| description: str |
| vector_score: float | None = None |
| bm25_score: float | None = None |
| retrieval_source: str | None = None |
| relevance: str = "relevant" |
|
|
| def wire(self) -> dict[str, Any]: |
| return { |
| "@type": "Item", |
| "url": self.document.url, |
| "name": self.document.name, |
| "site": self.document.site, |
| "score": self.score, |
| "relevance": self.relevance, |
| "description": self.description, |
| "vector_score": self.vector_score, |
| "bm25_score": self.bm25_score, |
| "retrieval_source": self.retrieval_source, |
| "schema_object": self.document.schema_object, |
| } |
|
|
|
|
| @dataclass(slots=True) |
| class Event: |
| type: str |
| data: Any = None |
| query_id: str = "" |
| sequence: int = 0 |
| timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) |
| event_id: str = field(default_factory=lambda: str(uuid.uuid4())) |
|
|
| def wire(self) -> dict[str, Any]: |
| return { |
| "message_id": self.event_id, |
| "message_type": self.type, |
| "query_id": self.query_id, |
| "sequence": self.sequence, |
| "timestamp": self.timestamp, |
| "content": self.data, |
| } |
|
|
| def json(self) -> str: |
| return json.dumps(self.wire(), separators=(",", ":")) |
|
|
