Spaces:
Sleeping
Sleeping
File size: 4,053 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | 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
@field_validator("type", mode="before")
@classmethod
def coerce_type(cls, value: object) -> EntityType:
return ENTITY_TYPES_BY_NAME.get(str(value).strip().lower(), EntityType.CONCEPT)
@property
def key(self) -> str:
return canonical_key(self.name)
class Relation(BaseModel):
source: NonEmptyStr
relation: str
target: NonEmptyStr
evidence: str = ""
@field_validator("relation", mode="before")
@classmethod
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
@field_validator("evidence", mode="before")
@classmethod
def normalize_evidence(cls, value: object) -> str:
return " ".join(str(value or "").split())
@property
def source_key(self) -> str:
return canonical_key(self.source)
@property
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)
@model_validator(mode="after")
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)
@property
def is_empty(self) -> bool:
return not self.chunks and not self.facts
@property
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
|