graph_rag_server/src/api/db/models/search_request_models.py

"""
搜索相关的Pydantic模型定义

该模块定义了用于图像和文本搜索的请求和响应模型。
"""
from pydantic import BaseModel, Field, field_validator
from typing import List, Optional


class SearchRequest(BaseModel):
    """
    搜索请求模型
    
    支持通过图像URL或文本进行相似性搜索。
    至少需要提供 image_url 或 matching_text 其中之一。
    
    Attributes:
        image_url: 用于图像搜索的图片URL (可选)
        matching_text: 用于文本匹配搜索的查询文本 (可选)
        topn: 返回的最大结果数量，默认为10，范围1-100
    """
    database_ids: List[str] = Field(
        default=[],
        description="知识库id",
        examples=["database_id1", "database_id2"]
    )
    database_ids_str: Optional[str] = Field(
        default=[],
        description="知识库id",
        examples=["database_id1", "database_id2"]
    )
    image_url: Optional[str] = Field(
        default=None,
        description="用于图像搜索的图片URL",
        examples=["https://example.com/image.jpg"]
    )
    matching_text: Optional[str] = Field(
        default=None,
        description="用于文本匹配搜索的查询文本",
        examples=["搜索关键词"]
    )
    topk: int = Field(
        default=10,
        ge=1,
        le=100,
        description="返回的最大结果数量"
    )
    score_threshold: float = Field(
        default=0.0,
        ge=0.0,
        le=1.0,
        description="相似度分数阈值"
    )
    kb_ids_str: Optional[str] = Field(
        default=None,
        description="知识库id",
    )
    kb_id: List[str] = Field(
        default=None,
        description="知识库id",
        examples=["kb_id1", "kb_id2"]
    )

    @field_validator('image_url')
    @classmethod
    def validate_image_url(cls, v: Optional[str]) -> Optional[str]:
        """验证图片URL格式"""
        if v is not None and v.strip() == "":
            return None
        return v
    
    @field_validator('matching_text')
    @classmethod
    def validate_matching_text(cls, v: Optional[str]) -> Optional[str]:
        """验证匹配文本"""
        if v is not None and v.strip() == "":
            return None
        return v
    
    def model_post_init(self, __context) -> None:
        """模型初始化后验证：确保至少提供一个搜索条件"""
        if self.image_url is None and self.matching_text is None:
            raise ValueError("必须提供 image_url 或 matching_text 其中至少一个")