construction-rag / app /models.py
Ashanasri's picture
Upload app/models.py with huggingface_hub
3c951fb verified
from pydantic import BaseModel, Field
from typing import List, Optional, Annotated
class EmbedQueryIn(BaseModel):
text: Annotated[str, Field(..., min_length=1)]
class EmbedQueryOut(BaseModel):
model: str
dim: int
embedding: List[float]
class SearchIn(BaseModel):
query: Annotated[str, Field(min_length=1)]
top_k: int = Field(5, ge=1, le=20)
rerank: bool = True
include_chunks: bool = False # New field to control chunk inclusion in response
class ChunkHit(BaseModel):
id: int
source: Optional[str] = None
page: Optional[int] = None
chunk_index: Optional[int] = None
score: float
rerank_score: Optional[float] = None
text: str
class SearchOut(BaseModel):
query: str
answer: str
hits: Optional[List[ChunkHit]] = None
# from pydantic import BaseModel, Field
# from typing import List, Optional, Annotated
# class EmbedQueryIn(BaseModel):
# text: Annotated[str, Field(..., min_length=1, description="User question or query text")]
# class EmbedQueryOut(BaseModel):
# model: str
# dim: int
# embedding: List[float]
# class SearchIn(BaseModel):
# query: Annotated[str, Field(min_length=1)]
# top_k: int = Field(5, ge=1, le=20)
# rerank: bool = True
# class ChunkHit(BaseModel):
# id: int
# source: str
# page: int
# chunk_index: int
# score: float
# rerank_score: Optional[float] = None
# text: str
# class SearchOut(BaseModel):
# query: str
# top_k: int
# rerank: bool
# hits: List[ChunkHit]
# class SearchAnswerOut(BaseModel):
# query: str
# answer: str