Providers Refactor

libs/core/llmstudio_core/providers/azure.py

@@ -1,5 +1,4 @@
 import ast
-import asyncio
 import json
 import os
 import time
@@ -60,57 +59,7 @@ def validate_request(self, request: ChatRequest):
 
     async def agenerate_client(self, request: ChatRequest) -> Any:
         """Generate an AzureOpenAI client"""
-
-        self.is_llama = "llama" in request.model.lower()
-        self.is_openai = "gpt" in request.model.lower()
-        self.has_tools = request.parameters.get("tools") is not None
-        self.has_functions = request.parameters.get("functions") is not None
-
-        try:
-            messages = self.prepare_messages(request)
-
-            # Prepare the optional tool-related arguments
-            tool_args = {}
-            if not self.is_llama and self.has_tools and self.is_openai:
-                tool_args = {
-                    "tools": request.parameters.get("tools"),
-                    "tool_choice": "auto" if request.parameters.get("tools") else None,
-                }
-
-            # Prepare the optional function-related arguments
-            function_args = {}
-            if not self.is_llama and self.has_functions and self.is_openai:
-                function_args = {
-                    "functions": request.parameters.get("functions"),
-                    "function_call": "auto"
-                    if request.parameters.get("functions")
-                    else None,
-                }
-
-            # Prepare the base arguments
-            base_args = {
-                "model": request.model,
-                "messages": messages,
-                "stream": True,
-            }
-
-            # Combine all arguments
-            combined_args = {
-                **base_args,
-                **tool_args,
-                **function_args,
-                **request.parameters,
-            }
-            # Perform the asynchronous call
-            return await asyncio.to_thread(
-                self._client.chat.completions.create, **combined_args
-            )
-
-        except openai._exceptions.APIConnectionError as e:
-            raise ProviderError(f"There was an error reaching the endpoint: {e}")
-
-        except openai._exceptions.APIStatusError as e:
-            raise ProviderError(e.response.json())
+        return self.generate_client(request=request)
 
     def generate_client(self, request: ChatRequest) -> Any:
         """Generate an AzureOpenAI client"""
@@ -185,17 +134,11 @@ def prepare_messages(self, request: ChatRequest):
     async def aparse_response(
         self, response: AsyncGenerator, **kwargs
     ) -> AsyncGenerator[str, None]:
-        if self.is_llama and (self.has_tools or self.has_functions):
-            for chunk in self.handle_tool_response(response, **kwargs):
-                if chunk:
-                    yield chunk
-        else:
-            for chunk in response:
-                c = chunk.model_dump()
-                if c.get("choices"):
-                    yield c
+        result = self.parse_response(response=response, **kwargs)
+        for chunk in result:
+            yield chunk
 
-    def parse_response(self, response: Generator, **kwargs) -> Any:
+    def parse_response(self, response: AsyncGenerator, **kwargs) -> Any:
         if self.is_llama and (self.has_tools or self.has_functions):
             for chunk in self.handle_tool_response(response, **kwargs):
                 if chunk:

libs/core/llmstudio_core/providers/openai.py

@@ -1,4 +1,3 @@
-import asyncio
 import os
 from typing import Any, AsyncGenerator, Coroutine, Generator
 
@@ -26,21 +25,7 @@ async def agenerate_client(
         self, request: ChatRequest
     ) -> Coroutine[Any, Any, Generator]:
         """Generate an OpenAI client"""
-
-        try:
-            return await asyncio.to_thread(
-                self._client.chat.completions.create,
-                model=request.model,
-                messages=(
-                    [{"role": "user", "content": request.chat_input}]
-                    if isinstance(request.chat_input, str)
-                    else request.chat_input
-                ),
-                stream=True,
-                **request.parameters,
-            )
-        except openai._exceptions.APIError as e:
-            raise ProviderError(str(e))
+        return self.generate_client(request=request)
 
     def generate_client(self, request: ChatRequest) -> Generator:
         """Generate an OpenAI client"""
@@ -62,8 +47,9 @@ def generate_client(self, request: ChatRequest) -> Generator:
     async def aparse_response(
         self, response: AsyncGenerator, **kwargs
     ) -> AsyncGenerator[str, None]:
-        for chunk in response:
-            yield chunk.model_dump()
+        result = self.parse_response(response=response, **kwargs)
+        for chunk in result:
+            yield chunk
 
     def parse_response(self, response: Generator, **kwargs) -> Generator:
         for chunk in response: