mirror of
https://github.com/james-m-jordan/morphik-core.git
synced 2025-05-09 19:32:38 +00:00
38 lines
1011 B
Python
38 lines
1011 B
Python
from typing import Any, Dict, List, Optional, Type, TypeVar, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
# Type variable for any Pydantic model
|
|
PydanticT = TypeVar("PydanticT", bound=BaseModel)
|
|
|
|
|
|
class ChunkSource(BaseModel):
|
|
"""Source information for a chunk used in completion"""
|
|
|
|
document_id: str
|
|
chunk_number: int
|
|
score: Optional[float] = None
|
|
|
|
|
|
class CompletionResponse(BaseModel):
|
|
"""Response from completion generation"""
|
|
|
|
completion: Union[str, PydanticT]
|
|
usage: Dict[str, int]
|
|
finish_reason: Optional[str] = None
|
|
sources: List[ChunkSource] = []
|
|
metadata: Optional[Dict] = None
|
|
|
|
|
|
class CompletionRequest(BaseModel):
|
|
"""Request for completion generation"""
|
|
|
|
query: str
|
|
context_chunks: List[str]
|
|
max_tokens: Optional[int] = 1000
|
|
temperature: Optional[float] = 0.7
|
|
prompt_template: Optional[str] = None
|
|
folder_name: Optional[str] = None
|
|
end_user_id: Optional[str] = None
|
|
schema: Optional[Union[Type[BaseModel], Dict[str, Any]]] = None
|