changed to uvicorn setup for HF v11

main/api.py

@@ -2,6 +2,7 @@ import httpx
 from typing import Optional, AsyncIterator, Dict, Any, Iterator, List
 import logging
 import asyncio
+import os
 from litserve import LitAPI
 from pydantic import BaseModel
 
@@ -40,6 +41,52 @@ class InferenceApi(LitAPI):
         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.
+
+        Args:
+            method: HTTP method ('GET' or 'POST')
+            endpoint: Endpoint name to get from config
+            params: Query parameters
+            json: JSON body for POST requests
+            stream: Whether to return a streaming response
+        """
+        access_token = os.environ.get("InferenceAPI")
+        headers = {"Authorization": f"Bearer {access_token}"} if access_token else {}
+
+        try:
+            async with await self._get_client() as client:
+                if stream:
+                    return await client.stream(
+                        method,
+                        self._get_endpoint(endpoint),
+                        params=params,
+                        json=json,
+                        headers=headers
+                    )
+                else:
+                    response = await client.request(
+                        method,
+                        self._get_endpoint(endpoint),
+                        params=params,
+                        json=json,
+                        headers=headers
+                    )
+                    response.raise_for_status()
+                    return response
+
+        except Exception as e:
+            self.logger.error(f"Error in request to {endpoint}: {str(e)}")
+            raise
+
     def predict(self, x: str, **kwargs) -> Iterator[str]:
         """Non-async prediction method that yields results."""
         loop = asyncio.get_event_loop()
@@ -63,19 +110,71 @@ class InferenceApi(LitAPI):
             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 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:
+            async with await self._make_request(
+                    "POST",
+                    "generate_stream",
+                    json={
+                        "prompt": prompt,
+                        "system_message": system_message,
+                        "max_new_tokens": max_new_tokens
+                    },
+                    stream=True
+            ) 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:
-            async with await self._get_client() as client:
-                response = await client.post(
-                    self._get_endpoint('embedding'),
-                    json={"text": text}
-                )
-                response.raise_for_status()
-                data = response.json()
-                return data["embedding"]
+            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)}")
@@ -86,12 +185,11 @@ class InferenceApi(LitAPI):
         self.logger.debug("Checking system status...")
 
         try:
-            async with await self._get_client() as client:
-                response = await client.get(
-                    self._get_endpoint('system_status')
-                )
-                response.raise_for_status()
-                return response.json()
+            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)}")
@@ -102,33 +200,27 @@ class InferenceApi(LitAPI):
         self.logger.debug(f"Forwarding model download request for: {model_name or 'default model'}")
 
         try:
-            async with await self._get_client() as client:
-                response = await client.post(
-                    self._get_endpoint('model_download'),
-                    params={"model_name": model_name} if model_name else None
-                )
-                response.raise_for_status()
-                return response.json()
+            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
 
-        except Exception as e:
-            self.logger.error(f"Error initiating model download: {str(e)}")
-            raise
-
     async def validate_system(self) -> Dict[str, Any]:
         """Validate system configuration and setup."""
         self.logger.debug("Validating system configuration...")
 
         try:
-            async with await self._get_client() as client:
-                response = await client.get(
-                    self._get_endpoint('system_validate')
-                )
-                response.raise_for_status()
-                return response.json()
+            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)}")
@@ -139,13 +231,12 @@ class InferenceApi(LitAPI):
         self.logger.debug(f"Initializing model: {model_name or 'default'}")
 
         try:
-            async with await self._get_client() as client:
-                response = await client.post(
-                    self._get_endpoint('model_initialize'),
-                    params={"model_name": model_name} if model_name else None
-                )
-                response.raise_for_status()
-                return response.json()
+            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)}")
@@ -156,13 +247,12 @@ class InferenceApi(LitAPI):
         self.logger.debug(f"Initializing embedding model: {model_name or 'default'}")
 
         try:
-            async with await self._get_client() as client:
-                response = await client.post(
-                    self._get_endpoint('model_initialize_embedding'),
-                    json={"model_name": model_name} if model_name else {}
-                )
-                response.raise_for_status()
-                return response.json()
+            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)}")
@@ -184,62 +274,6 @@ class InferenceApi(LitAPI):
         except StopIteration:
             return {"generated_text": ""}
 
-    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:
-            async with await self._get_client() as client:
-                response = await client.post(
-                    self._get_endpoint('generate'),
-                    json={
-                        "prompt": prompt,
-                        "system_message": system_message,
-                        "max_new_tokens": max_new_tokens
-                    }
-                )
-                response.raise_for_status()
-                data = response.json()
-                return data["generated_text"]
-
-        except Exception as e:
-            self.logger.error(f"Error in generate_response: {str(e)}")
-            raise
-
-    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 = await self._get_client()
-            async with client.stream(
-                    "POST",
-                    self._get_endpoint('generate_stream'),
-                    json={
-                        "prompt": prompt,
-                        "system_message": system_message,
-                        "max_new_tokens": max_new_tokens
-                    }
-            ) as response:
-                response.raise_for_status()
-                async for chunk in response.aiter_text():
-                    yield chunk
-            await client.aclose()
-
-        except Exception as e:
-            self.logger.error(f"Error in generate_stream: {str(e)}")
-            raise
-
     async def cleanup(self):
         """Cleanup method - no longer needed as clients are created per-request"""
         pass

main/config.yaml

@@ -5,7 +5,7 @@ server:
   max_batch_size: 1
 
 llm_server:
-  base_url: "https://teamgenki-llmserver.hf.space:7860"
+  base_url: "https://teamgenki-llmserver.hf.space:7860" # The base URL of the LLM server
   timeout: 60.0
   api_prefix: "/api/v1"  # This will be used for route prefixing
   endpoints: