diff --git a/requirements.txt b/requirements.txt index f7055242..a853cfa4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ langchain-ibm==0.3.10 langchain_mcp_adapters==0.0.9 langgraph==0.3.34 langchain-community +langchain-ollama \ No newline at end of file diff --git a/src/agent/intent_classifier/agent.py b/src/agent/intent_classifier/agent.py index 8e10fcce..168ae508 100644 --- a/src/agent/intent_classifier/agent.py +++ b/src/agent/intent_classifier/agent.py @@ -24,7 +24,8 @@ def run_agent(self) -> IntentClassifierOutput: agents_prompt=self.agent_prompt, input_to_prompt={ "input": self.user_prompt - } + }, + model_name=self.llm ) logger.info(f"Intent Classifier Agent finished. Output: {output}") diff --git a/src/agent/intent_classifier/prompt.py b/src/agent/intent_classifier/prompt.py index 0eeeced1..f26cf7d5 100644 --- a/src/agent/intent_classifier/prompt.py +++ b/src/agent/intent_classifier/prompt.py @@ -18,14 +18,18 @@ Output: {{ "intent": true }} Example 3: +User query: "please buy this on the website" +Output: {{ "intent": true }} + +Example 4: User query: "How do I use OpenAI's API to generate text?" Output: {{ "intent": false }} -Example 4: +Example 5: User query: "Can you summarize this PDF document?" Output: {{ "intent": false }} -Example 5: +Example 6: User query: "Please test if the color of the button is blue". Output: {{ @@ -33,7 +37,7 @@ intent": false }} -Example 5: +Example 7: User query: "Please test the submit button along with the hoover animation". Output: {{ @@ -47,6 +51,8 @@ Now classify the following user query: User query: "{input}" +- Always make sure that performing any sought of functionality would also be considered as SQA because user might sometimes just ask simply to do this/that on the website. + Output Structre: Return the result as a JSON object in the following format ONLY: { diff --git a/src/agent/main_agent/agent.py b/src/agent/main_agent/agent.py index d686f9ff..27e04054 100644 --- a/src/agent/main_agent/agent.py +++ b/src/agent/main_agent/agent.py @@ -5,9 +5,8 @@ import os import logging from langchain_openai import ChatOpenAI +from langchain_ollama.chat_models import ChatOllama from dotenv import load_dotenv -from langchain.prompts import PromptTemplate -from langchain.output_parsers import PydanticOutputParser from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate from ...outputdata.output_data import write_data_to_file load_dotenv() @@ -27,17 +26,30 @@ def get_chatprompt_templates(agents_prompt: str, input_keys: list[str]) -> ChatP agents_prompt_template, user_input_template ]) + +def get_llm_model( + model_name: Optional[AIModel], + output_pydantic_class +): + if AIModel.Ollama_DeepSeek_14b.value in model_name.model: + return ChatOllama(model=AIModel.Ollama_DeepSeek_14b.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) + + def run_main_agent( output_pydantic_class, agents_name: str, agents_prompt: str, input_to_prompt: dict, - model_name: Optional[AIModel] = AIModel.GPT_4O, + model_name: Optional[str], ) -> Any: - llm_model = ChatOpenAI(model_name=model_name.value).with_structured_output(output_pydantic_class) - logger.info(f"Running agent '{agents_name}' with model: {model_name.value}") + llm_model = get_llm_model(model_name, output_pydantic_class) + + logger.info(f"Running agent '{agents_name}' with model: {model_name}") chat_prompt = get_chatprompt_templates(agents_prompt, list(input_to_prompt.keys())) diff --git a/src/agent/orchestrator/agent_orchestrator.py b/src/agent/orchestrator/agent_orchestrator.py index c72260ef..10121ea4 100644 --- a/src/agent/orchestrator/agent_orchestrator.py +++ b/src/agent/orchestrator/agent_orchestrator.py @@ -69,8 +69,6 @@ def __init__( self.planner_llm = planner_llm self.use_vision_for_planner = use_vision_for_planner - - self.builder = StateGraph(State) self.builder.add_node("intent_classifier", self.intent_classifier) @@ -151,15 +149,26 @@ def _get_output_value(self, output, key, default=None): async def snippet_extractor(self, state: State) -> State: logger.info("\n\n SNIPPET EXTRACTOR NODE...\n") - output = await SnippetExtractorAgent(user_prompt=self.user_query, url=self.url).run_agent() + + output = await SnippetExtractorAgent( + llm=self.llm, + user_prompt=self.user_query, + url=self.url + + ).run_agent() state['extracted_snippet_agent_msg'] = self._get_output_value(output, "agent_msg", "") state['snippet_check'] = self._get_output_value(output, "snippet_check", False) state['extracted_snippet'] = self._get_output_value(output, "extracted_snippet", "") + + print("\n\n\n EXTRACTED SNIPPET: ", output.extracted_snippet) + + return state def QA_possibility(self, state: State) -> State: logger.info("\n\n QA POSSIBILTY CHECKER AGENT...\n") output = QAPossibilityChecker( + llm=self.llm, user_prompt=self.user_query, extracted_snippet=state['extracted_snippet'] ).run_agent() @@ -170,6 +179,7 @@ def QA_possibility(self, state: State) -> State: def prompt_enhancer(self, state: State) -> State: logger.info("\n\n PROMPT ENHANCER AGENT...\n") output = PromptEnhancerAgent( + llm=self.llm, user_prompt=self.user_query, extracted_snippet=state['extracted_snippet'] ).run_agent() diff --git a/src/agent/prompt_enahncer/agent.py b/src/agent/prompt_enahncer/agent.py index 3107a816..e07f3893 100644 --- a/src/agent/prompt_enahncer/agent.py +++ b/src/agent/prompt_enahncer/agent.py @@ -6,12 +6,13 @@ logger = logging.getLogger(__name__) class PromptEnhancerAgent: - def __init__(self, user_prompt: str, extracted_snippet: str) -> None: + def __init__(self,llm: str, user_prompt: str, extracted_snippet: str) -> None: logger.info("Initializing PromptEnhancerAgent...") self.output_pydantic_class = PromptEnhancerOutput self.user_prompt = user_prompt self.agent_prompt = agents_prompt self.extracted_snippet = extracted_snippet + self.llm = llm def run_agent(self) -> PromptEnhancerOutput: output = run_main_agent( @@ -21,7 +22,8 @@ def run_agent(self) -> PromptEnhancerOutput: input_to_prompt={ "input": self.user_prompt, "extracted_snippet": self.extracted_snippet - } + }, + model_name=self.llm ) logger.info(f"Prompt Enhancers Agent finished...") diff --git a/src/agent/qa_possibilty_checker/agent.py b/src/agent/qa_possibilty_checker/agent.py index 1c539b54..e04429bf 100644 --- a/src/agent/qa_possibilty_checker/agent.py +++ b/src/agent/qa_possibilty_checker/agent.py @@ -6,12 +6,13 @@ logger = logging.getLogger(__name__) class QAPossibilityChecker: - def __init__(self, user_prompt: str, extracted_snippet: str) -> None: + def __init__(self, llm: str, user_prompt: str, extracted_snippet: str) -> None: logger.info("Initializing QAPossibilityChecker") self.output_pydantic_class = QAPossibiltyChecker self.user_prompt = user_prompt self.agent_prompt = agents_prompt self.extracted_snippet = extracted_snippet + self.llm = llm def run_agent(self) -> QAPossibiltyChecker: @@ -22,7 +23,8 @@ def run_agent(self) -> QAPossibiltyChecker: input_to_prompt={ "input": self.user_prompt, "extracted_snippet": self.extracted_snippet - } + }, + model_name=self.llm ) logger.info(f"QA Possibilty Checker finished Output...: {output}") diff --git a/src/agent/qa_possibilty_checker/prompt.py b/src/agent/qa_possibilty_checker/prompt.py index a2c7fcf4..966f79d4 100644 --- a/src/agent/qa_possibilty_checker/prompt.py +++ b/src/agent/qa_possibilty_checker/prompt.py @@ -2,7 +2,7 @@ You are an amazing QA possibilty checker agent and your job is to classify if the QA test are possibile for a web page snippet or not. You will be provided with a - -1. user prompt which are going to be tests. +1. user prompt which are going to be tests/qa/functionality to be tested. 2. web page snippet which will be the extracted snippet from the web page that is to be tested. Analyse the user prompt and the extracted webpage snippet and classify if the QA enterd by the user are possible on the web snippet or not. @@ -11,14 +11,14 @@ ' Extracte Web Page Snippet ': ' {extracted_snippet} ' Output Structure: -eturn the result as a JSON object in the following format ONLY: +return the result as a JSON object in the following format ONLY: { "agent_msg": "QA is possible on the extracted Snippet" "qa_possibilty": true } or { - "agent_msg": "QA is not possible on the extracted Snippet" + "agent_msg": "QA is not possible on the extracted Snippet (due to this reasom)" "qa_possibilty": false } """ \ No newline at end of file diff --git a/src/agent/snippet_extractor/agent.py b/src/agent/snippet_extractor/agent.py index 6ed08086..3b59ec60 100644 --- a/src/agent/snippet_extractor/agent.py +++ b/src/agent/snippet_extractor/agent.py @@ -10,12 +10,13 @@ logger = logging.getLogger(__name__) class SnippetExtractorAgent: - def __init__(self, user_prompt: str, url: str) -> None: + def __init__(self, llm: str, user_prompt: str, url: str) -> None: logger.info("Initializing SnippetExtractorAgent") self.output_pydantic_class = SnippetExtractorOutput self.user_prompt = user_prompt self.agent_prompt = agents_prompt self.url = url + self.llm = llm def store_webpage_code(self) -> bool: logger.info("Storing webpage code into webpage_code.html under webpage directory....") @@ -158,7 +159,8 @@ async def run_agent(self) -> SnippetExtractorOutput: input_to_prompt={ 'input': self.user_prompt, 'webpage_code': chunk - } + }, + model_name=self.llm ) if output.snippet_check: logger.info(f"Snippet found in chunk {i+1}") diff --git a/src/models/models.py b/src/models/models.py index afa76f71..4f67dfc6 100644 --- a/src/models/models.py +++ b/src/models/models.py @@ -4,4 +4,5 @@ class AIModel(Enum): GPT_4 = "gpt-4" - GPT_4O = "gpt-4o" \ No newline at end of file + GPT_4O = "gpt-4o" + Ollama_DeepSeek_14b = "deepseek-r1:14b" \ No newline at end of file