From e40c0af2f7491d219e1fd3268fd3b7f5db6e89c8 Mon Sep 17 00:00:00 2001
From: KKlochko <kklochko@protonmail.com>
Date: Fri, 14 Nov 2025 14:22:51 +0200
Subject: [PATCH] Add models.

---
 gliner_inference_server/models.py | 119 ++++++++++++++++++++++++++++++
 1 file changed, 119 insertions(+)
 create mode 100644 gliner_inference_server/models.py

diff --git a/gliner_inference_server/models.py b/gliner_inference_server/models.py
new file mode 100644
index 0000000..9bfe56b
--- /dev/null
+++ b/gliner_inference_server/models.py
@@ -0,0 +1,119 @@
+from pydantic import BaseModel, Field, field_validator
+from typing import List, Optional, Union, Dict, Any, Tuple
+import os
+
+
+# ==================== Request Models ====================
+class GeneralRequest(BaseModel):
+    text: Optional[str] = None
+    texts: Optional[List[str]] = None
+    entities: List[str]
+    prompt: Optional[str] = None
+    threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
+
+    @field_validator('entities')
+    @classmethod
+    def validate_entities(cls, v):
+        if not v or len(v) == 0:
+            raise ValueError("entities list cannot be empty")
+        return v
+
+    @field_validator('text', 'texts')
+    @classmethod
+    def validate_text_length(cls, v):
+        max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
+        if isinstance(v, str) and len(v) > max_length:
+            raise ValueError(f"text exceeds maximum length of {max_length}")
+        if isinstance(v, list):
+            for text in v:
+                if len(text) > max_length:
+                    raise ValueError(f"text exceeds maximum length of {max_length}")
+        return v
+
+
+class RelationRule(BaseModel):
+    relation: str
+    pairs_filter: List[Tuple[str, str]]
+    distance: int
+
+
+class RelationExtractionRequest(BaseModel):
+    text: Optional[str] = None
+    texts: Optional[List[str]] = None
+    relations: List[str]
+    entities: Optional[List[str]] = None
+    rules: Optional[List[RelationRule]] = None
+    prompt: Optional[str] = None
+    threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
+
+    @field_validator('relations')
+    @classmethod
+    def validate_relations(cls, v):
+        if not v or len(v) == 0:
+            raise ValueError("relations list cannot be empty")
+        return v
+
+    @field_validator('text', 'texts')
+    @classmethod
+    def validate_text_length(cls, v):
+        max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
+        if isinstance(v, str) and len(v) > max_length:
+            raise ValueError(f"text exceeds maximum length of {max_length}")
+        if isinstance(v, list):
+            for text in v:
+                if len(text) > max_length:
+                    raise ValueError(f"text exceeds maximum length of {max_length}")
+        return v
+
+
+class SummarizationRequest(BaseModel):
+    text: Optional[str] = None
+    texts: Optional[List[str]] = None
+    prompt: Optional[str] = None
+    threshold: Optional[float] = Field(default=0.5, ge=0.0, le=1.0)
+
+    @field_validator('text', 'texts')
+    @classmethod
+    def validate_text_length(cls, v):
+        max_length = int(os.getenv("MAX_TEXT_LENGTH", "10000"))
+        if isinstance(v, str) and len(v) > max_length:
+            raise ValueError(f"text exceeds maximum length of {max_length}")
+        if isinstance(v, list):
+            for text in v:
+                if len(text) > max_length:
+                    raise ValueError(f"text exceeds maximum length of {max_length}")
+        return v
+
+
+# ==================== Response Models ====================
+class EntityOutput(BaseModel):
+    entity: str
+    span: str
+    start: int
+    end: int
+    score: float
+
+
+class RelationOutput(BaseModel):
+    source: EntityOutput
+    relation: str
+    target: EntityOutput
+    score: float
+
+
+class SummaryOutput(BaseModel):
+    text: str
+    start: int
+    end: int
+    score: float
+
+# ==================== Error Response Model ====================
+class ErrorDetail(BaseModel):
+    code: str
+    message: str
+    details: Dict[str, Any] = Field(default_factory=dict)
+
+
+class ErrorResponse(BaseModel):
+    error: ErrorDetail
+