Spaces:
Sleeping
Sleeping
| import re | |
| import unicodedata | |
| from enum import StrEnum | |
| from typing import Annotated | |
| from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator, model_validator | |
| NonEmptyStr = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] | |
| RESERVED_RELATIONS = frozenset({"MENTIONS", "HAS_CHUNK"}) | |
| FALLBACK_RELATION = "RELATED_TO" | |
| def to_ascii(value: str) -> str: | |
| return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode() | |
| def canonical_key(name: str) -> str: | |
| return re.sub(r"[^a-z0-9]", "", to_ascii(name).lower()) or name.strip().casefold() | |
| class EntityType(StrEnum): | |
| COMPONENT = "Component" | |
| CLASS = "Class" | |
| FUNCTION = "Function" | |
| LIBRARY = "Library" | |
| CONCEPT = "Concept" | |
| FEATURE = "Feature" | |
| ERROR = "Error" | |
| DATABASE = "Database" | |
| API = "API" | |
| ENTITY_TYPES_BY_NAME = {entity_type.value.lower(): entity_type for entity_type in EntityType} | |
| class RetrievalMode(StrEnum): | |
| VECTOR = "vector" | |
| GRAPH = "graphrag" | |
| class Chunk(BaseModel): | |
| model_config = ConfigDict(frozen=True) | |
| id: str | |
| document: str | |
| position: int | |
| text: str | |
| class RetrievedChunk(Chunk): | |
| score: float | |
| entity_keys: list[str] = Field(default_factory=list) | |
| class Entity(BaseModel): | |
| name: NonEmptyStr | |
| type: EntityType = EntityType.CONCEPT | |
| def coerce_type(cls, value: object) -> EntityType: | |
| return ENTITY_TYPES_BY_NAME.get(str(value).strip().lower(), EntityType.CONCEPT) | |
| def key(self) -> str: | |
| return canonical_key(self.name) | |
| class Relation(BaseModel): | |
| source: NonEmptyStr | |
| relation: str | |
| target: NonEmptyStr | |
| evidence: str = "" | |
| def normalize_relation(cls, value: object) -> str: | |
| normalized = re.sub(r"[^A-Z0-9]+", "_", to_ascii(str(value)).upper()).strip("_") | |
| if not normalized or not normalized[0].isalpha() or normalized in RESERVED_RELATIONS: | |
| return FALLBACK_RELATION | |
| return normalized | |
| def normalize_evidence(cls, value: object) -> str: | |
| return " ".join(str(value or "").split()) | |
| def source_key(self) -> str: | |
| return canonical_key(self.source) | |
| def target_key(self) -> str: | |
| return canonical_key(self.target) | |
| class GraphExtraction(BaseModel): | |
| entities: list[Entity] = Field(default_factory=list) | |
| relations: list[Relation] = Field(default_factory=list) | |
| def reconcile(self) -> "GraphExtraction": | |
| self.relations = [ | |
| relation for relation in self.relations if relation.source_key != relation.target_key | |
| ] | |
| entities = {entity.key: entity for entity in self.entities} | |
| for relation in self.relations: | |
| entities.setdefault(relation.source_key, Entity(name=relation.source)) | |
| entities.setdefault(relation.target_key, Entity(name=relation.target)) | |
| self.entities = list(entities.values()) | |
| return self | |
| class GraphFact(BaseModel): | |
| source: str | |
| relation: str | |
| target: str | |
| evidence: list[str] = Field(default_factory=list) | |
| documents: list[str] = Field(default_factory=list) | |
| distance: int = 1 | |
| class RetrievalContext(BaseModel): | |
| chunks: list[RetrievedChunk] = Field(default_factory=list) | |
| facts: list[GraphFact] = Field(default_factory=list) | |
| seed_entities: list[str] = Field(default_factory=list) | |
| def is_empty(self) -> bool: | |
| return not self.chunks and not self.facts | |
| def sources(self) -> list[str]: | |
| documents = {chunk.document for chunk in self.chunks} | |
| documents.update(document for fact in self.facts for document in fact.documents) | |
| return sorted(documents) | |
| class Answer(BaseModel): | |
| question: str | |
| text: str | |
| mode: RetrievalMode | |
| context: RetrievalContext | |