added openAI schema based endpoint and response

main/api.py

@@ -9,13 +9,15 @@ class GenerationResponse(BaseModel):
     generated_text: str
 
 class InferenceApi(LitAPI):
-    def __init__(self):
+    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  # Add stream flag for compatibility with LitAPI
+        self.stream = False
+        self.config = config
+        self.llm_config = config.get('llm_server', {})
 
     async def setup(self, device: Optional[str] = None):
         """Setup method required by LitAPI"""
@@ -25,15 +27,19 @@ class InferenceApi(LitAPI):
     async def _get_client(self):
         """Get or create HTTP client as needed"""
         return httpx.AsyncClient(
-            base_url="http://localhost:8002",
-            timeout=60.0
+            base_url=self.llm_config.get('base_url', 'http://localhost:8002'),
+            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}"
+
     def predict(self, x: str, **kwargs) -> Iterator[str]:
-        """
-        Non-async prediction method that yields results.
-        Implements required LitAPI method.
-        """
+        """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):
@@ -47,9 +53,7 @@ class InferenceApi(LitAPI):
                 break
 
     async def _async_predict(self, x: str, **kwargs) -> AsyncIterator[str]:
-        """
-        Internal async prediction method.
-        """
+        """Internal async prediction method."""
         if self.stream:
             async for chunk in self.generate_stream(x, **kwargs):
                 yield chunk
@@ -58,19 +62,13 @@ class InferenceApi(LitAPI):
             yield response
 
     def decode_request(self, request: Any, **kwargs) -> str:
-        """
-        Convert the request payload to input format.
-        Implements required LitAPI method.
-        """
+        """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.
-        Implements required LitAPI method.
-        """
+        """Convert the model output to a response payload."""
         if self.stream:
             return {"generated_text": output}
         try:
@@ -91,7 +89,7 @@ class InferenceApi(LitAPI):
         try:
             async with await self._get_client() as client:
                 response = await client.post(
-                    "/api/v1/generate",
+                    self._get_endpoint('generate'),
                     json={
                         "prompt": prompt,
                         "system_message": system_message,
@@ -119,7 +117,7 @@ class InferenceApi(LitAPI):
             client = await self._get_client()
             async with client.stream(
                     "POST",
-                    "/api/v1/generate/stream",
+                    self._get_endpoint('generate_stream'),
                     json={
                         "prompt": prompt,
                         "system_message": system_message,

main/config.yaml

@@ -1,11 +1,18 @@
 server:
+  host: "localhost"
   port: 8001
   timeout: 60
+  max_batch_size: 1
 
 llm_server:
-  base_url: "https://teamgenki-llmserver.hf.space:7680"  # URL of your LLM Server
-  timeout: 60  # Timeout for requests to LLM Server
-
-logging:
-  level: "INFO"
-  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+  base_url: "https://teamgenki-llmserver.hf.space:7860"
+  timeout: 60.0
+  api_prefix: "/api/v1"  # This will be used for route prefixing
+  endpoints:
+    generate: "/generate"
+    generate_stream: "/generate/stream"
+    embedding: "/embedding"
+    system_status: "/system/status"
+    system_validate: "/system/validate"
+    model_initialize: "/model/initialize"
+    model_initialize_embedding: "/model/initialize/embedding"

main/main.py

@@ -31,17 +31,18 @@ async def async_main():
     try:
         # Load configuration
         config = load_config()
+        server_config = config.get('server', {})
 
-        # Initialize API and router with await
-        api = InferenceApi()
-        await api.setup()  # Properly await the setup
-        await init_router(config)  # Modified to be async
+        # Initialize API with config and await setup
+        api = InferenceApi(config)
+        await api.setup()
+        await init_router(config)
 
-        # Create LitServer instance
+        # Create LitServer instance with config
         server = ls.LitServer(
             api,
-            timeout=config.get("server", {}).get("timeout", 60),
-            max_batch_size=1,
+            timeout=server_config.get('timeout', 60),
+            max_batch_size=server_config.get('max_batch_size', 1),
             track_requests=True
         )
 
@@ -54,13 +55,16 @@ async def async_main():
             allow_headers=["*"],
         )
 
-        # Add our routes to the server's FastAPI app
-        server.app.include_router(router, prefix="/api/v1")
+        # Add routes with configured prefix
+        api_prefix = config.get('llm_server', {}).get('api_prefix', '/api/v1')
+        server.app.include_router(router, prefix=api_prefix)
 
-        port = config.get("server", {}).get("port", 8001)
+        # Get configured port
+        port = server_config.get('port', 8001)
+        host = server_config.get('host', 'localhost')
 
-        # Get the server instance from our app and run it
-        server.run(port=port)
+        # Run server
+        server.run(host=host, port=port)
 
     except Exception as e:
         logger.error(f"Server initialization failed: {str(e)}")

main/routes.py

@@ -6,7 +6,9 @@ from .schemas import (
     EmbeddingRequest,
     EmbeddingResponse,
     SystemStatusResponse,
-    ValidationResponse
+    ValidationResponse,
+    ChatCompletionRequest,
+    ChatCompletionResponse
 )
 import logging
 
@@ -14,11 +16,68 @@ router = APIRouter()
 logger = logging.getLogger(__name__)
 api = None
 
+
+@router.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest):
+    """OpenAI-compatible chat completion endpoint"""
+    logger.info(f"Received chat completion request with {len(request.messages)} messages")
+
+    try:
+        # Extract the last user message, or combine messages if needed
+        last_message = request.messages[-1].content
+
+        if request.stream:
+            # For streaming, we need to create a generator that yields OpenAI-compatible chunks
+            async def generate_stream():
+                async for chunk in api.generate_stream(
+                        prompt=last_message,
+                ):
+                    # Create a streaming response chunk in OpenAI format
+                    response_chunk = {
+                        "id": "chatcmpl-123",
+                        "object": "chat.completion.chunk",
+                        "created": int(time()),
+                        "model": request.model,
+                        "choices": [{
+                            "index": 0,
+                            "delta": {
+                                "content": chunk
+                            },
+                            "finish_reason": None
+                        }]
+                    }
+                    yield f"data: {json.dumps(response_chunk)}\n\n"
+
+                # Send the final chunk
+                yield f"data: [DONE]\n\n"
+
+            return StreamingResponse(
+                generate_stream(),
+                media_type="text/event-stream"
+            )
+
+        else:
+            # For non-streaming, generate the full response
+            response_text = await api.generate_response(
+                prompt=last_message,
+            )
+
+            # Convert to OpenAI format
+            return ChatCompletionResponse.from_response(
+                content=response_text,
+                model=request.model
+            )
+
+    except Exception as e:
+        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()
-    await api.setup()  # Properly await the setup
+    api = InferenceApi(config)
+    await api.setup()
     logger.info("Router initialized with Inference API instance")
 
 @router.post("/generate")

main/schemas.py

@@ -1,5 +1,57 @@
-from pydantic import BaseModel
-from typing import Optional, List, Dict, Union
+from pydantic import BaseModel, Field
+from typing import List, Optional, Dict, Union, Literal
+from time import time
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    stream: bool = False
+
+class ChatCompletionMessage(BaseModel):
+    role: str = "assistant"
+    content: str
+
+class ChatCompletionChoice(BaseModel):
+    index: int = 0
+    message: ChatCompletionMessage
+    logprobs: Optional[None] = None
+    finish_reason: str = "stop"
+
+class CompletionTokenDetails(BaseModel):
+    reasoning_tokens: int = 0
+    accepted_prediction_tokens: int = 0
+    rejected_prediction_tokens: int = 0
+
+class CompletionUsage(BaseModel):
+    prompt_tokens: int = 9  # Placeholder values
+    completion_tokens: int = 12
+    total_tokens: int = 21
+    completion_tokens_details: CompletionTokenDetails = Field(default_factory=CompletionTokenDetails)
+
+class ChatCompletionResponse(BaseModel):
+    id: str = Field(default="chatcmpl-123")
+    object: str = "chat.completion"
+    created: int = Field(default_factory=lambda: int(time()))
+    model: str = "gpt-4o-mini"
+    system_fingerprint: str = "fp_44709d6fcb"
+    choices: List[ChatCompletionChoice]
+    usage: CompletionUsage = Field(default_factory=CompletionUsage)
+
+    @classmethod
+    def from_response(cls, content: str, model: str = "gpt-4o-mini") -> "ChatCompletionResponse":
+        """Create a ChatCompletionResponse from a simple response string"""
+        return cls(
+            model=model,
+            choices=[
+                ChatCompletionChoice(
+                    message=ChatCompletionMessage(content=content)
+                )
+            ]
+        )
 
 class GenerateRequest(BaseModel):
     prompt: str

requirements.txt

@@ -9,6 +9,7 @@ httptools==0.6.4
 httpx==0.28.1
 idna==3.10
 litserve==0.2.5
+numpy==2.2.1
 pydantic==2.10.4
 pydantic_core==2.27.2
 python-dotenv==1.0.1