From 46658c0d34f11731ff1953850e033809659b6eab Mon Sep 17 00:00:00 2001 From: Mohsin Ali Irfan Date: Tue, 17 Jun 2025 11:41:04 +0500 Subject: [PATCH] Local gemma model in ollama integratuion --- src/agent/main_agent/agent.py | 5 +++-- src/agent/prompt_enahncer/prompt.py | 4 ++-- src/models/models.py | 3 ++- src/utils/config.py | 2 +- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/agent/main_agent/agent.py b/src/agent/main_agent/agent.py index 27e04054b..92c7115c8 100644 --- a/src/agent/main_agent/agent.py +++ b/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( diff --git a/src/agent/prompt_enahncer/prompt.py b/src/agent/prompt_enahncer/prompt.py index 8e22463a9..9b2a4f387 100644 --- a/src/agent/prompt_enahncer/prompt.py +++ b/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: { diff --git a/src/models/models.py b/src/models/models.py index 4f67dfc6a..1a20f78fe 100644 --- a/src/models/models.py +++ b/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" \ No newline at end of file + Ollama_DeepSeek_14b = "deepseek-r1:14b" + Ollama_Gemma_7b = "gemma:7b" \ No newline at end of file diff --git a/src/utils/config.py b/src/utils/config.py index de82bb9e9..b3a5a9d24 100644 --- a/src/utils/config.py +++ b/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"],