Spaces:
Runtime error
Runtime error
File size: 13,302 Bytes
eb5a3fb 47031d7 eb5a3fb 47031d7 63fdbaa a4e24d4 db5664e eb5a3fb db5664e d0b5a4b db5664e 47031d7 a4e24d4 02fd6bb 47031d7 a4e24d4 47031d7 da1009f 02fd6bb a4e24d4 799409f db5664e da1009f 799409f da1009f 2a73022 9a5d1e4 da1009f 9a5d1e4 02fd6bb 47031d7 02fd6bb 73ca5b8 1f1d641 ab616bd 73ca5b8 ab616bd 73ca5b8 94c7e5b 73ca5b8 ab616bd 73ca5b8 63fdbaa 02fd6bb 63fdbaa 02fd6bb 8f2f662 a4e24d4 73ca5b8 eb5a3fb 73ca5b8 ab616bd 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b 73ca5b8 d0b5a4b a4e24d4 02fd6bb a4e24d4 63fdbaa 02fd6bb 63fdbaa 47031d7 a4e24d4 da1009f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 |
import json
from pathlib import Path
import httpx
from typing import Optional, AsyncIterator, Dict, Any, Iterator, List, Callable
import logging
import asyncio
from litserve import LitAPI
from pydantic import BaseModel
from .utils import extract_json
class GenerationResponse(BaseModel):
generated_text: str
class InferenceApi(LitAPI):
def __init__(self, config: Dict[str, Any]):
"""Initialize the Inference API with configuration."""
super().__init__()
self.logger = logging.getLogger(__name__)
self.logger.info("Initializing Inference API")
self._device = None
self.stream = False
self.config = config
self.llm_config = config.get('llm_server', {})
def setup(self, device: Optional[str] = None):
"""Synchronous setup method required by LitAPI"""
self._device = device
self.logger.info(f"Inference API setup completed on device: {device}")
return self # It's common for setup methods to return self for chaining
async def _get_client(self):
"""Get or create HTTP client as needed"""
host = self.llm_config.get('host', 'localhost')
port = self.llm_config.get('port', 8002)
# Construct base URL, omitting port for HF spaces
if 'hf.space' in host:
base_url = f"https://{host}"
else:
base_url = f"http://{host}:{port}"
return httpx.AsyncClient(
base_url=base_url,
timeout=float(self.llm_config.get('timeout', 60.0))
)
def _get_endpoint(self, endpoint_name: str) -> str:
"""Get full endpoint path including prefix"""
endpoints = self.llm_config.get('endpoints', {})
api_prefix = self.llm_config.get('api_prefix', '')
endpoint = endpoints.get(endpoint_name, '')
return f"{api_prefix}{endpoint}"
async def _make_request(
self,
method: str,
endpoint: str,
*,
params: Optional[Dict[str, Any]] = None,
json: Optional[Dict[str, Any]] = None,
stream: bool = False
) -> Any:
"""Make an authenticated request to the LLM Server."""
base_url = self.llm_config.get('host', 'http://localhost:8001')
full_endpoint = f"{base_url.rstrip('/')}/{self._get_endpoint(endpoint).lstrip('/')}"
try:
self.logger.info(f"Making {method} request to: {full_endpoint}")
# Create client outside the with block for streaming
client = await self._get_client()
if stream:
# For streaming, return both client and response context managers
return client, client.stream(
method,
self._get_endpoint(endpoint),
params=params,
json=json
)
else:
# For non-streaming, use context manager
async with client as c:
response = await c.request(
method,
self._get_endpoint(endpoint),
params=params,
json=json
)
response.raise_for_status()
return response
except Exception as e:
self.logger.error(f"Error in request to {full_endpoint}: {str(e)}")
raise
def predict(self, x: str, **kwargs) -> Iterator[str]:
"""Non-async prediction method that yields results."""
loop = asyncio.get_event_loop()
async def async_gen():
async for item in self._async_predict(x, **kwargs):
yield item
gen = async_gen()
while True:
try:
yield loop.run_until_complete(gen.__anext__())
except StopAsyncIteration:
break
async def _async_predict(self, x: str, **kwargs) -> AsyncIterator[str]:
"""Internal async prediction method."""
if self.stream:
async for chunk in self.generate_stream(x, **kwargs):
yield chunk
else:
response = await self.generate_response(x, **kwargs)
yield response
async def generate_response(
self,
prompt: str,
system_message: Optional[str] = None,
max_new_tokens: Optional[int] = None
) -> str:
"""Generate a complete response by forwarding the request to the LLM Server."""
self.logger.debug(f"Forwarding generation request for prompt: {prompt[:50]}...")
try:
response = await self._make_request(
"POST",
"generate",
json={
"prompt": prompt,
"system_message": system_message,
"max_new_tokens": max_new_tokens
}
)
data = response.json()
return data["generated_text"]
except Exception as e:
self.logger.error(f"Error in generate_response: {str(e)}")
raise
async def structured_llm_query(
self,
template_name: str,
input_text: str,
additional_context: Optional[Dict[str, Any]] = None,
pre_hooks: Optional[List[Callable]] = None,
post_hooks: Optional[List[Callable]] = None
) -> Dict[str, Any]:
"""Execute a structured LLM query using a template."""
template_path = Path(__file__).parent / "prompt_templates" / f"{template_name}.json"
try:
# Load and parse template
with open(template_path) as f:
template = json.load(f)
# Apply pre-processing hooks
processed_input = input_text
if pre_hooks:
for hook in pre_hooks:
processed_input = hook(processed_input)
# Format the prompt with the context
context = {"input_text": processed_input}
if additional_context:
context.update(additional_context)
prompt = template["prompt_template"].format(**context)
# Make the request to the LLM
response = await self._make_request(
"POST",
"generate",
json={
"prompt": prompt,
"system_message": template.get("system_message"),
"max_new_tokens": 1000
}
)
# Extract JSON from response
data = response.json()
result = extract_json(data["generated_text"])
# Apply any additional post-processing hooks
if post_hooks:
for hook in post_hooks:
result = hook(result)
return result
except FileNotFoundError:
raise ValueError(f"Template {template_name} not found")
except Exception as e:
self.logger.error(f"Error in structured_llm_query: {str(e)}")
raise
async def expand_query(
self,
query: str,
system_message: Optional[str] = None
) -> Dict[str, Any]:
"""Expand a query for RAG processing."""
return await self.structured_llm_query(
template_name="query_expansion",
input_text=query,
additional_context={"system_message": system_message} if system_message else None
)
async def rerank_chunks(
self,
query: str,
chunks: List[str],
system_message: Optional[str] = None
) -> Dict[str, Any]:
"""Rerank text chunks based on their relevance to the query."""
# Format chunks as numbered list for better LLM processing
formatted_chunks = "\n".join(f"{i+1}. {chunk}" for i, chunk in enumerate(chunks))
return await self.structured_llm_query(
template_name="chunk_rerank",
input_text=query,
additional_context={
"chunks": formatted_chunks,
"system_message": system_message
}
)
async def generate_stream(
self,
prompt: str,
system_message: Optional[str] = None,
max_new_tokens: Optional[int] = None
) -> AsyncIterator[str]:
"""Generate a streaming response by forwarding the request to the LLM Server."""
self.logger.debug(f"Forwarding streaming request for prompt: {prompt[:50]}...")
try:
client, stream_cm = await self._make_request(
"POST",
"generate_stream",
json={
"prompt": prompt,
"system_message": system_message,
"max_new_tokens": max_new_tokens
},
stream=True
)
async with client:
async with stream_cm as response:
async for chunk in response.aiter_text():
yield chunk
except Exception as e:
self.logger.error(f"Error in generate_stream: {str(e)}")
raise
async def generate_embedding(self, text: str) -> List[float]:
"""Generate embedding vector from input text."""
self.logger.debug(f"Forwarding embedding request for text: {text[:50]}...")
try:
response = await self._make_request(
"POST",
"embedding",
json={"text": text}
)
data = response.json()
return data["embedding"]
except Exception as e:
self.logger.error(f"Error in generate_embedding: {str(e)}")
raise
async def check_system_status(self) -> Dict[str, Any]:
"""Check system status of the LLM Server."""
self.logger.debug("Checking system status...")
try:
response = await self._make_request(
"GET",
"system_status"
)
return response.json()
except Exception as e:
self.logger.error(f"Error in check_system_status: {str(e)}")
raise
async def download_model(self, model_name: Optional[str] = None) -> Dict[str, str]:
"""Download model files from the LLM Server."""
self.logger.debug(f"Forwarding model download request for: {model_name or 'default model'}")
try:
response = await self._make_request(
"POST",
"model_download",
params={"model_name": model_name} if model_name else None
)
return response.json()
except Exception as e:
self.logger.error(f"Error in download_model: {str(e)}")
raise
async def validate_system(self) -> Dict[str, Any]:
"""Validate system configuration and setup."""
self.logger.debug("Validating system configuration...")
try:
response = await self._make_request(
"GET",
"system_validate"
)
return response.json()
except Exception as e:
self.logger.error(f"Error in validate_system: {str(e)}")
raise
async def initialize_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
"""Initialize specified model or default model."""
self.logger.debug(f"Initializing model: {model_name or 'default'}")
try:
response = await self._make_request(
"POST",
"model_initialize",
params={"model_name": model_name} if model_name else None
)
return response.json()
except Exception as e:
self.logger.error(f"Error in initialize_model: {str(e)}")
raise
async def initialize_embedding_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
"""Initialize embedding model."""
self.logger.debug(f"Initializing embedding model: {model_name or 'default'}")
try:
response = await self._make_request(
"POST",
"model_initialize_embedding",
json={"model_name": model_name} if model_name else {}
)
return response.json()
except Exception as e:
self.logger.error(f"Error in initialize_embedding_model: {str(e)}")
raise
def decode_request(self, request: Any, **kwargs) -> str:
"""Convert the request payload to input format."""
if isinstance(request, dict) and "prompt" in request:
return request["prompt"]
return request
def encode_response(self, output: Iterator[str], **kwargs) -> Dict[str, Any]:
"""Convert the model output to a response payload."""
if self.stream:
return {"generated_text": output}
try:
result = next(output)
return {"generated_text": result}
except StopIteration:
return {"generated_text": ""}
async def cleanup(self):
"""Cleanup method - no longer needed as clients are created per-request"""
pass |