feat: Use streaming setup at Query level (#214)

embedchain/config/ChatConfig.py

@@ -4,5 +4,11 @@ class ChatConfig(QueryConfig):
     """
     Config for the `chat` method, inherits from `QueryConfig`.
     """
-    def __init__(self):
-        pass
+    def __init__(self, stream: bool = False):
+        """
+        Initializes the QueryConfig instance.
+
+        :param stream: Optional. Control if response is streamed back to the user
+        :raises ValueError: If the template is not valid as template should contain $context and $query
+        """
+        super().__init__(stream=stream)

embedchain/config/InitConfig.py

@@ -6,7 +6,7 @@ class InitConfig(BaseConfig):
     """
     Config to initialize an embedchain `App` instance.
     """
-    def __init__(self, ef=None, db=None, stream_response=False):
+    def __init__(self, ef=None, db=None):
         """
         :param ef: Optional. Embedding function to use.
         :param db: Optional. (Vector) database to use for embeddings.
@@ -27,10 +27,6 @@ class InitConfig(BaseConfig):
             self.db = ChromaDB(ef=self.ef)
         else:
             self.db = db
-        
-        if not isinstance(stream_response, bool):
-            raise ValueError("`stream_respone` should be bool")
-        self.stream_response = stream_response
 
         return
 

embedchain/config/QueryConfig.py

@@ -22,11 +22,12 @@ class QueryConfig(BaseConfig):
     """
     Config for the `query` method.
     """
-    def __init__(self, template: Template = None):
+    def __init__(self, template: Template = None, stream: bool = False):
         """
         Initializes the QueryConfig instance.
 
         :param template: Optional. The `Template` instance to use as a template for prompt.
+        :param stream: Optional. Control if response is streamed back to the user
         :raises ValueError: If the template is not valid as template should contain $context and $query
         """
         if template is None:
@@ -35,3 +36,7 @@ class QueryConfig(BaseConfig):
             and re.search(context_re, template.template)):
             raise ValueError("`template` should have `query` and `context` keys")
         self.template = template
+
+        if not isinstance(stream, bool):
+            raise ValueError("`stream` should be bool")
+        self.stream = stream