Implement model managers for Ollama, Gemini, and Mistral; update requirements.txt with new dependencies

models/llm_models.py

@@ -0,0 +1,137 @@
+from abc import ABC, abstractmethod
+import ollama
+from pydantic import BaseModel
+from pathlib import Path
+from google import genai
+from google.genai import types
+from mistralai import Mistral
+
+
+class AbstractModelManager(ABC):
+    def __init__(self, model_name, system_prompt_file="system.prompt"):
+        self.model_name = model_name
+        script_dir = Path(__file__).parent
+        self.system_prompt_file = script_dir / system_prompt_file
+
+    @abstractmethod
+    def is_model_loaded(self, model):
+        pass
+
+    @abstractmethod
+    def create_model(self, base_model, context_window=4096, temperature=0):
+        pass
+
+    @abstractmethod
+    def request(self, prompt):
+        pass
+
+    @abstractmethod
+    def delete(self):
+        pass
+
+class OllamaModelManager(AbstractModelManager):
+    def is_model_loaded(self, model):
+        loaded_models = [m.model for m in ollama.list().models]
+        return model in loaded_models or f'{model}:latest' in loaded_models
+
+    def create_model(self, base_model, context_window=4096, temperature=0):
+        with open(self.system_prompt_file, 'r') as f:
+            system = f.read()
+        
+        if not self.is_model_loaded(self.model_name):
+            print(f"Creating model {self.model_name}")
+            ollama.create(
+                model=self.model_name,
+                from_=base_model,
+                system=system,
+                parameters={
+                    "num_ctx": context_window,
+                    "temperature": temperature
+                }
+            )
+
+    def request(self, prompt):
+        response = ollama.chat(
+            model=self.model_name, 
+            messages=[{"role": "user", "content": prompt}],
+        )
+        response = response['message']['content']
+        return response
+    
+    def delete(self):
+        if self.is_model_loaded("C2Rust:latest"):
+            print(f"Deleting model {self.model_name}")
+            ollama.delete("C2Rust:latest")
+        else:
+            print(f"Model {self.model_name} not found, skipping deletion.")
+
+class GeminiModelManager(AbstractModelManager):
+    def __init__(self, api_key):
+        super().__init__()
+        self.client = genai.Client(api_key=api_key)
+        self.model = "gemini-2.0-flash"
+        # read system prompt from file
+        with open(self.system_prompt_file, 'r') as f:
+            self.system_instruction = f.read()
+
+
+    def is_model_loaded(self, model):
+        # Check if the specified model is the one set in the manager
+        return model == self.model
+
+    def create_model(self, base_model=None, context_window=4096, temperature=0):
+        # Initialize the Gemini model settings (if applicable)
+        self.model = base_model if base_model else "gemini-2.0-flash"
+
+    def request(self, prompt, temperature=0, context_window=4096):
+        # Request response from the Gemini model
+        response = self.client.models.generate_content(
+            model=self.model,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=temperature,
+                max_output_tokens=context_window,
+                system_instruction=self.system_instruction,
+            )
+        )
+        return response.text
+
+    def delete(self):
+        # Implement model deletion logic (if applicable)
+        self.model = None
+
+class MistralModelManager(AbstractModelManager):
+    def __init__(self, api_key, model_name="mistral-small-latest", system_prompt_file="system.prompt"):
+        super().__init__()
+        self.client = Mistral(api_key=api_key)
+        self.model = model_name
+        # read system prompt from file
+        with open(self.system_prompt_file, 'r') as f:
+            self.system_instruction = f.read()
+
+    def is_model_loaded(self, model):
+        # Check if the specified model is the one set in the manager
+        return model == self.model
+
+    def create_model(self, base_model=None, context_window=4096, temperature=0):
+        # Initialize the Mistral model settings (if applicable)
+        self.model = base_model if base_model else "mistral-small-latest"
+
+    def request(self, prompt, temperature=0, context_window=4096):
+        # Request response from the Mistral model
+        response = self.client.chat.complete(
+            messages=[
+            {
+                "role":"user",
+                "content": self.system_instruction + "\n" + prompt,
+            }
+            ],
+            model=self.model,
+            temperature=temperature,
+            max_tokens=context_window,
+        )
+        return response.text
+
+    def delete(self):
+        # Implement model deletion logic (if applicable)
+        self.model = None

requirements.txt

@@ -1,19 +1,40 @@
 annotated-types==0.7.0
 anyio==4.9.0
 beautifulsoup4==4.13.3
+cachetools==5.5.2
 certifi==2025.1.31
 charset-normalizer==3.4.1
-googlesearch-python==1.3.0
+google-ai-generativelanguage==0.6.15
+google-api-core==2.24.2
+google-api-python-client==2.166.0
+google-auth==2.38.0
+google-auth-httplib2==0.2.0
+google-genai==1.9.0
+googleapis-common-protos==1.69.2
+grpcio==1.71.0
+grpcio-status==1.71.0
 h11==0.14.0
 httpcore==1.0.7
+httplib2==0.22.0
 httpx==0.28.1
 idna==3.10
 ollama==0.4.7
+pathlib==1.0.1
+proto-plus==1.26.1
+protobuf==5.29.4
+pyasn1==0.6.1
+pyasn1_modules==0.4.2
 pydantic==2.11.1
 pydantic_core==2.33.0
+pyparsing==3.2.3
+python-dotenv==1.1.0
 requests==2.32.3
+rsa==4.9
 sniffio==1.3.1
 soupsieve==2.6
+tqdm==4.67.1
 typing-inspection==0.4.0
 typing_extensions==4.13.0
+uritemplate==4.1.1
 urllib3==2.3.0
+websockets==15.0.1