from pydantic import BaseModel
from typing import List, Optional

class RegisterRequest(BaseModel):
    url: str

class PendingRequest(BaseModel):
    task_id: int
    limit: Optional[int] = 10

class AddUrlsRequest(BaseModel):
    task_id: int
    urls_obj: dict

# schemas.py
class CrawlResult(BaseModel):
    source_url: str
    chunk_index: int  # 新增字段
    title: Optional[str] = None
    content: Optional[str] = None
    embedding: Optional[List[float]] = None

class SaveResultsRequest(BaseModel):
    task_id: int
    results: List[CrawlResult]

class SearchRequest(BaseModel):
    # 如果不传 task_id，则进行全库搜索
    task_id: Optional[int] = None
    query_embedding: dict
    limit: Optional[int] = 5