Refactor OllamaModel to remove context from request data and update m…

…essage format for LLM interaction
Sedrowow · Sep 16, 2024 · 1362dac · 1362dac
1 parent 9a6a4d1
commit 1362dac
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -0,0 +1,3 @@
+{
+    "python.REPL.enableREPLSmartSend": false
+}
diff --git a/app/models/ollama_model.py b/app/models/ollama_model.py
@@ -28,7 +28,6 @@ def format_user_request_for_llm(self, original_user_request, step_num, screensho
             str: The formatted request as a JSON string.
         """
         request_data = {
-            'context': self.context,
             'original_user_request': original_user_request,
             'step_num': step_num
         }
@@ -38,9 +37,18 @@ def format_user_request_for_llm(self, original_user_request, step_num, screensho
 
     def send_message_to_llm(self, formatted_user_request: str) -> Any:
         try:
-            response = ollama.generate(
+            response = ollama.chat(
                 model=self.model_name,
-                prompt=formatted_user_request
+                messages=[
+                    {
+                        'role': 'instructions',
+                        'content': self.context
+                    },
+                    {
+                        'role': 'prompt',
+                        'content': formatted_user_request
+                    },
+                     ]
             )
             return response
         except Exception as e: