Integrated ollama locally to run llm locally

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

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}")

src/agent/intent_classifier/prompt.py

@@ -18,22 +18,26 @@
 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:
 {{ 
 agent_msg: "UI elements cannot be tested"
 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:
 { 

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()))
 

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()

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...")

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}")

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 
 }
 """

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}")

src/models/models.py

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