|
|
|
@ -8,6 +8,7 @@ import threading
|
|
|
|
|
import sys
|
|
|
|
|
import json
|
|
|
|
|
import os
|
|
|
|
|
import traceback
|
|
|
|
|
from typing import Optional
|
|
|
|
|
import logging
|
|
|
|
|
import datetime
|
|
|
|
@ -17,13 +18,20 @@ 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
|
|
|
|
|
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"),
|
|
|
|
@ -33,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':
|
|
|
|
@ -48,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')
|
|
|
|
|
(dst_ip, dst_port, _, _) = conn.get_extra_info('sockname')
|
|
|
|
|
peername = conn.get_extra_info('peername')
|
|
|
|
|
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
|
|
|
|
@ -63,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
|
|
|
|
@ -89,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:
|
|
|
|
@ -136,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
|
|
|
|
|
)
|
|
|
|
@ -151,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:
|
|
|
|
@ -164,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:
|
|
|
|
|
line = line.rstrip('\n')
|
|
|
|
|
logger.info("User input", extra={"details": b64encode(line.encode('utf-8')).decode('utf-8')})
|
|
|
|
|
|
|
|
|
|
# Send the command to the LLM and give the response to the user
|
|
|
|
|
if process.command:
|
|
|
|
|
# Handle non-interactive command execution
|
|
|
|
|
command = process.command
|
|
|
|
|
logger.info("User input", extra={"details": b64encode(command.encode('utf-8')).decode('utf-8'), "interactive": False})
|
|
|
|
|
llm_response = await with_message_history.ainvoke(
|
|
|
|
|
{
|
|
|
|
|
"messages": [HumanMessage(content=line)],
|
|
|
|
|
"username": process.get_extra_info('username')
|
|
|
|
|
"messages": [HumanMessage(content=command)],
|
|
|
|
|
"username": process.get_extra_info('username'),
|
|
|
|
|
"interactive": False
|
|
|
|
|
},
|
|
|
|
|
config=llm_config
|
|
|
|
|
)
|
|
|
|
|
if llm_response.content == "XXX-END-OF-SESSION-XXX":
|
|
|
|
|
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')})
|
|
|
|
|
process.stdout.write(f"{llm_response.content}")
|
|
|
|
|
logger.info("LLM response", extra={"details": b64encode(llm_response.content.encode('utf-8')).decode('utf-8'), "interactive": False})
|
|
|
|
|
await session_summary(process, llm_config, with_message_history, server)
|
|
|
|
|
process.exit(0)
|
|
|
|
|
else:
|
|
|
|
|
# Handle interactive session
|
|
|
|
|
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 == "XXX-END-OF-SESSION-XXX":
|
|
|
|
|
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
|
|
|
|
@ -237,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', '-')
|
|
|
|
@ -245,7 +286,7 @@ class ContextFilter(logging.Filter):
|
|
|
|
|
record.dst_port = thread_local.__dict__.get('dst_port', '-')
|
|
|
|
|
|
|
|
|
|
record.task_name = task_name
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return True
|
|
|
|
|
|
|
|
|
|
def llm_get_session_history(session_id: str) -> BaseChatMessageHistory:
|
|
|
|
@ -264,15 +305,19 @@ def get_user_accounts() -> dict:
|
|
|
|
|
|
|
|
|
|
return accounts
|
|
|
|
|
|
|
|
|
|
def choose_llm():
|
|
|
|
|
llm_provider_name = config['llm'].get("llm_provider", "openai")
|
|
|
|
|
def choose_llm(llm_provider: Optional[str] = None, model_name: Optional[str] = None):
|
|
|
|
|
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-3.5-turbo")
|
|
|
|
|
|
|
|
|
|
if llm_provider_name == 'openai':
|
|
|
|
|
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,
|
|
|
|
@ -312,90 +357,150 @@ def get_prompts(prompt: Optional[str], prompt_file: Optional[str]) -> dict:
|
|
|
|
|
|
|
|
|
|
#### MAIN ####
|
|
|
|
|
|
|
|
|
|
# Parse command line arguments
|
|
|
|
|
parser = argparse.ArgumentParser(description='Start the SSH honeypot server.')
|
|
|
|
|
parser.add_argument('-c', '--config', type=str, default='config.ini', 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')
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
try:
|
|
|
|
|
# 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('-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()
|
|
|
|
|
|
|
|
|
|
# Check if the config file exists
|
|
|
|
|
if not os.path.exists(args.config):
|
|
|
|
|
print(f"Error: The specified config file '{args.config}' does not exist.", file=sys.stderr)
|
|
|
|
|
# Determine which config file to load
|
|
|
|
|
config = ConfigParser()
|
|
|
|
|
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', 'system_prompt': ''}
|
|
|
|
|
config['user_accounts'] = {}
|
|
|
|
|
|
|
|
|
|
# Override config values with command line arguments if provided
|
|
|
|
|
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.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
|
|
|
|
|
|
|
|
|
|
# Merge command-line user accounts into the config
|
|
|
|
|
if args.user_account:
|
|
|
|
|
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()] = ''
|
|
|
|
|
|
|
|
|
|
# Read the user accounts from the configuration
|
|
|
|
|
accounts = get_user_accounts()
|
|
|
|
|
|
|
|
|
|
# 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"))
|
|
|
|
|
|
|
|
|
|
# Get the sensor name from the config or use the system's hostname
|
|
|
|
|
sensor_name = config['honeypot'].get('sensor_name', socket.gethostname())
|
|
|
|
|
|
|
|
|
|
# Set up the honeypot logger
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
|
|
|
|
log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
|
|
|
|
|
logger.addHandler(log_file_handler)
|
|
|
|
|
|
|
|
|
|
log_file_handler.setFormatter(JSONFormatter(sensor_name))
|
|
|
|
|
|
|
|
|
|
f = ContextFilter()
|
|
|
|
|
logger.addFilter(f)
|
|
|
|
|
|
|
|
|
|
# Now get access to the LLM
|
|
|
|
|
|
|
|
|
|
prompts = get_prompts(args.prompt, args.prompt_file)
|
|
|
|
|
llm_system_prompt = prompts["system_prompt"]
|
|
|
|
|
llm_user_prompt = prompts["user_prompt"]
|
|
|
|
|
|
|
|
|
|
llm = choose_llm(config['llm'].get("llm_provider"), config['llm'].get("model_name"))
|
|
|
|
|
|
|
|
|
|
llm_sessions = dict()
|
|
|
|
|
|
|
|
|
|
llm_trimmer = trim_messages(
|
|
|
|
|
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(
|
|
|
|
|
[
|
|
|
|
|
(
|
|
|
|
|
"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)
|
|
|
|
|
|
|
|
|
|
# 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"))
|
|
|
|
|
|
|
|
|
|
# Read our configuration file
|
|
|
|
|
config = ConfigParser()
|
|
|
|
|
config.read(args.config)
|
|
|
|
|
|
|
|
|
|
# Read the user accounts from the configuration file
|
|
|
|
|
accounts = get_user_accounts()
|
|
|
|
|
|
|
|
|
|
# Set up the honeypot logger
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
|
|
|
|
log_file_handler = logging.FileHandler(config['honeypot'].get("log_file", "ssh_log.log"))
|
|
|
|
|
logger.addHandler(log_file_handler)
|
|
|
|
|
|
|
|
|
|
log_file_handler.setFormatter(JSONFormatter())
|
|
|
|
|
|
|
|
|
|
f = ContextFilter()
|
|
|
|
|
logger.addFilter(f)
|
|
|
|
|
|
|
|
|
|
# Now get access to the LLM
|
|
|
|
|
|
|
|
|
|
prompts = get_prompts(args.prompt, args.prompt_file)
|
|
|
|
|
llm_system_prompt = prompts["system_prompt"]
|
|
|
|
|
llm_user_prompt = prompts["user_prompt"]
|
|
|
|
|
|
|
|
|
|
llm = choose_llm()
|
|
|
|
|
|
|
|
|
|
llm_sessions = dict()
|
|
|
|
|
|
|
|
|
|
llm_trimmer = trim_messages(
|
|
|
|
|
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(
|
|
|
|
|
[
|
|
|
|
|
(
|
|
|
|
|
"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()
|
|
|
|
|
|
|
|
|
|