added openAI schema based endpoint and response v3

main/main.py

@@ -33,10 +33,12 @@ async def async_main():
         config = load_config()
         server_config = config.get('server', {})
 
-        # Initialize API with config and await setup
+        # Initialize API with config
         api = InferenceApi(config)
+        # Setup API first
         await api.setup()
-        await init_router(config)
+        # Initialize router with the already setup API instance
+        await init_router(api)
 
         # Create LitServer instance with config
         server = ls.LitServer(
@@ -61,7 +63,7 @@ async def async_main():
 
         # Get configured port
         port = server_config.get('port', 8001)
-        host = server_config.get('host', 'localhost')
+        host = server_config.get('host', '0.0.0.0')
 
         # Run server
         server.run(host=host, port=port)

main/routes.py

@@ -1,5 +1,9 @@
 from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
 from typing import Optional
+import json
+from time import time
+import logging
 from .api import InferenceApi
 from .schemas import (
     GenerateRequest,
@@ -10,12 +14,49 @@ from .schemas import (
     ChatCompletionRequest,
     ChatCompletionResponse
 )
-import logging
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
 api = None
 
+async def init_router(inference_api: InferenceApi):
+    """Initialize router with an already setup API instance"""
+    global api
+    api = inference_api
+    logger.info("Router initialized with Inference API instance")
+
+@router.post("/generate")
+async def generate_text(request: GenerateRequest):
+    """Generate text response from prompt"""
+    logger.info(f"Received generation request for prompt: {request.prompt[:50]}...")
+    try:
+        response = await api.generate_response(
+            prompt=request.prompt,
+            system_message=request.system_message,
+            max_new_tokens=request.max_new_tokens
+        )
+        logger.info("Successfully generated response")
+        return {"generated_text": response}
+    except Exception as e:
+        logger.error(f"Error in generate_text endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/generate/stream")
+async def generate_stream(request: GenerateRequest):
+    """Generate streaming text response from prompt"""
+    logger.info(f"Received streaming generation request for prompt: {request.prompt[:50]}...")
+    try:
+        return StreamingResponse(
+            api.generate_stream(
+                prompt=request.prompt,
+                system_message=request.system_message,
+                max_new_tokens=request.max_new_tokens
+            ),
+            media_type="text/event-stream"
+        )
+    except Exception as e:
+        logger.error(f"Error in generate_stream endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 @router.post("/v1/chat/completions")
 async def create_chat_completion(request: ChatCompletionRequest):
@@ -72,44 +113,6 @@ async def create_chat_completion(request: ChatCompletionRequest):
         logger.error(f"Error in chat completion endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
-async def init_router(config: dict):
-    """Initialize router with config and Inference API instance"""
-    global api
-    api = InferenceApi(config)
-    await api.setup()
-    logger.info("Router initialized with Inference API instance")
-
-@router.post("/generate")
-async def generate_text(request: GenerateRequest):
-    """Generate text response from prompt"""
-    logger.info(f"Received generation request for prompt: {request.prompt[:50]}...")
-    try:
-        response = await api.generate_response(
-            prompt=request.prompt,
-            system_message=request.system_message,
-            max_new_tokens=request.max_new_tokens
-        )
-        logger.info("Successfully generated response")
-        return {"generated_text": response}
-    except Exception as e:
-        logger.error(f"Error in generate_text endpoint: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@router.post("/generate/stream")
-async def generate_stream(request: GenerateRequest):
-    """Generate streaming text response from prompt"""
-    logger.info(f"Received streaming generation request for prompt: {request.prompt[:50]}...")
-    try:
-        return api.generate_stream(
-            prompt=request.prompt,
-            system_message=request.system_message,
-            max_new_tokens=request.max_new_tokens
-        )
-    except Exception as e:
-        logger.error(f"Error in generate_stream endpoint: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
-
 @router.post("/embedding", response_model=EmbeddingResponse)
 async def generate_embedding(request: EmbeddingRequest):
     """Generate embedding vector from text"""
@@ -175,4 +178,4 @@ async def initialize_embedding_model(model_name: Optional[str] = None):
 async def shutdown_event():
     """Clean up resources on shutdown"""
     if api:
-        await api.close()
+        await api.cleanup()