Merge pull request #17 from db-agent/gemini

kchandan · web-flow · commit 00342b16a85d · 2025-01-10T18:04:53.000-05:00
Added ollama changes and supported model
diff --git a/helpers/supported_models.py b/helpers/supported_models.py
@@ -8,6 +8,7 @@
                     "defog/sqlcoder-70b-alpha",
                     "meta-llama/Llama-3.3-70B-Instruct"],
     "ollama": ["llama3.2:1b",
+               "llama3.2:latest",
                "llama3.3",
                "hf.co/defog/sqlcoder-7b-2"],
     "vllm": ["microsoft/Phi-3.5-mini-instruct", 
diff --git a/textgen/ollama.py b/textgen/ollama.py
@@ -13,6 +13,8 @@ def construct_sql_payload(self, user_question, db_schema):
         prompt = (
             "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n"
             f"Generate a SQL query only to answer this question without explanation: `{user_question}`\n"
+            "considering values like true,TRUE,yes,Yes,YES in a case-insensitive manner.\n"
+            "considering values like false,FALSE,no,No,NO in a case-insensitive manner.\n"
             "DDL statements:\n"
             f"{db_schema}<|start_header_id|>assistant<|end_header_id|>\n\n"
         )
@@ -22,7 +24,21 @@ def construct_sql_payload(self, user_question, db_schema):
             "stream": False,
             "temperature": 0
         }
+    def construct_generic_payload(self, user_question):
+        prompt = user_question
+        return {
+            "model": self.model_name,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.7,
+            "max_new_tokens": 4096,
+            "max_length": 4096,
+            "top_p": 0.9,
+            "repetition_penalty": 1.1,
+            "stream": False,
+            "max_tokens": 1024
 
+        }
+        
     def parse_response(self, response):
         logger.info(f"Parsing response for Ollama {response}")
         return response.get('message', {}).get('content', '').strip()
