Add initial implementation of OpenAI assistant within a tmux session

metacoma · Oct 28, 2023 · 2cc0d45 · 2cc0d45
1 parent f740af3
commit 2cc0d45
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 98 deletions.
diff --git a/files/ai-helper/ai-helper.py b/files/ai-helper/ai-helper.py
@@ -1,18 +1,14 @@
 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"
 neo4j_url = "bolt://192.168.49.2:31237"
 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
+    # <user>@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
@@ -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
diff --git 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<<OPENAI_EOF
+#*************
+#{chat_message}
+#*************
+#OPENAI_EOF
+#"""
+#      }
+
       tmux_cmd = {
-        "tmux_pane_id": "13",
-        "session_name": "ai-helper",
-        "cmd": f"""cat<<OPENAI_EOF
-*************
-{chat_message}
-*************
-OPENAI_EOF
-"""
+        "tmux_pane_id": tmux_pane_id,
+        "session_name": tmux_session_name,
+        "cmd": f"""echo{chat_message}"""
       }
       rabbitmq_channel.basic_publish(exchange='tmux', routing_key='tmux', body=json.dumps(tmux_cmd))
       break

diff --git a/inventory/targets/vector.yml b/inventory/targets/vector.yml
@@ -275,83 +275,13 @@ parameters:
 #
 
         sinks:
-#          kafka-tmux-pane-bytestream:
-#            type: kafka
-#            inputs:
-#              - tmux-pane-bytestream
-#            bootstrap_servers: "${kafka_cluster_name}-kafka-external-bootstrap.${kafka_k8s_namespace}:${kafka_port}"
-#            topic: "${kafka_topic_tmux_bytestream}"
-#            compression: "none"
-#            encoding:
-#              codec: json
           console:
             type: console
             inputs:
               - http_server
-#              - io-context
-#              # - tmux-pane-bytestream
-#              # - tmux-stream-split-line
-#              # - kafka-linestream-to-elasticsearch
-#              # - kafka-tmux-pane-linestream
-#              # - kafka-io-context-to-elasticsearch
+              - tmux-pane-raw-io-context
             encoding:
               codec: json
-#          kafka-tmux-pane-raw-linestream:
-#            type: kafka
-#            inputs:
-#              -  tmux-stream-split-line
-#            bootstrap_servers: "${kafka_cluster_name}-kafka-external-bootstrap.${kafka_k8s_namespace}:${kafka_port}"
-#            # topic: "${kafka_topic_tmux_linestream}"
-#            topic: "${kafka_topic_tmux_raw_linestream}" # linestream with shell color codes
-#            compression: "none"
-#            encoding:
-#              codec: json
-#          opensearch-raw-tmux-lines:
-#            type: elasticsearch
-#            inputs:
-#              - tmux-stream-split-line
-#            endpoint: https://${opensearch_helm_release}-cluster-master.${opensearch_k8s_namespace}:9200
-#            api_version: "v7" # https://opensearch.org/faq/
-#            auth: # TODO (@metacoma) use opensearch credentials from opensearch class
-#              user: ${opensearch_vector_user}
-#              password: ${opensearch_vector_password}
-#              strategy: basic
-#            tls:
-#              verify_certificate: false
-#              verify_hostname: false
-#            healthcheck:
-#              enabled: false
-#            bulk:
-#              index: "${opensearch_raw_tmux_lines_index}-%F"
-
-#          opensearch-raw-tmux-io-context:
-#            type: elasticsearch
-#            inputs:
-#              - tmux-pane-raw-io-context
-#            endpoint: https://${opensearch_helm_release}-cluster-master.${opensearch_k8s_namespace}:9200
-#            api_version: "v7" # https://opensearch.org/faq/
-#            auth: # TODO (@metacoma) use opensearch credentials from opensearch class
-#              user: ${opensearch_vector_user}
-#              password: ${opensearch_vector_password}
-#              strategy: basic
-#            tls:
-#              verify_certificate: false
-#              verify_hostname: false
-#            healthcheck:
-#              enabled: false
-#            bulk:
-#              index: "${opensearch_raw_tmux_io_context_index}-%F"
-
-#          kafka-tmux-pane-raw-io-context:
-#            type: kafka
-#            inputs:
-#              - tmux-pane-raw-io-context
-#              # - tmux-pane-io-context
-#            bootstrap_servers: "${kafka_cluster_name}-kafka-external-bootstrap.${kafka_k8s_namespace}:${kafka_port}"
-#            topic: "${kafka_topic_tmux_raw_io_context}"
-#            compression: "none"
-#            encoding:
-#              codec: json
 
           rabbitmq-tmux-pane-raw-io-context:
             type: amqp
@@ -363,7 +293,7 @@ parameters:
             encoding:
               codec: json
 
-          rabbitmq-tmux-pane-raw-io-context:
+          rabbitmq-tmux-pane-raw-io-line:
             type: amqp
             inputs:
               - tmux-stream-split-line