""" File: /requests.py Created Date: Monday October 20th 2025 Author: Christian Nonis ----- Last Modified: Sunday April 12th 2026 1:35:36 pm Modified By: Christian Nonis ----- """ from typing import Any, Dict, List, Literal, Optional, Tuple, Union from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator from src.constants.data import Observation, PartialNode, PartialPredicate, TextChunk from src.constants.kg import ( EntitySynergy, IdentificationParams, Node, Predicate, Relationship, ) from src.constants.tasks.ingestion import IngestionTaskArgs from src.core.search.entity_info import MatchPath class IngestionRequestBody(IngestionTaskArgs): """ Request body for the ingestion endpoint. """ class IngestionStructuredDataElement(BaseModel): """ Element for the structured ingestion endpoint. """ json_data: dict = Field( default={}, description="The data rapresenting the structured element.", ) metadata: Optional[dict] = Field( default={}, description="The metadata of the structured element. The information here will be appended to the entity but not analyzed.", ) types: List[str] = Field( default=[], description="A list of types, used to categorize the data.", ) identification_params: Optional[IdentificationParams] = Field( default=None, description="The parameters used to identify the structured element.", ) textual_data: Optional[dict] = Field( default={}, description="The textual descriptive data rapresenting the structured element. Could be a description, a summary, some notes, etc.", ) class RequestPartialNode(PartialNode): uuid: Optional[str] = Field(None, description="The id of the node.") type: str = Field(description="The type of the node.") labels: Optional[List[str]] = Field( default=None, description="Optional labels; defaults to [type] when omitted.", ) description: Optional[str] = Field( default=None, description="The description of the node.", ) properties: Optional[dict] = Field( default_factory=dict, description="The properties of the node.", ) happened_at: Optional[str] = Field( None, description="The date and time the node happened at if known otherwise None. Mostly used for event nodes.", ) @model_validator(mode="after") def default_labels_from_type(self): if not self.labels: self.labels = [self.type] if self.properties is None: self.properties = {} return self class RequestPartialPredicate(PartialPredicate): uuid: Optional[str] = Field(None, description="The id of the relationship.") description: Optional[str] = Field( default=None, description="The description of the relationship.", ) amount: Optional[Union[int, float]] = Field( None, description="The amount of the relationship, if it is a quantitative relationship.", ) class IngestionTripleSet(BaseModel): subject: Optional[RequestPartialNode] = Field( None, description="The optional actor of the triple.", ) subj_event: Optional[RequestPartialPredicate] = Field( None, description="The optional event predicate of the triple.", ) event: Optional[RequestPartialNode] = Field( None, description="The optional event node. Omit for a direct subject-predicate-object edge.", ) event_obj: Optional[RequestPartialPredicate] = Field( None, description="The optional event-to-object predicate. Paired with event when present.", ) object: RequestPartialNode @model_validator(mode="after") def require_event_pair_or_direct_edge(self): has_event = self.event is not None has_event_obj = self.event_obj is not None if has_event != has_event_obj: raise ValueError("event and event_obj must both be set or both omitted") if has_event: return self if not self.subject or not self.subj_event: raise ValueError( "direct triples require subject, subj_event, and object" ) return self class PartialNodeFilter(BaseModel): """ Partial node filter model. """ name: Optional[str] = Field( None, description="The name of the node to filter by.", ) type: Optional[str] = Field( None, description="The type of the node to filter by.", ) uuid: Optional[str] = Field( None, description="The uuid of the node to filter by.", ) meta_description: Optional[str] = Field( None, description="An extra description that explains what this node is about, helpful for entity resolution.", ) @model_validator(mode="after") def require_uuid_or_name_and_type(self): if self.uuid: return self if self.name and self.type: return self raise ValueError( "Anchor must include either uuid, or both name and type." ) class IngestionStructuredRequestBody(BaseModel): """ Request body for the structured ingestion endpoint. """ data: List[IngestionTripleSet] = Field( ..., description="The list of event-centric information triples to ingest.", ) anchor: Optional[PartialNodeFilter] = Field( None, description="The related information to connect the structured data to.", ) text: Optional[str] = Field( default=None, description="Additional text context for the structured data.", ) mode: Optional[Literal["deterministic", "hybrid", "enrich"]] = Field( default=None, description=( "Ingest mode. deterministic: no LLM (triples only; UUID anchor or " "exact name+type match). hybrid: persist submitted triples then " "optionally enrich from text. enrich: same as hybrid with LLM " "backfill. When omitted, inferred as hybrid if text is set else " "deterministic." ), ) brain_id: str = Field( default="default", description="The brain identifier to store the data in." ) @field_validator("anchor", mode="before") @classmethod def reject_string_anchor(cls, value): if isinstance(value, str): raise ValueError( "Anchor must be an object with uuid or name+type, not a string." ) return value def resolved_mode(self) -> Literal["deterministic", "hybrid", "enrich"]: if self.mode is not None: return self.mode return "hybrid" if self.text else "deterministic" class RetrieveRequestResponse(BaseModel): """ Response for the retrieve endpoint. """ data: List[TextChunk] observations: List[Observation] relationships: List[dict] @field_serializer("relationships", when_used="json") def _serialize_relationships(self, value: List[Any]): try: from neo4j.graph import ( Node as NeoNode, ) from neo4j.graph import ( Path as NeoPath, ) from neo4j.graph import ( Relationship as NeoRel, ) except Exception: NeoNode = NeoRel = NeoPath = tuple() from src.constants.kg import Node as KGNode def _ser(obj): if obj is None: return None if isinstance(obj, (str, int, float, bool)): return obj if isinstance(obj, KGNode): return obj.model_dump(mode="json") if isinstance(obj, dict): return {k: _ser(v) for k, v in obj.items()} if isinstance(obj, (list, tuple, set)): return [_ser(v) for v in obj] if NeoNode and isinstance(obj, NeoNode): d = dict(obj) labels = getattr(obj, "labels", None) if labels is not None: d["_labels"] = list(labels) return d if NeoRel and isinstance(obj, NeoRel): d = dict(obj) rel_type = getattr(obj, "type", None) if rel_type is not None: d["_type"] = rel_type try: nodes_attr = getattr(obj, "nodes", None) if nodes_attr: d["_start_id"] = getattr(nodes_attr[0], "id", None) d["_end_id"] = getattr(nodes_attr[1], "id", None) except Exception: pass return d if NeoPath and isinstance(obj, NeoPath): return { "nodes": [_ser(n) for n in obj.nodes], "relationships": [_ser(r) for r in obj.relationships], } return str(obj) return [_ser(v) for v in value] if value is not None else [] class RetrievedNeighborNode(BaseModel): """ Node model for the retrieve neighbors endpoint. """ neighbor: Node relationship: Predicate most_common: Node similarity_score: float = Field( 0.0, description="The similarity score of the neighbor to the main node." ) class RetrieveNeighborsRequestResponse(BaseModel): """ Response for the retrieve neighbors endpoint. """ count: int = Field(0, description="The number of neighbors found.") main_node: Node = Field(..., description="The main node of the neighbors.") neighbors: List[RetrievedNeighborNode] class RetrieveNeighborsAiModeRequestBody(BaseModel): """ Request body for the retrieve neighbors AI mode endpoint. """ identification_params: IdentificationParams = Field( ..., description="The identification parameters of the entity to get neighbors for.", ) limit: int = Field(10, description="The number of neighbors to return.") looking_for: Optional[list[str]] = Field( ..., description="The description of the neighbors to look for.", ) brain_id: str = Field( default="default", description="The brain identifier to store the data in." ) class RetrieveNeighborsWithIdentificationParamsRequestBody(BaseModel): """ Request body for the retrieve neighbors with identification params endpoint. """ identification_params: IdentificationParams = Field( ..., description="The identification parameters of the entity to get neighbors for.", ) limit: int = Field(10, description="The number of neighbors to return.") brain_id: str = Field( default="default", description="The brain identifier to store the data in." ) look_for: Optional[str] = Field( None, description="Optional filter for what type of neighbors to look for." ) class CreateBrainRequest(BaseModel): """ Request body for the create brain endpoint. """ brain_id: str class AddEntityRequest(BaseModel): """Request model for adding a new entity to the graph.""" name: str brain_id: str = "default" labels: list[str] = [] description: Optional[str] = None properties: Optional[dict] = None identification_params: Optional[dict] = None metadata: Optional[dict] = None class UpdateEntityRequest(BaseModel): """Request model for updating an existing entity in the graph.""" uuid: str brain_id: str = "default" new_name: Optional[str] = None new_description: Optional[str] = None new_labels: Optional[list[str]] = None new_properties: Optional[dict] = None properties_to_remove: Optional[list[str]] = None class AddRelationshipRequest(BaseModel): """Request model for adding a new relationship between two entities.""" subject_uuid: str predicate_name: str predicate_description: str object_uuid: str brain_id: str = "default" class UpdateRelationshipRequest(BaseModel): """Request model for updating an existing relationship's properties.""" uuid: str brain_id: str = "default" new_properties: Optional[dict] = None properties_to_remove: Optional[list[str]] = None class GetEntityInfoResponse(BaseModel): """Response model for the get entity info endpoint.""" target_node: Optional[Node] = None path: MatchPath class GetEntityContextResponse(BaseModel): """Response model for the get entity context endpoint.""" neighborhood: list[dict] target_node: Optional[Node] = None text_contexts: list[str] = [] natural_language_web: list[dict] = [] class GetEntitySibilingsResponse(BaseModel): """Response model for the get entity siblings endpoint.""" target_node: Node synergies: List[EntitySynergy] anchors: Optional[List[Node]] = None potential_anchors: Optional[List[Node]] = None class RecommendRequestBody(BaseModel): """Request body for event-graph recommendations.""" target: str brain_id: str = "default" polarity: Literal["same", "opposite"] = "same" top_k: int = Field(20, ge=1, le=200) labels: Optional[List[str]] = None include_asymmetric: bool = True include_multi_interest: bool = True include_attribute_pref: bool = False diversify: bool = True asymmetric_direction: Literal["outbound", "inbound", "both"] = "outbound" exclude_seen: bool = False recency_half_life_days: Optional[float] = Field(None, ge=0) dampen_degree: bool = False behavior_weights: Optional[Dict[str, float]] = None class RecommendItem(BaseModel): node: Node score: float connected_by: List[Node] channel: str class RecommendResponse(BaseModel): target_node: Node recommendations: List[RecommendItem] class GetEntityStatusResponse(BaseModel): """Response model for the get entity status endpoint.""" node: Optional[Node] = None exists: bool has_relationships: bool relationships: List[Tuple[Predicate, Node]] observations: List[Observation] class GetContextRequestBody(BaseModel): """Request body for the get context endpoint.""" text: str brain_id: str = "default" historical_limit: int = 10 max_facts: int = Field(40, ge=0) max_passages: int = 8 apply_fact_filter: bool = True use_ppr: bool = True sufficiency_retry: bool = False profile_stages: bool = False cross_event_bridges: int = Field(3, ge=0) class GetContextTriple(BaseModel): """Triple for the get context endpoint.""" identified_entity: str triple: Tuple[Node, Predicate, Node, Predicate, Node] source_chunk_ids: Optional[List[str]] = None source_session_ids: Optional[List[str]] = None class GetContextResponse(BaseModel): """Response for the get context endpoint.""" text_context: str triples: List[GetContextTriple] historical_context: List[str] = [] source_passages: List[str] = [] graph_session_ids: Optional[List[str]] = None temporal_conflicts: Optional[List[dict[str, Any]]] = None paths: Optional[List[dict[str, Any]]] = None topics: Optional[List[dict[str, Any]]] = None stage_timings: Optional[dict[str, Any]] = None class SearchRequestBody(BaseModel): query: str brain_id: str = "default" k: int = Field(10, ge=1, le=200) channels: List[str] = Field(default_factory=lambda: ["passages"]) node_labels: Optional[List[str]] = None community_labels: Optional[List[str]] = None expand: Literal["none", "neighbors"] = "none" fusion: Optional[Literal["rrf", "cc"]] = None fusion_alpha: Optional[float] = Field(default=None, ge=0.0, le=1.0) rerank: Optional[str] = Field( default=None, description="none or plugin:. Unknown plugin names return 400.", ) mode: Literal["default", "catalog"] = "default" profile_stages: bool = False extras: Optional[Dict[str, str]] = None target: Optional[str] = None @field_validator("target", mode="before") @classmethod def empty_target(cls, value): if value is None: return None text = str(value).strip() return text or None @field_validator("extras", mode="before") @classmethod def coerce_extras_filter(cls, value): if value is None or value == {}: return None if not isinstance(value, dict): raise ValueError("extras must be an object of string keys and values") return { str(key): str(item) for key, item in value.items() if item is not None } or None class SearchHitScores(BaseModel): bm25: Optional[float] = None dense: Optional[float] = None rrf: Optional[float] = None cc: Optional[float] = None rerank: Optional[float] = None plugin: Optional[dict[str, float]] = None graph: Optional[float] = None personalize: Optional[float] = None class SearchHit(BaseModel): id: str channel: str score: float scores: SearchHitScores snippet: str labels: List[str] = Field(default_factory=list) extras: Optional[dict[str, Any]] = None node_id: Optional[str] = None class SearchResponse(BaseModel): hits: List[SearchHit] stage_timings: Optional[dict[str, Any]] = None channel_lists: Optional[dict[str, list[str]]] = None facets: Optional[dict[str, dict[str, int]]] = None node_ids: List[str] = Field(default_factory=list)