WIP: Refactor generate client to include usage from providers

examples/core.py

@@ -5,18 +5,21 @@
 from pprint import pprint
 import os
 from dotenv import load_dotenv
+import asyncio
 load_dotenv()
 
 def run_provider(provider, model, api_key, **kwargs):
-    llm = LLMCore(provider=provider, api_key=api_key, **kwargs)
 
+    llm = LLMCore(provider=provider, api_key=api_key, **kwargs)
     latencies = {}
+
+    """
     chat_request = build_chat_request(model, chat_input="Hello, my name is Jason Json", is_stream=False)
 
-    import asyncio
     response_async = asyncio.run(llm.achat(**chat_request))
     pprint(response_async)
     latencies["async (ms)"]= response_async.metrics["latency_s"]*1000
+    """
 
     # stream
     print("\nasync stream")
@@ -114,6 +117,7 @@ def build_chat_request(model: str, chat_input: str, is_stream: bool, max_tokens:
 #     pprint(latencies)
 # # we need credits
 
+"""     
 provider = "azure"
 model = "gpt-4o-mini"
 for _ in range(1):
@@ -122,24 +126,28 @@ def build_chat_request(model: str, chat_input: str, is_stream: bool, max_tokens:
                             api_version=os.environ["AZURE_API_VERSION"],
                             api_endpoint=os.environ["AZURE_API_ENDPOINT"])
     pprint(latencies)
+"""
 
-provider = "azure"
-model = "o1-preview"
-for _ in range(1):
-    latencies = run_provider(provider=provider, model=model, 
-                            api_key=os.environ["AZURE_API_KEY"], 
-                            api_version=os.environ["AZURE_API_VERSION"],
-                            api_endpoint=os.environ["AZURE_API_ENDPOINT"])
-    pprint(latencies)
+# provider = "azure"
+# model = "o1-preview"
+# for _ in range(1):
+#     latencies = run_provider(provider=provider, model=model, 
+#                             api_key=os.environ["AZURE_API_KEY"], 
+#                             api_version=os.environ["AZURE_API_VERSION"],
+#                             api_endpoint=os.environ["AZURE_API_ENDPOINT"])
+#     pprint(latencies)
+# # we need a deployment
+
 
-provider = "azure"
-model = "o1-mini"
-for _ in range(1):
-    latencies = run_provider(provider=provider, model=model, 
-                            api_key=os.environ["AZURE_API_KEY"], 
-                            api_version=os.environ["AZURE_API_VERSION"],
-                            api_endpoint=os.environ["AZURE_API_ENDPOINT"])
-    pprint(latencies)
+# provider = "azure"
+# model = "o1-mini"
+# for _ in range(1):
+#     latencies = run_provider(provider=provider, model=model, 
+#                             api_key=os.environ["AZURE_API_KEY"], 
+#                             api_version=os.environ["AZURE_API_VERSION"],
+#                             api_endpoint=os.environ["AZURE_API_ENDPOINT"])
+#     pprint(latencies)
+# # we need a deployment
 
 # provider = "azure"
 # model = "gpt-4o"

libs/core/llmstudio_core/providers/azure.py

@@ -109,7 +109,7 @@ def generate_client(self, request: ChatRequest) -> Any:
             base_args = {
                 "model": request.model,
                 "messages": messages,
-                "stream": True,
+                "stream": request.is_stream,
             }
 
             combined_args = {
@@ -145,14 +145,14 @@ def prepare_messages(self, request: ChatRequest):
                 else request.chat_input
             )
 
-    async def aparse_response(
+    async def _aparse_response(
         self, response: AsyncGenerator, **kwargs
     ) -> AsyncGenerator[str, None]:
-        result = self.parse_response(response=response, **kwargs)
+        result = self._parse_response(response=response, **kwargs)
         for chunk in result:
             yield chunk
 
-    def parse_response(self, response: AsyncGenerator, **kwargs) -> Any:
+    def _parse_response(self, response: AsyncGenerator, **kwargs) -> Any:
         """
         Processes a generator response and yields processed chunks.
 

libs/core/llmstudio_core/providers/bedrock/anthropic.py

@@ -80,12 +80,12 @@ def generate_client(self, request: ChatRequest) -> Coroutine[Any, Any, Generator
         except Exception as e:
             raise ProviderError(str(e))
 
-    async def aparse_response(
+    async def _aparse_response(
         self, response: Any, **kwargs
     ) -> AsyncGenerator[Any, None]:
-        return self.parse_response(response=response, **kwargs)
+        return self._parse_response(response=response, **kwargs)
 
-    def parse_response(self, response: AsyncGenerator[Any, None], **kwargs) -> Any:
+    def _parse_response(self, response: AsyncGenerator[Any, None], **kwargs) -> Any:
         tool_name = None
         tool_arguments = ""
         tool_id = None

libs/core/llmstudio_core/providers/bedrock/provider.py

@@ -32,12 +32,12 @@ def generate_client(self, request: ChatRequest) -> Coroutine[Any, Any, Generator
         self.selected_model = self._get_provider(request.model)
         return self.selected_model.generate_client(request=request)
 
-    async def aparse_response(
+    async def _aparse_response(
         self, response: Any, **kwargs
     ) -> AsyncGenerator[Any, None]:
-        result = await self.selected_model.aparse_response(response=response, **kwargs)
+        result = await self.selected_model._aparse_response(response=response, **kwargs)
         for chunk in result:
             yield chunk
 
-    def parse_response(self, response: AsyncGenerator[Any, None], **kwargs) -> Any:
-        return self.selected_model.parse_response(response=response, **kwargs)
+    def _parse_response(self, response: AsyncGenerator[Any, None], **kwargs) -> Any:
+        return self.selected_model._parse_response(response=response, **kwargs)

libs/core/llmstudio_core/providers/openai.py

@@ -38,19 +38,19 @@ def generate_client(self, request: ChatRequest) -> Generator:
                     if isinstance(request.chat_input, str)
                     else request.chat_input
                 ),
-                stream=True,
+                stream=request.is_stream,
                 **request.parameters,
             )
         except openai._exceptions.APIError as e:
             raise ProviderError(str(e))
 
-    async def aparse_response(
+    async def _aparse_response(
         self, response: AsyncGenerator, **kwargs
     ) -> AsyncGenerator[str, None]:
-        result = self.parse_response(response=response, **kwargs)
+        result = self._parse_response(response=response, **kwargs)
         for chunk in result:
             yield chunk
 
-    def parse_response(self, response: Generator, **kwargs) -> Generator:
+    def _parse_response(self, response: Generator, **kwargs) -> Generator:
         for chunk in response:
             yield chunk.model_dump()