feat: customize OpenAI parameters (#215)

embedchain/config/QueryConfig.py

@@ -34,18 +34,21 @@ query_re = re.compile(r"\$\{*query\}*")
 context_re = re.compile(r"\$\{*context\}*")
 history_re = re.compile(r"\$\{*history\}*")
 
-
 class QueryConfig(BaseConfig):
     """
     Config for the `query` method.
     """
 
-    def __init__(self, template: Template = None, history=None, stream: bool = False):
+    def __init__(self, template: Template = None, model = None, temperature = None, max_tokens = None, top_p = None, history = None, stream: bool = False):
         """
         Initializes the QueryConfig instance.
 
-        :param template: Optional. The `Template` instance to use as a
-        template for prompt.
+        :param template: Optional. The `Template` instance to use as a template for prompt.
+        :param model: Optional. Controls the OpenAI model used.
+        :param temperature: Optional. Controls the randomness of the model's output. 
+                            Higher values (closer to 1) make output more random, lower values make it more deterministic.
+        :param max_tokens: Optional. Controls how many tokens are generated.
+        :param top_p: Optional. Controls the diversity of words. Higher values (closer to 1) make word selection more diverse, lower values make words less diverse.
         :param history: Optional. A list of strings to consider as history.
         :param stream: Optional. Control if response is streamed back to user
         :raises ValueError: If the template is not valid as template should
@@ -65,6 +68,12 @@ class QueryConfig(BaseConfig):
             else:
                 template = DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE
 
+      
+        self.temperature = temperature if temperature else 0
+        self.max_tokens = max_tokens if max_tokens else 1000
+        self.model = model if model else "gpt-3.5-turbo-0613"
+        self.top_p = top_p if top_p else 1
+
         if self.validate_template(template):
             self.template = template
         else: