From f1f6c294e9a7a2c1c9f21f50e5bc87b5cf0b562e Mon Sep 17 00:00:00 2001
From: Paulslewis66 <35843981+Paulslewis66@users.noreply.github.com>
Date: Sat, 8 Feb 2025 21:04:02 +0000
Subject: [PATCH] Update ssh_server.py

Added Ollama Model
---
 SSH/ssh_server.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/SSH/ssh_server.py b/SSH/ssh_server.py
index f2c3a25..342fdcb 100755
--- a/SSH/ssh_server.py
+++ b/SSH/ssh_server.py
@@ -18,6 +18,7 @@ from operator import itemgetter
 from langchain_openai import ChatOpenAI
 from langchain_aws import ChatBedrock, ChatBedrockConverse
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_ollama import ChatOllama 
 from langchain_core.messages import HumanMessage, SystemMessage, trim_messages
 from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
 from langchain_core.runnables.history import RunnableWithMessageHistory
@@ -313,6 +314,10 @@ def choose_llm():
         llm_model = ChatOpenAI(
             model=model_name
         )
+    elif llm_provider_name == 'ollama':
+            llm_model = ChatOllama(
+            model=model_name
+        )
     elif llm_provider_name == 'aws':
         llm_model = ChatBedrockConverse(
             model=model_name,