From 2cc0d4522232e27b669acf1c789221c8238fcea6 Mon Sep 17 00:00:00 2001 From: Ryabin Sergey Date: Sat, 28 Oct 2023 18:52:19 +0200 Subject: [PATCH] Add initial implementation of OpenAI assistant within a tmux session --- files/ai-helper/ai-helper.py | 55 ++++++++++---- .../freeplane-io-context.py | 3 +- .../freeplane_grpc/python/freeplane_pb2.py | 1 - files/openai-io-context/openai-io-context.py | 34 ++++++--- inventory/targets/vector.yml | 74 +------------------ 5 files changed, 69 insertions(+), 98 deletions(-) diff --git a/files/ai-helper/ai-helper.py b/files/ai-helper/ai-helper.py index 959dff2..ee8af20 100644 --- a/files/ai-helper/ai-helper.py +++ b/files/ai-helper/ai-helper.py @@ -1,9 +1,7 @@ import pika import pprint import json -import grpc import re -import libtmux from neo4j import GraphDatabase rabbitmq_url = "amqp://user:password@192.168.49.2:30466/%2f" @@ -11,8 +9,6 @@ neo4j_username = 'neo4j' neo4j_password = 'password' -FP_DETAILS_MAX_LINE = 4 - # Create a connection to RabbitMQ connection = pika.BlockingConnection(pika.URLParameters(rabbitmq_url)) rabbitmq_channel = connection.channel() @@ -23,18 +19,51 @@ rabbitmq_channel.queue_bind(exchange=exchange_name, queue=queue_name) +def get_last_io_document(host, tmux_session_name, tmux_pane_id, n): + query = """ +MATCH (host:Node {node_type: "host", name: "%s"})-[:HAS_TMUX_SESSION]->(session:Node {node_type: "tmux_session", name: "%s"})-[:HAS_PANE]->(pane:Node {node_type: "tmux_pane", name: "%d"})-[:HAS_IO_CONTEXT]->(io_context:Node {node_type: "io_context"})-[:IO_DOCUMENT]->(io_document:Node {node_type: "io_document"}) +WITH io_document +ORDER BY ID(io_document) DESC +LIMIT %d +RETURN io_document +""" % (host, tmux_session_name, int(tmux_pane_id), int(n)) + print(query) + with GraphDatabase.driver(neo4j_url, auth=(neo4j_username, neo4j_password)) as driver: + with driver.session() as session: + result = session.run(query) + record = result.single() + return record + + return None + + def callback(ch, method, properties, body): - pprint.pprint(body.decode()) - #data = json.loads(body.decode()) - #pprint.pprint(data) - #session_name = data['session_name'] - #tmux_pane_id = data['tmux_pane_id'] - #cmd = data['cmd'] + data = json.loads(body.decode()) - #server = libtmux.Server() + # trigger pattern + # @host:/home/user $ # ai request + pattern = r".*@.*\$ #.*" + if re.match(pattern, data['message']): + host = data['host'] + tmux_session_name = data['metadata']['tmux']['session_name'] + tmux_pane_id = data['metadata']['tmux']['pane_id'] + last_document = get_last_io_document(host, tmux_session_name, tmux_pane_id, 1) + cmd_input = last_document['io_document'].get('name') + cmd_output = last_document['io_document'].get('output_data') - #session = server.find_where({"session_name": session_name}) - #session.panes.filter(pane_id = "%" + tmux_pane_id)[0].send_keys(cmd) + question = re.sub(r'^.*?#[ ]?', '', data['message']) + data = { + "tmux_pane_id": tmux_pane_id, + "tmux_session_name": tmux_session_name, + "host": host, + "question": question, + "terminal-command-sequence": [{ + "input": cmd_input, + "output": cmd_output + }] + } + rabbitmq_channel.basic_publish(exchange='openai', routing_key='openai', body=json.dumps(data)) + pprint.pprint(data) rabbitmq_channel.basic_consume(queue=queue_name, on_message_callback=callback, auto_ack=True) print("Waiting for messages. To exit, press Ctrl+C") diff --git a/files/freeplane-io-context/freeplane-io-context.py b/files/freeplane-io-context/freeplane-io-context.py index 8301494..88dad08 100644 --- a/files/freeplane-io-context/freeplane-io-context.py +++ b/files/freeplane-io-context/freeplane-io-context.py @@ -1,4 +1,3 @@ -{%- raw -%} import pika import pprint import json @@ -29,6 +28,7 @@ output_queue_name = "io-document" exchange_name = os.getenv('EXCHANGE_NAME') +exchange_name = "io-context" result = rabbitmq_channel.queue_declare(queue="", exclusive=True) queue_name = result.method.queue @@ -473,4 +473,3 @@ def callback(ch, method, properties, body): print("Waiting for messages. To exit, press Ctrl+C") rabbitmq_channel.start_consuming() -{%- endraw -%} diff --git a/files/grpc/freeplane_grpc/python/freeplane_pb2.py b/files/grpc/freeplane_grpc/python/freeplane_pb2.py index a0d4244..5ddd7d3 100644 --- a/files/grpc/freeplane_grpc/python/freeplane_pb2.py +++ b/files/grpc/freeplane_grpc/python/freeplane_pb2.py @@ -19,7 +19,6 @@ _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'freeplane_pb2', _globals) if _descriptor._USE_C_DESCRIPTORS == False: - DESCRIPTOR._options = None DESCRIPTOR._serialized_options = b'\n\031org.freeplane.plugin.grpcB\tfreeplaneP\001\242\002\002FP' _globals['_CREATECHILDREQUEST']._serialized_start=30 diff --git a/files/openai-io-context/openai-io-context.py b/files/openai-io-context/openai-io-context.py index cfcf81e..2d35512 100644 --- a/files/openai-io-context/openai-io-context.py +++ b/files/openai-io-context/openai-io-context.py @@ -35,28 +35,42 @@ def callback(ch, method, properties, body): """} ] pprint.pprint(body.decode()) + data = json.loads(body.decode()) - data["question"] = "What is the uptime?" - messages.append({"role": "user", "content": json.dumps(data)}) + + openai_payload = { + "question": data['question'], + "terminal-command-sequence": data['terminal-command-sequence'] + } + + messages.append({"role": "user", "content": json.dumps(openai_payload)}) response = openai.ChatCompletion.create( model="gpt-3.5-turbo", messages=messages ) + tmux_pane_id = data['tmux_pane_id'] + tmux_session_name = data['tmux_session_name'] # Print the response and add it to the messages list for i in range(len(response['choices'])): chat_message = response['choices'][i]['message']['content'] print(f"Bot: #{i} {chat_message}") +# tmux_cmd = { +# "tmux_pane_id": "{tmux_pane_id}", +# "session_name": "{tmux_session_name}", +# "cmd": f"""cat<