Files
privateGPT/private_gpt/components/text_processing/models.py
Javier Martinez ce25f73a48 fix: citations
2026-07-16 16:04:18 +02:00

60 lines
1.2 KiB
Python

from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any
class Action(Enum):
PASS = auto()
DROP = auto()
UNWRAP = auto()
REPLACE = auto()
HOLD = auto()
class ProbeStatus(Enum):
NO_MATCH = auto()
MATCH = auto()
NEED_MORE = auto()
@dataclass(frozen=True)
class ProbeResult:
status: ProbeStatus
consumed: int = 0
action: Action = Action.PASS
replacement: str | None = None
metadata: tuple[Any, ...] = ()
state_updates: dict[str, Any] = field(default_factory=dict)
state_deletes: tuple[str, ...] = ()
@classmethod
def no_match(cls) -> ProbeResult:
return cls(status=ProbeStatus.NO_MATCH)
@classmethod
def need_more(cls) -> ProbeResult:
return cls(status=ProbeStatus.NEED_MORE, action=Action.HOLD)
@dataclass
class ProcessingContext:
final: bool = False
state: dict[str, Any] = field(default_factory=dict)
@dataclass(frozen=True)
class ProcessResult:
text: str
metadata: tuple[Any, ...]
pending: str
consumed: int
@dataclass(frozen=True)
class ProcessDelta:
text: str
metadata: tuple[Any, ...]
pending: str