Changed default LLM to 'gpt-4o-mini'

Only used if the model isn't specified in the config or on the command line

README.md

@@ -1,5 +1,7 @@
 # DECEIVE
 
+<img align="right" src="DECEIVE.png" alt="A cybercriminal interacts with a ghostly, AI-driven honeypot system">
+
 DECEIVE, the **DECeption with Evaluative Integrated Validation Engine**, is a high-interaction, low-effort honeypot system. Unlike most high-interaction honeypots, DECEIVE doesn't provide attackers with access to any actual system. AI actually does all the work of simulating a realistic honeypot system based on a configurable system prompt that describes what type of system you want to simulate. Unlike many other high-interaction honeypots which require substantial effort to seed with realistic users, data, and applications, DECEIVE's AI backend will do all this for you, automatically.
 
 This version of DECEIVE simulates a Linux server via the SSH protocol. It will log all the user inputs, the outputs returned by the LLM backend, as well as a summary of each session after they end. It'll even tell you if it thinks a users' session was benign, suspicious, or outright malicious.
@@ -91,8 +93,8 @@ Things to note:
     * `Session summary`
     * `SSH connection closed`
 * Several of these message types also feature a `details` field with additional information
-    * `User input` messages contain a base64-encoded copy of the entire user input
+    * `User input` messages contain a base64-encoded copy of the entire user input in the `details` field, as well as an `interactive` field (true/false) that tells you whether this was an interactive or non-interactive command (i.e., whether they logged in with a terminal session or provided a command on the SSH command-line).
-    * `LLM response` messages contain a base64-encoded copy of the entire simulated response
+    * `LLM response` messages contain a base64-encoded copy of the entire simulated response in the `details` field.
     * `Session summary` messages contain not only a summary of the commands, but also a guess as to what they might have been intended to accomplish. There will also be a `judgement` field that contains one of "BENIGN", "SUSPICIOUS", or "MALICIOUS"
 * Since this is a honeypot and not intended for use by real users, IT WILL LOG USERNAMES AND PASSWORDS! These are found in the `Authentication success` messages, in the `username` and `password` fields.
 

SSH/config.ini.TEMPLATE

@@ -6,6 +6,10 @@
 # The name of the file you wish to write the honeypot log to.
 log_file = ssh_log.log 
 
+# The name of the sensor, used to identify this honeypot in the logs.
+# If you leave this blank, the honeypot will use the system's hostname.
+sensor_name = deceive
+
 # Settings for the SSH honeypot
 [ssh]
 # The port the SSH honeypot will listen on. You will probably want to set 
@@ -14,8 +18,10 @@ port = 8022
 # The host key to use for the SSH server. This should be a private key.
 # See the README for how to generate this key.
 host_priv_key = ssh_host_key
-# The server version string to send to clients.
+# The server version string to send to clients. The SSH server automatically 
-server_version_string = SSH-2.0-OpenSSH_8.2p1 Ubuntu-4ubuntu0.3
+# prepends "SSH-2.0-" to this. So "OpenSSH_8.2p1 Ubuntu-4ubuntu0.3" will 
+# be transformed to "SSH-2.0-OpenSSH_8.2p1 Ubuntu-4ubuntu0.3".
+server_version_string = OpenSSH_8.2p1 Ubuntu-4ubuntu0.3
 
 # Settings to configure which LLM backend to use.  Only one stanza 
 # should be uncommented at a time.
@@ -24,6 +30,17 @@ server_version_string = SSH-2.0-OpenSSH_8.2p1 Ubuntu-4ubuntu0.3
 llm_provider = openai
 model_name = gpt-4o
 
+##### Azure OpenAI
+#llm_provider = azure 
+#azure_deployment = gpt-4o 
+#azure_api_version = 2025-01-01-preview
+#azure_endpoint = <your endpoint url>
+#model_name = gpt-4o
+
+##### ollama llama3
+#llm_provider = ollama
+#model_name = llama3.3
+
 ##### Any model via Amazon Bedrock
 # Valid AWS model names can be found here: https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
 #llm_provider = AWS
@@ -36,6 +53,12 @@ model_name = gpt-4o
 #llm_provider = gemini
 #model_name = gemini-1.5-pro
 
+# Temperature controls randomness in LLM responses. Values usually range from 0.0 to 2.0.
+# Lower values (e.g., 0.2) make responses more focused and deterministic.
+# Higher values (e.g., 0.8) make responses more creative and variable.
+# Default is 0.2.
+temperature = 0.2
+
 # The maximum number of tokens to send to the LLM backend in a single 
 # request. This includes the message history for the session, so should
 # be fairly high. Not all models support large token counts, so be sure 
@@ -51,10 +74,11 @@ system_prompt = Interpret all inputs as though they were SSH commands and provid
     responses do not indicate that the user is on a virtual machine, container, 
     or emulation at all. 
 
-    On the first call, be sure to include a realistic MOTD.
+    If {interactive} is True, make sure to include a realistic MOTD on the
+    first call, and to end all responses with a realistic shell prompt to display to 
+    the user, including a space at the end. 
 
-    End all responses with a realistic shell prompt to display to the user, 
+    If {interactive} is False, do not include a shell prompt or MOTD.
-    including a space at the end.
 
     Include ANSI color codes for the terminal with the output of ls commands 
     (including any flags), or in any other situation where it is appropriate, but 
@@ -74,9 +98,12 @@ system_prompt = Interpret all inputs as though they were SSH commands and provid
 # The valid user accounts and passwords for the SSH server, in the
 # form "username = password".  Note that you can enable login without 
 # a password by leaving that field blank (e.g., "guest =" on a line by
-# itself)
+# itself). You can set an account to accept ANY password, including an empty
+# password, by setting the password to "*"
 [user_accounts]
 guest = 
 user1 = secretpw
 user2 = password123
+root = *
+
 

SSH/ssh_server.py

@@ -1,27 +1,37 @@
 #!/usr/bin/env python3
 
 from configparser import ConfigParser
+import argparse
 import asyncio
 import asyncssh
 import threading
 import sys
 import json
+import os
+import traceback
 from typing import Optional
 import logging
 import datetime
 import uuid
 from base64 import b64encode
 from operator import itemgetter
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, AzureChatOpenAI
 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
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.runnables import RunnablePassthrough
+from asyncssh.misc import ConnectionLost
+import socket
 
 class JSONFormatter(logging.Formatter):
+    def __init__(self, sensor_name, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.sensor_name = sensor_name
+
     def format(self, record):
         log_record = {
             "timestamp": datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).isoformat(sep="T", timespec="milliseconds"),
@@ -31,8 +41,12 @@ class JSONFormatter(logging.Formatter):
             "src_port": record.src_port,
             "dst_ip": record.dst_ip,
             "dst_port": record.dst_port,
-            "message": record.getMessage()
+            "message": record.getMessage(),
+            "sensor_name": self.sensor_name,
+            "sensor_protocol": "ssh"
         }
+        if hasattr(record, 'interactive'):
+            log_record["interactive"] = record.interactive
         # Include any additional fields from the extra dictionary
         for key, value in record.__dict__.items():
             if key not in log_record and key != 'args' and key != 'msg':
@@ -46,8 +60,18 @@ class MySSHServer(asyncssh.SSHServer):
 
     def connection_made(self, conn: asyncssh.SSHServerConnection) -> None:
         # Get the source and destination IPs and ports
-        (src_ip, src_port, _, _) = conn.get_extra_info('peername')
+        peername = conn.get_extra_info('peername')
-        (dst_ip, dst_port, _, _) = conn.get_extra_info('sockname')
+        sockname = conn.get_extra_info('sockname')
+
+        if peername is not None:
+            src_ip, src_port = peername[:2]
+        else:
+            src_ip, src_port = '-', '-'
+
+        if sockname is not None:
+            dst_ip, dst_port = sockname[:2]
+        else:
+            dst_ip, dst_port = '-', '-'
 
         # Store the connection details in thread-local storage
         thread_local.src_ip = src_ip
@@ -61,6 +85,8 @@ class MySSHServer(asyncssh.SSHServer):
     def connection_lost(self, exc: Optional[Exception]) -> None:
         if exc:
             logger.error('SSH connection error', extra={"error": str(exc)})
+            if not isinstance(exc, ConnectionLost):
+                traceback.print_exception(exc)
         else:
             logger.info("SSH connection closed")
         # Ensure session summary is called on connection loss if attributes are set
@@ -87,7 +113,7 @@ class MySSHServer(asyncssh.SSHServer):
     def validate_password(self, username: str, password: str) -> bool:
         pw = accounts.get(username, '*')
         
-        if ((pw != '*') and (password == pw)):
+        if pw == '*' or (pw != '*' and password == pw):
             logger.info("Authentication success", extra={"username": username, "password": password})
             return True
         else:
@@ -134,7 +160,8 @@ representative examples.
     llm_response = await session.ainvoke(
         {
             "messages": [HumanMessage(content=prompt)],
-            "username": process.get_extra_info('username')
+            "username": process.get_extra_info('username'),
+            "interactive": True  # Ensure interactive flag is passed
         },
             config=llm_config
     )
@@ -149,6 +176,7 @@ representative examples.
         judgement = "MALICIOUS"
 
     logger.info("Session summary", extra={"details": llm_response.content, "judgement": judgement})
+
     server.summary_generated = True
 
 async def handle_client(process: asyncssh.SSHServerProcess, server: MySSHServer) -> None:
@@ -162,42 +190,57 @@ async def handle_client(process: asyncssh.SSHServerProcess, server: MySSHServer)
 
     llm_config = {"configurable": {"session_id": task_uuid}}
 
-    llm_response = await with_message_history.ainvoke(
-        {
-            "messages": [HumanMessage(content="ignore this message")],
-            "username": process.get_extra_info('username')
-        },
-            config=llm_config
-    )
-
-    process.stdout.write(f"{llm_response.content}")
-    logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8')})
-
-    # Store process, llm_config, and session in the MySSHServer instance
-    server._process = process
-    server._llm_config = llm_config
-    server._session = with_message_history
-
     try:
-        async for line in process.stdin:
+        if process.command:
-            line = line.rstrip('\n')
+            # Handle non-interactive command execution
-            logger.info("User input", extra={"details": b64encode(line.encode('utf-8')).decode('utf-8')})
+            command = process.command
-
+            logger.info("User input", extra={"details": b64encode(command.encode('utf-8')).decode('utf-8'), "interactive": False})
-            # Send the command to the LLM and give the response to the user
             llm_response = await with_message_history.ainvoke(
                 {
-                    "messages": [HumanMessage(content=line)],
+                    "messages": [HumanMessage(content=command)],
-                    "username": process.get_extra_info('username')
+                    "username": process.get_extra_info('username'),
+                    "interactive": False
                 },
                     config=llm_config
             )
-            if llm_response.content == "XXX-END-OF-SESSION-XXX":
+            process.stdout.write(f"{llm_response.content}")
-                await session_summary(process, llm_config, with_message_history, server)
+            logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": False})
-                process.exit(0)
+            await session_summary(process, llm_config, with_message_history, server)
-                return
+            process.exit(0)
-            else:
+        else:
-                process.stdout.write(f"{llm_response.content}")
+            # Handle interactive session
-                logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8')})
+            llm_response = await with_message_history.ainvoke(
+                {
+                    "messages": [HumanMessage(content="ignore this message")],
+                    "username": process.get_extra_info('username'),
+                    "interactive": True
+                },
+                    config=llm_config
+            )
+
+            process.stdout.write(f"{llm_response.content}")
+            logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": True})
+
+            async for line in process.stdin:
+                line = line.rstrip('\n')
+                logger.info("User input", extra={"details": b64encode(line.encode('utf-8')).decode('utf-8'), "interactive": True})
+
+                # Send the command to the LLM and give the response to the user
+                llm_response = await with_message_history.ainvoke(
+                    {
+                        "messages": [HumanMessage(content=line)],
+                        "username": process.get_extra_info('username'),
+                        "interactive": True
+                    },
+                        config=llm_config
+                )
+                if llm_response.content == "YYY-END-OF-SESSION-YYY":
+                    await session_summary(process, llm_config, with_message_history, server)
+                    process.exit(0)
+                    return
+                else:
+                    process.stdout.write(f"{llm_response.content}")
+                    logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": True})
 
     except asyncssh.BreakReceived:
         pass
@@ -235,7 +278,7 @@ class ContextFilter(logging.Filter):
         if task:
             task_name = task.get_name()
         else:
-            task_name = "-"
+            task_name = thread_local.__dict__.get('session_id', '-')
 
         record.src_ip = thread_local.__dict__.get('src_ip', '-')
         record.src_port = thread_local.__dict__.get('src_port', '-')   
@@ -262,33 +305,67 @@ def get_user_accounts() -> dict:
 
     return accounts
 
-def choose_llm():
+def choose_llm(llm_provider: Optional[str] = None, model_name: Optional[str] = None):
-    llm_provider_name = config['llm'].get("llm_provider", "openai")
+    llm_provider_name = llm_provider or config['llm'].get("llm_provider", "openai")
     llm_provider_name = llm_provider_name.lower()
-    model_name = config['llm'].get("model_name", "gpt-3.5-turbo")
+    model_name = model_name or config['llm'].get("model_name", "gpt-4o-mini")
+    
+    # Get temperature parameter from config, default to 0.2 if not specified
+    temperature = config['llm'].getfloat("temperature", 0.2)
 
     if llm_provider_name == 'openai':
         llm_model = ChatOpenAI(
-            model=model_name
+            model=model_name,
+            temperature=temperature
+        )
+    elif llm_provider_name == 'azure':
+        llm_model = AzureChatOpenAI(
+            azure_deployment=config['llm'].get("azure_deployment"),
+            azure_endpoint=config['llm'].get("azure_endpoint"),
+            api_version=config['llm'].get("azure_api_version"),
+            model=config['llm'].get("model_name"),  # Ensure model_name is passed here
+            temperature=temperature
+        )
+    elif llm_provider_name == 'ollama':
+        llm_model = ChatOllama(
+            model=model_name,
+            temperature=temperature
         )
     elif llm_provider_name == 'aws':
         llm_model = ChatBedrockConverse(
             model=model_name,
             region_name=config['llm'].get("aws_region", "us-east-1"),
-            credentials_profile_name=config['llm'].get("aws_credentials_profile", "default")        )
+            credentials_profile_name=config['llm'].get("aws_credentials_profile", "default"),
+            temperature=temperature
+        )
     elif llm_provider_name == 'gemini':
         llm_model = ChatGoogleGenerativeAI(
             model=model_name,
+            temperature=temperature
         )
     else:
         raise ValueError(f"Invalid LLM provider {llm_provider_name}.")
 
     return llm_model
 
-def get_prompts() -> dict:
+def get_prompts(prompt: Optional[str], prompt_file: Optional[str]) -> dict:
     system_prompt = config['llm']['system_prompt']
