Merge pull request #1 from muxi-ai/fix/event-loop-management

fix: resolve event loop management issues for async/sync interoperability

Author: ranaroussi

onellm/audio.py

@@ -26,11 +26,11 @@
 with support for fallback models if the primary model fails.
 """
 
-import asyncio
 from typing import Any, Dict, IO, List, Optional, Union
 
 from .providers.base import get_provider_with_fallbacks
 from .utils.fallback import FallbackConfig
+from .utils.async_helpers import run_async
 
 class AudioTranscription:
     """
@@ -113,8 +113,8 @@ def create_sync(
         Returns:
             Transcription result
         """
-        # Use asyncio.run to execute the async create method in a new event loop
-        return asyncio.run(
+        # Use our safe async runner to execute the async create method
+        return run_async(
             cls.create(
                 file=file,
                 model=model,
@@ -204,8 +204,8 @@ def create_sync(
         Returns:
             Translation result with text in English
         """
-        # Use asyncio.run to execute the async create method in a new event loop
-        return asyncio.run(
+        # Use our safe async runner to execute the async create method
+        return run_async(
             cls.create(
                 file=file,
                 model=model,

onellm/chat_completion.py

@@ -25,15 +25,21 @@
 completions from various providers in a manner compatible with OpenAI's API.
 """
 
-import asyncio
 import logging
 import warnings
 from typing import Any, AsyncGenerator, Dict, List, Optional, Union
 
-from .providers.base import get_provider_with_fallbacks
+from .providers.base import get_provider_with_fallbacks, parse_model_name
 from .models import ChatCompletionResponse, ChatCompletionChunk
 from .utils.fallback import FallbackConfig
-from .validators import validate_model_name, validate_messages, validate_stream
+from .utils.async_helpers import run_async
+from .validators import (
+    validate_model_name,
+    validate_messages,
+    validate_stream,
+    validate_chat_params,
+    validate_provider_model,
+)
 
 class ChatCompletion:
     """Class for creating chat completions with various providers."""
@@ -262,11 +268,20 @@ def create(
         validate_model_name(model)
         validate_messages(messages)
         validate_stream(stream)
+        
+        # Validate all parameters
+        validate_chat_params(**kwargs)
+        
+        # Parse model and validate for provider
+        provider_name, model_without_prefix = parse_model_name(model)
+        validate_provider_model(model_without_prefix, provider_name)
 
         # Validate fallback models if provided
         if fallback_models:
             for i, fallback_model in enumerate(fallback_models):
                 validate_model_name(fallback_model)
+                fb_provider, fb_model = parse_model_name(fallback_model)
+                validate_provider_model(fb_model, fb_provider)
 
         # Process fallback configuration
         fb_config = None
@@ -294,24 +309,13 @@ def create(
             provider, messages, stream, kwargs
         )
 
-        # Call the provider's method synchronously
-        if stream:
-            # For streaming, we need to use async properly
-            # Create a new event loop to run the async code
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            return loop.run_until_complete(
-                provider.create_chat_completion(
-                    messages=messages, model=model_name, stream=stream, **processed_kwargs
-                )
-            )
-        else:
-            # For non-streaming, we can just run and get the result
-            return asyncio.run(
-                provider.create_chat_completion(
-                    messages=messages, model=model_name, stream=stream, **processed_kwargs
-                )
+        # Call the provider's method synchronously using our safe async runner
+        # This handles edge cases like Jupyter notebooks, existing event loops, etc.
+        return run_async(
+            provider.create_chat_completion(
+                messages=messages, model=model_name, stream=stream, **processed_kwargs
             )
+        )
 
     @classmethod
     async def acreate(
@@ -359,11 +363,20 @@ async def acreate(
         validate_model_name(model)
         validate_messages(messages)
         validate_stream(stream)
+        
+        # Validate all parameters
+        validate_chat_params(**kwargs)
+        
+        # Parse model and validate for provider
+        provider_name, model_without_prefix = parse_model_name(model)
+        validate_provider_model(model_without_prefix, provider_name)
 
         # Validate fallback models if provided
         if fallback_models:
             for i, fallback_model in enumerate(fallback_models):
                 validate_model_name(fallback_model)
+                fb_provider, fb_model = parse_model_name(fallback_model)
+                validate_provider_model(fb_model, fb_provider)
 
         # Process fallback configuration
         fb_config = None

onellm/completion.py

@@ -25,12 +25,19 @@
 completions from various providers in a manner compatible with OpenAI's API.
 """
 
-import asyncio
 from typing import Any, AsyncGenerator, List, Optional, Union
 
-from .providers.base import get_provider_with_fallbacks
+from .providers.base import get_provider_with_fallbacks, parse_model_name
 from .models import CompletionResponse
 from .utils.fallback import FallbackConfig
+from .utils.async_helpers import run_async
+from .validators import (
+    validate_model_name,
+    validate_prompt,
+    validate_stream,
+    validate_completion_params,
+    validate_provider_model,
+)
 
 class Completion:
     """Class for creating text completions with various providers."""
@@ -76,9 +83,24 @@ def create(
             ... )
             >>> print(response.choices[0].text)
         """
-        # Validate prompt
-        if not prompt or not prompt.strip():
-            raise ValueError("Prompt cannot be empty")
+        # Validate inputs
+        validate_model_name(model)
+        validate_prompt(prompt)
+        validate_stream(stream)
+        
+        # Validate all parameters
+        validate_completion_params(**kwargs)
+        
+        # Parse model and validate for provider
+        provider_name, model_without_prefix = parse_model_name(model)
+        validate_provider_model(model_without_prefix, provider_name)
+        
+        # Validate fallback models if provided
+        if fallback_models:
+            for fallback_model in fallback_models:
+                validate_model_name(fallback_model)
+                fb_provider, fb_model = parse_model_name(fallback_model)
+                validate_provider_model(fb_model, fb_provider)
 
         # Process fallback configuration
         fb_config = None
@@ -104,25 +126,13 @@ def create(
             fallback_config=fb_config,
         )
 
-        # Call the provider's method synchronously
-        if stream:
-            # For streaming, we need to use async properly
-            # Create a new event loop to run the async code in a synchronous context
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            return loop.run_until_complete(
-                provider.create_completion(
-                    prompt=prompt, model=model_name, stream=stream, **kwargs
-                )
-            )
-        else:
-            # For non-streaming, we can just run and get the result
-            # asyncio.run creates a new event loop, runs the coroutine, and closes the loop
-            return asyncio.run(
-                provider.create_completion(
-                    prompt=prompt, model=model_name, stream=stream, **kwargs
-                )
+        # Call the provider's method synchronously using our safe async runner
+        # This handles edge cases like Jupyter notebooks, existing event loops, etc.
+        return run_async(
+            provider.create_completion(
+                prompt=prompt, model=model_name, stream=stream, **kwargs
             )
+        )
 
     @classmethod
     async def acreate(
@@ -166,9 +176,24 @@ async def acreate(
             ... )
             >>> print(response.choices[0].text)
         """
-        # Validate prompt
-        if not prompt or not prompt.strip():
-            raise ValueError("Prompt cannot be empty")
+        # Validate inputs
+        validate_model_name(model)
+        validate_prompt(prompt)
+        validate_stream(stream)
+        
+        # Validate all parameters
+        validate_completion_params(**kwargs)
+        
+        # Parse model and validate for provider
+        provider_name, model_without_prefix = parse_model_name(model)
+        validate_provider_model(model_without_prefix, provider_name)
+        
+        # Validate fallback models if provided
+        if fallback_models:
+            for fallback_model in fallback_models:
+                validate_model_name(fallback_model)
+                fb_provider, fb_model = parse_model_name(fallback_model)
+                validate_provider_model(fb_model, fb_provider)
 
         # Process fallback configuration
         fb_config = None

onellm/embedding.py

@@ -25,12 +25,12 @@
 from various providers in a manner compatible with OpenAI's API.
 """
 
-import asyncio
 from typing import List, Optional, Union
 
 from .providers.base import get_provider_with_fallbacks
 from .models import EmbeddingResponse
 from .utils.fallback import FallbackConfig
+from .utils.async_helpers import run_async
 from .errors import InvalidRequestError
 
 def validate_embedding_input(input_data: Union[str, List[str]]) -> None:
@@ -109,8 +109,8 @@ def create(
             fallback_config=fb_config,
         )
 
-        # Call the provider's method synchronously by running the async method in an event loop
-        return asyncio.run(
+        # Call the provider's method synchronously using our safe async runner
+        return run_async(
             provider.create_embedding(input=input, model=model_name, **kwargs)
         )