apply pydantic2 and add SearchRequest, DocumentUpdate

SlendGamer · Sep 29, 2024 · ddab46f · ddab46f
1 parent c780582
commit ddab46f
Showing 1 changed file with 58 additions and 11 deletions.
diff --git a/backend/app/schemas.py b/backend/app/schemas.py
@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import List, Optional
 
-from pydantic import BaseModel, EmailStr
+from pydantic import BaseModel, ConfigDict, EmailStr
 
 
 class TagBase(BaseModel):
@@ -15,8 +15,7 @@ class TagCreate(TagBase):
 class TagResponse(TagBase):
     id: int
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class UserBase(BaseModel):
@@ -32,8 +31,7 @@ class UserResponse(UserBase):
     id: int
     created_at: datetime
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class DocumentBase(BaseModel):
@@ -53,8 +51,15 @@ class DocumentResponse(DocumentBase):
     summary: Optional[str] = None
     tags: List[TagResponse] = []
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
+
+
+class DocumentUpdate(BaseModel):
+    title: Optional[str] = None
+    description: Optional[str] = None
+    tags: Optional[List[str]] = []
+
+    model_config = ConfigDict(from_attributes=True)
 
 
 class DocumentVersionResponse(BaseModel):
@@ -64,14 +69,56 @@ class DocumentVersionResponse(BaseModel):
     file_path: str
     uploaded_at: datetime
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class SearchIndexResponse(BaseModel):
     id: int
     document_id: int
     embedding_vector: str  # Adjust type as needed
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
+
+
+class SearchRequest(BaseModel):
+    query: str
+    top_k: Optional[int] = 5
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class SearchResult(BaseModel):
+    document_id: int
+    title: str
+    description: Optional[str]
+    summary: Optional[str]
+    relevance_score: float
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class SummaryRequest(BaseModel):
+    document_id: int
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class SummaryResponse(BaseModel):
+    document_id: int
+    summary: Optional[str]
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+
+class TokenData(BaseModel):
+    username: Optional[str] = None
+
+
+class UserLogin(BaseModel):
+    username: str
+    password: str