-    with open("prompt.txt", "r") as f:
+    if prompt is not None:
-        user_prompt = f.read()
+        if not prompt.strip():
+            print("Error: The prompt text cannot be empty.", file=sys.stderr)
+            sys.exit(1)
+        user_prompt = prompt
+    elif prompt_file:
+        if not os.path.exists(prompt_file):
+            print(f"Error: The specified prompt file '{prompt_file}' does not exist.", file=sys.stderr)
+            sys.exit(1)
+        with open(prompt_file, "r") as f:
+            user_prompt = f.read()
+    elif os.path.exists("prompt.txt"):
+        with open("prompt.txt", "r") as f:
+            user_prompt = f.read()
+    else:
+        raise ValueError("Either prompt or prompt_file must be provided.")
     return {
         "system_prompt": system_prompt,
         "user_prompt": user_prompt
@@ -296,78 +373,153 @@ def get_prompts() -> dict:
 
 #### MAIN ####
 
-# Always use UTC for logging
+try:
-logging.Formatter.formatTime = (lambda self, record, datefmt=None: datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).isoformat(sep="T",timespec="milliseconds"))
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description='Start the SSH honeypot server.')
+    parser.add_argument('-c', '--config', type=str, default=None, help='Path to the configuration file')
+    parser.add_argument('-p', '--prompt', type=str, help='The entire text of the prompt')
+    parser.add_argument('-f', '--prompt-file', type=str, default='prompt.txt', help='Path to the prompt file')
+    parser.add_argument('-l', '--llm-provider', type=str, help='The LLM provider to use')
+    parser.add_argument('-m', '--model-name', type=str, help='The model name to use')
+    parser.add_argument('-t', '--trimmer-max-tokens', type=int, help='The maximum number of tokens to send to the LLM backend in a single request')
+    parser.add_argument('-s', '--system-prompt', type=str, help='System prompt for the LLM')
+    parser.add_argument('-r', '--temperature', type=float, help='Temperature parameter for controlling randomness in LLM responses (0.0-2.0)')
+    parser.add_argument('-P', '--port', type=int, help='The port the SSH honeypot will listen on')
+    parser.add_argument('-k', '--host-priv-key', type=str, help='The host key to use for the SSH server')
+    parser.add_argument('-v', '--server-version-string', type=str, help='The server version string to send to clients')
+    parser.add_argument('-L', '--log-file', type=str, help='The name of the file you wish to write the honeypot log to')
+    parser.add_argument('-S', '--sensor-name', type=str, help='The name of the sensor, used to identify this honeypot in the logs')
+    parser.add_argument('-u', '--user-account', action='append', help='User account in the form username=password. Can be repeated.')
+    args = parser.parse_args()
 
-# Read our configuration file
+    # Determine which config file to load
-config = ConfigParser()
+    config = ConfigParser()
-config.read("config.ini")
+    if args.config is not None:
+        # User explicitly set a config file; error if it doesn't exist.
+        if not os.path.exists(args.config):
+            print(f"Error: The specified config file '{args.config}' does not exist.", file=sys.stderr)
+            sys.exit(1)
+        config.read(args.config)
+    else:
+        default_config = "config.ini"
+        if os.path.exists(default_config):
+            config.read(default_config)
+        else:
+            # Use defaults when no config file found.
+            config['honeypot'] = {'log_file': 'ssh_log.log', 'sensor_name': socket.gethostname()}
+            config['ssh'] = {'port': '8022', 'host_priv_key': 'ssh_host_key', 'server_version_string': 'SSH-2.0-OpenSSH_8.2p1 Ubuntu-4ubuntu0.3'}
+            config['llm'] = {'llm_provider': 'openai', 'model_name': 'gpt-3.5-turbo', 'trimmer_max_tokens': '64000', 'temperature': '0.7', 'system_prompt': ''}
+            config['user_accounts'] = {}
 
-# Read the user accounts from the configuration file
+    # Override config values with command line arguments if provided
-accounts = get_user_accounts()
+    if args.llm_provider:
+        config['llm']['llm_provider'] = args.llm_provider
+    if args.model_name:
+        config['llm']['model_name'] = args.model_name
+    if args.trimmer_max_tokens:
+        config['llm']['trimmer_max_tokens'] = str(args.trimmer_max_tokens)
+    if args.system_prompt:
+        config['llm']['system_prompt'] = args.system_prompt
+    if args.temperature is not None:
+        config['llm']['temperature'] = str(args.temperature)
+    if args.port:
+        config['ssh']['port'] = str(args.port)
+    if args.host_priv_key:
+        config['ssh']['host_priv_key'] = args.host_priv_key
+    if args.server_version_string:
+        config['ssh']['server_version_string'] = args.server_version_string
+    if args.log_file:
+        config['honeypot']['log_file'] = args.log_file
+    if args.sensor_name:
+        config['honeypot']['sensor_name'] = args.sensor_name
 
-# Set up the honeypot logger
+    # Merge command-line user accounts into the config
-logger = logging.getLogger(__name__)  
+    if args.user_account:
-logger.setLevel(logging.INFO)  
+        if 'user_accounts' not in config:
+            config['user_accounts'] = {}
+        for account in args.user_account:
+            if '=' in account:
+                key, value = account.split('=', 1)
+                config['user_accounts'][key.strip()] = value.strip()
+            else:
+                config['user_accounts'][account.strip()] = ''
 
-log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
+    # Read the user accounts from the configuration
-logger.addHandler(log_file_handler)
+    accounts = get_user_accounts()
 
-log_file_handler.setFormatter(JSONFormatter())
+    # Always use UTC for logging
+    logging.Formatter.formatTime = (lambda self, record, datefmt=None: datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc).isoformat(sep="T",timespec="milliseconds"))
 
-f = ContextFilter()
+    # Get the sensor name from the config or use the system's hostname
-logger.addFilter(f)
+    sensor_name = config['honeypot'].get('sensor_name', socket.gethostname())
 
-# Now get access to the LLM
+    # Set up the honeypot logger
+    logger = logging.getLogger(__name__)  
+    logger.setLevel(logging.INFO)  
 
-prompts = get_prompts()
+    log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
-llm_system_prompt = prompts["system_prompt"]
+    logger.addHandler(log_file_handler)
-llm_user_prompt = prompts["user_prompt"]
 
-llm = choose_llm()
+    log_file_handler.setFormatter(JSONFormatter(sensor_name))
 
-llm_sessions = dict()
+    f = ContextFilter()
+    logger.addFilter(f)
 
-llm_trimmer = trim_messages(
+    # Now get access to the LLM
-    max_tokens=config['llm'].getint("trimmer_max_tokens", 64000),
-    strategy="last",
-    token_counter=llm,
-    include_system=True,
-    allow_partial=False,
-    start_on="human",
-)
 
-llm_prompt = ChatPromptTemplate.from_messages(
+    prompts = get_prompts(args.prompt, args.prompt_file)
-    [
+    llm_system_prompt = prompts["system_prompt"]
-        (
+    llm_user_prompt = prompts["user_prompt"]
-            "system",
-            llm_system_prompt
-        ),
-        (
-            "system",
-            llm_user_prompt
-        ),
-        MessagesPlaceholder(variable_name="messages"),
-    ]
-)
 
-llm_chain = (
+    llm = choose_llm(config['llm'].get("llm_provider"), config['llm'].get("model_name"))
-    RunnablePassthrough.assign(messages=itemgetter("messages") | llm_trimmer)
-    | llm_prompt
-    | llm
-)
 
-with_message_history = RunnableWithMessageHistory(
+    llm_sessions = dict()
-    llm_chain, 
-    llm_get_session_history,
-    input_messages_key="messages"
-)
-# Thread-local storage for connection details
-thread_local = threading.local()
 
-# Kick off the server!
+    llm_trimmer = trim_messages(
-loop = asyncio.new_event_loop()
+        max_tokens=config['llm'].getint("trimmer_max_tokens", 64000),
-asyncio.set_event_loop(loop)
+        strategy="last",
-loop.run_until_complete(start_server())
+        token_counter=llm,
-loop.run_forever()
+        include_system=True,
+        allow_partial=False,
+        start_on="human",
+    )
+
+    llm_prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                llm_system_prompt
+            ),
+            (
+                "system",
+                llm_user_prompt
+            ),
+            MessagesPlaceholder(variable_name="messages"),
+        ]
+    )
+
+    llm_chain = (
+        RunnablePassthrough.assign(messages=itemgetter("messages") | llm_trimmer)
+        | llm_prompt
+        | llm
+    )
+
+    with_message_history = RunnableWithMessageHistory(
+        llm_chain, 
+        llm_get_session_history,
+        input_messages_key="messages"
+    )
+    # Thread-local storage for connection details
+    thread_local = threading.local()
+
+    # Kick off the server!
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(start_server())
+    loop.run_forever()
+
+except Exception as e:
+    print(f"Error: {e}", file=sys.stderr)
+    traceback.print_exc()
+    sys.exit(1)
 

requirements.txt

@@ -7,7 +7,8 @@ langchain_community
 langchain_openai
 # For Google's Gemini models
 langchain_google_genai
-
+# For Ollama models
+langchain_ollama
 # For AWS
 langchain_aws
 transformers