Spaces:
Sleeping
Sleeping
| from fastapi import Query, Depends, HTTPException | |
| from pydantic import BaseModel, model_validator | |
| from typing import Dict, Any | |
| import json, os, logging | |
| from app.utils import ( | |
| is_safe_path, | |
| _prompt_base_path, | |
| _prompt_path, | |
| ) | |
| # Configure logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| class Prompt(BaseModel): | |
| name: str = '' | |
| prompt: str = '' | |
| params: Dict[str, Any] = {} | |
| source: str = '' | |
| def prompt_validate(self) -> "Prompt": | |
| # If name is provided, validate it and prefer it over prompt | |
| if self.name: | |
| if is_safe_path(_prompt_base_path(), _prompt_path(self.name)): | |
| if os.path.isfile(_prompt_path(self.name)): | |
| self.source = 'name' | |
| return self | |
| else: | |
| logger.error(f"The specified prompt file at '{_prompt_path(self.name)}' was not found.") | |
| raise HTTPException(status_code=400, detail=f"The specified prompt file at '{_prompt_path(self.name)}' was not found.") | |
| else: | |
| logger.error("Unsafe path access detected; Got path: " + _prompt_path(self.name)) | |
| raise HTTPException(status_code=400, detail="Unsafe path access detected; Got path: " + _prompt_path(self.name)) | |
| # If prompt is provided (and name is not valid), use prompt | |
| if self.prompt: | |
| self.source = 'prompt' | |
| return self | |
| # If neither is valid, raise error | |
| logger.error("At least one parameter is required: 'prompt_name' or 'prompt_prompt'") | |
| raise HTTPException(status_code=400, detail="At least one parameter is required: 'prompt_name' or 'prompt_prompt'") | |
| class RequestModel(BaseModel): | |
| prompt: Prompt | |
| # AI_model: AIModel # AI_model.name, AI_model.params | |
| class QueryParams(BaseModel): | |
| prompt_name: str = Query('') | |
| prompt_prompt: str = Query('') | |
| prompt_params: str = Query('{}') | |
| def parse_query(query: QueryParams = Depends()) -> RequestModel: | |
| # Transform prompt_params (JSON string -> dict) | |
| if query.prompt_params: | |
| try: | |
| prompt_params = json.loads(query.prompt_params) | |
| except Exception as e: | |
| logger.exception(f"Invalid JSON for prompt_params: {e}") | |
| raise HTTPException(status_code=400, detail=f"Invalid JSON for prompt_params: {e}") | |
| # Make Prompt model | |
| prompt = Prompt( | |
| name = query.prompt_name, | |
| prompt=query.prompt_prompt, | |
| params=prompt_params | |
| ) | |
| logger.info(f"RequestModel has been created and verified") | |
| return RequestModel(prompt=prompt) |