Local gemma model in ollama integratuion

src/agent/main_agent/agent.py

@@ -33,10 +33,11 @@ def get_llm_model(
 ):
     if AIModel.Ollama_DeepSeek_14b.value in model_name.model:
         return ChatOllama(model=AIModel.Ollama_DeepSeek_14b.value).with_structured_output(output_pydantic_class)
-
+    elif AIModel.Ollama_Gemma_7b.value in model_name.model:
+        return ChatOllama(model=AIModel.Ollama_Gemma_7b.value).with_structured_output(output_pydantic_class)
     else:
         #we could specify which model to use here always, for now if not deepseek then just use gpt4-o
-        return ChatOpenAI(model=AIModel.GPT_4O.value).with_structured_output(output_pydantic_class)
+        return ChatOpenAI(model_name=AIModel.GPT_4O.value).with_structured_output(output_pydantic_class)
 
 
 def run_main_agent(

src/agent/prompt_enahncer/prompt.py

@@ -1,7 +1,5 @@
 agents_prompt = """ You are an prompt enahncer agent with over 10 years of experince in ehancing user prompt for SQA purposes.
-
 You will be provided with a user query which primarily will be a qa test to perform and extracted web snippet on which the tests will be performed.
-
 Your job is to analyse the user prompt, and enahnce it to perform best sqa on the web snippet accordingly.
 
 ' User QA prompt ': "{input}"
@@ -22,6 +20,8 @@
 - Do not enhance the prompt to test UI/animations of a webpage please.
 - If the user has prompt to test the UI/animation, please get rid of that part in the prompt, and enhance only the functionality testing part.
 
+- Make sure that the url is always mentioned in the prompt.
+
 Output Structure:
 Return the result as a JSON object in the following format ONLY:
 { 

src/models/models.py

@@ -5,4 +5,5 @@
 class AIModel(Enum):
     GPT_4 = "gpt-4"
     GPT_4O = "gpt-4o"
-    Ollama_DeepSeek_14b = "deepseek-r1:14b"
+    Ollama_DeepSeek_14b = "deepseek-r1:14b"
+    Ollama_Gemma_7b = "gemma:7b"

src/utils/config.py

@@ -20,7 +20,7 @@
                "gemini-1.5-flash-8b-latest", "gemini-2.0-flash-thinking-exp-01-21", "gemini-2.0-pro-exp-02-05",
                "gemini-2.5-pro-preview-03-25", "gemini-2.5-flash-preview-04-17"],
     "ollama": ["qwen2.5:7b", "qwen2.5:14b", "qwen2.5:32b", "qwen2.5-coder:14b", "qwen2.5-coder:32b", "llama2:7b",
-               "deepseek-r1:14b", "deepseek-r1:32b"],
+               "deepseek-r1:14b", "deepseek-r1:32b", "gemma:7b"],
     "azure_openai": ["gpt-4o", "gpt-4", "gpt-3.5-turbo"],
     "mistral": ["pixtral-large-latest", "mistral-large-latest", "mistral-small-latest", "ministral-8b-latest"],
     "alibaba": ["qwen-plus", "qwen-max", "qwen-vl-max", "qwen-vl-plus", "qwen-turbo", "qwen-long"],