Merge pull request #345 from edenai/SD2-1328-important-streaming-should-work-with-the-new-llm-chat-endoints

[Fix] stream for llm completion

Author: juandavidcruzgomez

edenai_apis/features/llm/chat/chat_dataclass.py

@@ -1,6 +1,7 @@
-from typing import List, Optional, Union, Dict, Any, Literal
+from typing import List, Optional, Union, Dict, Any, Literal, Generator
 from enum import Enum
 from pydantic import BaseModel, Field, model_validator
+from litellm import ModelResponseStream
 
 
 class ChatRole(str, Enum):
@@ -211,3 +212,7 @@ class ChatDataClass(BaseModel):
     system_fingerprint: Optional[str] = Field(
         None, description="Identifier for the system version that processed the request"
     )
+
+
+class StreamChat(BaseModel):
+    stream: Generator[ModelResponseStream, None, None]

edenai_apis/llmengine/llm_engine.py

@@ -1,8 +1,6 @@
 import base64
 import json
 import mimetypes
-import os
-import re
 import uuid
 from io import BytesIO
 from typing import Dict, List, Literal, Optional, Type, Union
@@ -53,6 +51,9 @@
     AutomaticTranslationDataClass,
     LanguageDetectionDataClass,
 )
+from edenai_apis.features.llm.chat.chat_dataclass import (
+    StreamChat as StreamChatCompletion,
+)
 from edenai_apis.llmengine.clients import LLM_COMPLETION_CLIENTS
 from edenai_apis.llmengine.clients.completion import CompletionClient
 from edenai_apis.llmengine.mapping import Mappings
@@ -845,7 +846,10 @@ def completion(
             completion_params = completion_params
             call_params = self._prepare_args(**completion_params)
             response = self.completion_client.completion(**call_params, **kwargs)
-            response = ResponseModel.model_validate(response)
-            return response
+            if stream:
+                return StreamChatCompletion(stream=response)
+            else:
+                response = ResponseModel.model_validate(response)
+                return response
         except Exception as ex:
             raise ex