fixing pydantic issue v6

main/api.py

@@ -1,6 +1,7 @@
 import httpx
-from typing import Optional, AsyncIterator, Dict, Any
+from typing import Optional, AsyncIterator, Dict, Any, Iterator
 import logging
+import asyncio
 from litserve import LitAPI
 from pydantic import BaseModel
 
@@ -24,10 +25,25 @@ class InferenceApi(LitAPI):
         )
         self.logger.info(f"Inference API setup completed on device: {device}")
 
-    async def predict(self, x: str, **kwargs) -> AsyncIterator[str]:
+    def predict(self, x: str, **kwargs) -> Iterator[str]:
         """
-        Main prediction method required by LitAPI.
-        Always yields, either chunks in streaming mode or complete response in non-streaming mode.
+        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):
@@ -42,12 +58,17 @@ class InferenceApi(LitAPI):
             return request["prompt"]
         return request
 
-    def encode_response(self, output: AsyncIterator[str], **kwargs) -> AsyncIterator[Dict[str, str]]:
+    def encode_response(self, output: Iterator[str], **kwargs) -> Dict[str, Any]:
         """Convert the model output to a response payload."""
-        async def wrapper():
-            async for chunk in output:
-                yield {"generated_text": chunk}
-        return wrapper()
+        # For streaming responses
+        if self.stream:
+            return {"generated_text": output}
+        # For non-streaming, take the first (and only) item from the iterator
+        try:
+            result = next(output)
+            return {"generated_text": result}
+        except StopIteration:
+            return {"generated_text": ""}
 
     async def generate_response(
             self,