TeamGenKI/Inference-API
0
1import json2from pathlib import Path3 4import httpx5from typing import Optional, AsyncIterator, Dict, Any, Iterator, List, Callable6import logging7import asyncio8from litserve import LitAPI9from pydantic import BaseModel10from .utils import extract_json11 12 13class GenerationResponse(BaseModel):14 generated_text: str15 16class InferenceApi(LitAPI):17 def __init__(self, config: Dict[str, Any]):18 """Initialize the Inference API with configuration."""19 super().__init__()20 self.logger = logging.getLogger(__name__)21 self.logger.info("Initializing Inference API")22 self._device = None23 self.stream = False24 self.config = config25 self.llm_config = config.get('llm_server', {})26 27 def setup(self, device: Optional[str] = None):28 """Synchronous setup method required by LitAPI"""29 self._device = device30 self.logger.info(f"Inference API setup completed on device: {device}")31 return self # It's common for setup methods to return self for chaining32 33 async def _get_client(self):34 """Get or create HTTP client as needed"""35 host = self.llm_config.get('host', 'localhost')36 port = self.llm_config.get('port', 8002)37 38 # Construct base URL, omitting port for HF spaces39 if 'hf.space' in host:40 base_url = f"https://{host}"41 else:42 base_url = f"http://{host}:{port}"43 44 return httpx.AsyncClient(45 base_url=base_url,46 timeout=float(self.llm_config.get('timeout', 60.0))47 )48 49 def _get_endpoint(self, endpoint_name: str) -> str:50 """Get full endpoint path including prefix"""51 endpoints = self.llm_config.get('endpoints', {})52 api_prefix = self.llm_config.get('api_prefix', '')53 endpoint = endpoints.get(endpoint_name, '')54 return f"{api_prefix}{endpoint}"55 56 async def _make_request(57 self,58 method: str,59 endpoint: str,60 *,61 params: Optional[Dict[str, Any]] = None,62 json: Optional[Dict[str, Any]] = None,63 stream: bool = False64 ) -> Any:65 """Make an authenticated request to the LLM Server."""66 base_url = self.llm_config.get('host', 'http://localhost:8001')67 full_endpoint = f"{base_url.rstrip('/')}/{self._get_endpoint(endpoint).lstrip('/')}"68 69 try:70 self.logger.info(f"Making {method} request to: {full_endpoint}")71 # Create client outside the with block for streaming72 client = await self._get_client()73 74 if stream:75 # For streaming, return both client and response context managers76 return client, client.stream(77 method,78 self._get_endpoint(endpoint),79 params=params,80 json=json81 )82 else:83 # For non-streaming, use context manager84 async with client as c:85 response = await c.request(86 method,87 self._get_endpoint(endpoint),88 params=params,89 json=json90 )91 response.raise_for_status()92 return response93 94 except Exception as e:95 self.logger.error(f"Error in request to {full_endpoint}: {str(e)}")96 raise97 98 def predict(self, x: str, **kwargs) -> Iterator[str]:99 """Non-async prediction method that yields results."""100 loop = asyncio.get_event_loop()101 async def async_gen():102 async for item in self._async_predict(x, **kwargs):103 yield item104 105 gen = async_gen()106 while True:107 try:108 yield loop.run_until_complete(gen.__anext__())109 except StopAsyncIteration:110 break111 112 async def _async_predict(self, x: str, **kwargs) -> AsyncIterator[str]:113 """Internal async prediction method."""114 if self.stream:115 async for chunk in self.generate_stream(x, **kwargs):116 yield chunk117 else:118 response = await self.generate_response(x, **kwargs)119 yield response120 121 async def generate_response(122 self,123 prompt: str,124 system_message: Optional[str] = None,125 max_new_tokens: Optional[int] = None126 ) -> str:127 """Generate a complete response by forwarding the request to the LLM Server."""128 self.logger.debug(f"Forwarding generation request for prompt: {prompt[:50]}...")129 130 try:131 response = await self._make_request(132 "POST",133 "generate",134 json={135 "prompt": prompt,136 "system_message": system_message,137 "max_new_tokens": max_new_tokens138 }139 )140 data = response.json()141 return data["generated_text"]142 143 except Exception as e:144 self.logger.error(f"Error in generate_response: {str(e)}")145 raise146 147 async def structured_llm_query(148 self,149 template_name: str,150 input_text: str,151 additional_context: Optional[Dict[str, Any]] = None,152 pre_hooks: Optional[List[Callable]] = None,153 post_hooks: Optional[List[Callable]] = None154 ) -> Dict[str, Any]:155 """Execute a structured LLM query using a template."""156 template_path = Path(__file__).parent / "prompt_templates" / f"{template_name}.json"157 158 try:159 # Load and parse template160 with open(template_path) as f:161 template = json.load(f)162 163 # Apply pre-processing hooks164 processed_input = input_text165 if pre_hooks:166 for hook in pre_hooks:167 processed_input = hook(processed_input)168 169 # Format the prompt with the context170 context = {"input_text": processed_input}171 if additional_context:172 context.update(additional_context)173 174 prompt = template["prompt_template"].format(**context)175 176 # Make the request to the LLM177 response = await self._make_request(178 "POST",179 "generate",180 json={181 "prompt": prompt,182 "system_message": template.get("system_message"),183 "max_new_tokens": 1000184 }185 )186 187 # Extract JSON from response188 data = response.json()189 result = extract_json(data["generated_text"])190 191 # Apply any additional post-processing hooks192 if post_hooks:193 for hook in post_hooks:194 result = hook(result)195 196 return result197 198 except FileNotFoundError:199 raise ValueError(f"Template {template_name} not found")200 except Exception as e:201 self.logger.error(f"Error in structured_llm_query: {str(e)}")202 raise203 204 async def expand_query(205 self,206 query: str,207 system_message: Optional[str] = None208 ) -> Dict[str, Any]:209 """Expand a query for RAG processing."""210 return await self.structured_llm_query(211 template_name="query_expansion",212 input_text=query,213 additional_context={"system_message": system_message} if system_message else None214 )215 216 async def rerank_chunks(217 self,218 query: str,219 chunks: List[str],220 system_message: Optional[str] = None221 ) -> Dict[str, Any]:222 """Rerank text chunks based on their relevance to the query."""223 # Format chunks as numbered list for better LLM processing224 formatted_chunks = "\n".join(f"{i+1}. {chunk}" for i, chunk in enumerate(chunks))225 226 return await self.structured_llm_query(227 template_name="chunk_rerank",228 input_text=query,229 additional_context={230 "chunks": formatted_chunks,231 "system_message": system_message232 }233 )234 235 236 async def generate_stream(237 self,238 prompt: str,239 system_message: Optional[str] = None,240 max_new_tokens: Optional[int] = None241 ) -> AsyncIterator[str]:242 """Generate a streaming response by forwarding the request to the LLM Server."""243 self.logger.debug(f"Forwarding streaming request for prompt: {prompt[:50]}...")244 245 try:246 client, stream_cm = await self._make_request(247 "POST",248 "generate_stream",249 json={250 "prompt": prompt,251 "system_message": system_message,252 "max_new_tokens": max_new_tokens253 },254 stream=True255 )256 257 async with client:258 async with stream_cm as response:259 async for chunk in response.aiter_text():260 yield chunk261 262 except Exception as e:263 self.logger.error(f"Error in generate_stream: {str(e)}")264 raise265 266 async def generate_embedding(self, text: str) -> List[float]:267 """Generate embedding vector from input text."""268 self.logger.debug(f"Forwarding embedding request for text: {text[:50]}...")269 270 try:271 response = await self._make_request(272 "POST",273 "embedding",274 json={"text": text}275 )276 data = response.json()277 return data["embedding"]278 279 except Exception as e:280 self.logger.error(f"Error in generate_embedding: {str(e)}")281 raise282 283 async def check_system_status(self) -> Dict[str, Any]:284 """Check system status of the LLM Server."""285 self.logger.debug("Checking system status...")286 287 try:288 response = await self._make_request(289 "GET",290 "system_status"291 )292 return response.json()293 294 except Exception as e:295 self.logger.error(f"Error in check_system_status: {str(e)}")296 raise297 298 async def download_model(self, model_name: Optional[str] = None) -> Dict[str, str]:299 """Download model files from the LLM Server."""300 self.logger.debug(f"Forwarding model download request for: {model_name or 'default model'}")301 302 try:303 response = await self._make_request(304 "POST",305 "model_download",306 params={"model_name": model_name} if model_name else None307 )308 return response.json()309 310 except Exception as e:311 self.logger.error(f"Error in download_model: {str(e)}")312 raise313 314 async def validate_system(self) -> Dict[str, Any]:315 """Validate system configuration and setup."""316 self.logger.debug("Validating system configuration...")317 318 try:319 response = await self._make_request(320 "GET",321 "system_validate"322 )323 return response.json()324 325 except Exception as e:326 self.logger.error(f"Error in validate_system: {str(e)}")327 raise328 329 async def initialize_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:330 """Initialize specified model or default model."""331 self.logger.debug(f"Initializing model: {model_name or 'default'}")332 333 try:334 response = await self._make_request(335 "POST",336 "model_initialize",337 params={"model_name": model_name} if model_name else None338 )339 return response.json()340 341 except Exception as e:342 self.logger.error(f"Error in initialize_model: {str(e)}")343 raise344 345 async def initialize_embedding_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:346 """Initialize embedding model."""347 self.logger.debug(f"Initializing embedding model: {model_name or 'default'}")348 349 try:350 response = await self._make_request(351 "POST",352 "model_initialize_embedding",353 json={"model_name": model_name} if model_name else {}354 )355 return response.json()356 357 except Exception as e:358 self.logger.error(f"Error in initialize_embedding_model: {str(e)}")359 raise360 361 def decode_request(self, request: Any, **kwargs) -> str:362 """Convert the request payload to input format."""363 if isinstance(request, dict) and "prompt" in request:364 return request["prompt"]365 return request366 367 def encode_response(self, output: Iterator[str], **kwargs) -> Dict[str, Any]:368 """Convert the model output to a response payload."""369 if self.stream:370 return {"generated_text": output}371 try:372 result = next(output)373 return {"generated_text": result}374 except StopIteration:375 return {"generated_text": ""}376 377 async def cleanup(self):378 """Cleanup method - no longer needed as clients are created per-request"""379 pass