Spaces:
Running
Running
| """ | |
| File: /requests.py | |
| Created Date: Monday October 20th 2025 | |
| Author: Christian Nonis <alch.infoemail@gmail.com> | |
| ----- | |
| Last Modified: Sunday April 12th 2026 1:35:36 pm | |
| Modified By: Christian Nonis <alch.infoemail@gmail.com> | |
| ----- | |
| """ | |
| 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.", | |
| ) | |
| 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 | |
| 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.", | |
| ) | |
| 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." | |
| ) | |
| 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] | |
| 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:<name>. Unknown plugin names return 400.", | |
| ) | |
| mode: Literal["default", "catalog"] = "default" | |
| profile_stages: bool = False | |
| extras: Optional[Dict[str, str]] = None | |
| target: Optional[str] = None | |
| def empty_target(cls, value): | |
| if value is None: | |
| return None | |
| text = str(value).strip() | |
| return text or None | |
| 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) | |