tensorwatch-api 0.1__py3-none-any.whl

tensorwatch_api-0.1.dist-info/METADATA

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: tensorwatch-api
+Version: 0.1
+Requires-Dist: pykafka
+Requires-Dist: tensorwatchext
+Requires-Dist: ipywidgets
+Requires-Dist: ipympl
+Dynamic: requires-dist

tensorwatch_api-0.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+twapi/Example_Senter.py,sha256=jPKi7f9teidzXKq-yXMW0sa2yzYGCob3WIHVi3J1uNA,1532
+twapi/__init__.py,sha256=wYkftb8fUjtf62Qs3yXegUNkMn4bSTaSWLyRCvv4yt0,34
+twapi/kafka_connector.py,sha256=RUcK9MYnfaMDO47QQC9-_SyqahRZmmoBtBwCSetAVFs,10829
+twapi/pykafka_connector.py,sha256=7e5fHEZyqJo1a43mHTpqmscbQJZb7CjZhpN8kyRw-ok,12977
+twapi/twapi.py,sha256=NkfcTz6u-YuJyCc9NsjwYuH6A2Pzkhn7lQVxZf0aXJQ,9158
+tensorwatch_api-0.1.dist-info/METADATA,sha256=vdsLmWDb-RRbGvXv7fdc1ajqJoUQwFTydttNhlMqdg4,189
+tensorwatch_api-0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tensorwatch_api-0.1.dist-info/entry_points.txt,sha256=xIxCPp_fRVck1fVJE1kqhxAEVEycrxOk8TM48EXeHWs,59
+tensorwatch_api-0.1.dist-info/top_level.txt,sha256=E48dj13nSJGzHKjYnAoCGtiUjM9AU2LZk3vvmdWtNhM,6
+tensorwatch_api-0.1.dist-info/RECORD,,

tensorwatch_api-0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tensorwatch_api-0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+twapi-senter = twapi.Example_Senter:main

tensorwatch_api-0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+twapi

twapi/Example_Senter.py

@@ -0,0 +1,51 @@
+from pykafka import KafkaClient
+import json
+import time
+import random
+import sys
+
+def main():
+    """
+    Benchmark producer to test throughput and latency.
+    """
+    try:
+        client = KafkaClient(hosts="127.0.0.1:9093")
+        topic = client.topics['gemini2']
+    except Exception as e:
+        print(f"Failed to connect to Kafka: {e}")
+        sys.exit(1)
+
+    parse_type = "json"
+    num_messages = 200000  # Number of messages to send
+
+    with topic.get_sync_producer() as producer:
+        print("Starting benchmark...")
+        print(f"Sending {num_messages} messages to topic '{topic.name.decode()}'...")
+
+        start_time = time.time()
+
+        for i in range(num_messages):
+            message = {
+                'seq': i,
+                'send_time': time.time(),
+                'data': random.randint(0, 1000)
+            }
+            
+            if parse_type == "json":
+                data = json.dumps(message)
+                producer.produce(data.encode('utf-8'))
+            
+            if (i + 1) % 1000 == 0:
+                print(f"Sent {i + 1}/{num_messages} messages...")
+
+        end_time = time.time()
+        duration = end_time - start_time
+        throughput = num_messages / duration if duration > 0 else float('inf')
+
+        print("\n--- BENCHMARK SUMMARY ---")
+        print(f"Sent {num_messages} messages in {duration:.2f} seconds.")
+        print(f"Producer throughput: {throughput:.2f} messages/sec.")
+        print("------------------------")
+
+if __name__ == "__main__":
+    main()

twapi/__init__.py

@@ -0,0 +1 @@
+from .tensorwatch-api import twapi

twapi/kafka_connector.py

@@ -0,0 +1,255 @@
+import time
+import threading
+from queue import Queue
+import random
+import json
+import pickle
+from typing import Dict
+from confluent_kafka import Consumer
+from tensorwatch import Watcher
+from probables import CountMinSketch
+import logging
+
+# Optional Parsers
+try:
+    import xmltodict
+except ImportError:
+    xmltodict = None
+
+try:
+    import avro.schema
+    from avro.io import DatumReader, BinaryDecoder
+    import io
+except ImportError:
+    avro = None
+
+try:
+    from protobuf_to_dict import protobuf_to_dict
+    from google.protobuf import message
+except ImportError:
+    protobuf_to_dict = None
+
+class KafkaConnector(threading.Thread):
+    """
+    A Kafka consumer that runs in a separate thread to consume messages from a Kafka topic.
+    It supports various message formats and integrates with TensorWatch for real-time data visualization.
+    """
+    def __init__(self, hosts="localhost:9092", topic=None, parsetype=None, avro_schema=None, queue_length=50000,
+                 cluster_size=1, consumer_config=None, poll=1.0, auto_offset="earliest", group_id="mygroup",
+                 decode="utf-8", schema_path=None, protobuf_message=None, random_sampling=None, countmin_width=None,
+                 countmin_depth=None, twapi_instance=None):
+        """
+        Initializes the KafkaConnector.
+
+        Args:
+            hosts (str): Comma-separated list of Kafka brokers.
+            topic (str): The Kafka topic to consume from.
+            parsetype (str): The format of the messages (e.g., "json", "pickle", "xml", "avro", "protobuf").
+            avro_schema (str): The Avro schema for message deserialization.
+            queue_length (int): The maximum number of messages to store in the internal queue.
+            cluster_size (int): The number of consumer threads to run.
+            consumer_config (dict): A dictionary of Kafka consumer configuration settings.
+            poll (float): The timeout for polling for new messages from Kafka.
+            auto_offset (str): The offset reset policy.
+            group_id (str): The consumer group ID.
+            decode (str): The encoding to use for decoding messages.
+            schema_path (str): The path to the Avro schema file.
+            protobuf_message (str): The name of the Protobuf message class.
+            random_sampling (int): The percentage of messages to sample (0-100).
+            countmin_width (int): The width of the Count-Min Sketch.
+            countmin_depth (int): The depth of the Count-Min Sketch.
+            twapi_instance: An instance of the TensorWatch API for updating metrics.
+        """
+        super().__init__()
+        self.hosts = hosts or "localhost:9092"
+        self.topic = topic
+        self.cluster_size = cluster_size
+        self.decode = decode
+        self.parsetype = parsetype
+        self.protobuf_message = protobuf_message
+        self.queue_length = queue_length
+        self.data = Queue(maxsize=queue_length)
+        self.cms = {}  # Count-Min Sketch table
+        self.countmin_width = countmin_width
+        self.countmin_depth = countmin_depth
+        self.random_sampling = random_sampling
+        self.poll = poll
+        self.consumer_config = consumer_config or {
+            "bootstrap.servers": self.hosts,
+            "group.id": group_id,
+            "auto.offset.reset": auto_offset,
+        }
+        self._quit = threading.Event()
+        self.size = 0
+        self.watcher = Watcher()
+        self.schema = None
+        self.reader = None
+
+        self.twapi_instance = twapi_instance
+        self.latencies = []
+        self.received_count = 0
+        self.last_report_time = time.time()
+        self.first_message_sent = False
+
+        # Load Avro Schema if needed
+        if parsetype == "avro" and avro:
+            try:
+                self.schema = avro.schema.parse(avro_schema)
+                self.reader = DatumReader(self.schema)
+            except Exception as e:
+                logging.error(f"Avro Schema Error: {e}, Avro may not work")
+                print(f"Avro Schema Error: {e}, Avro may not work")
+                return
+
+        # Load Protobuf if needed
+        if parsetype == "protobuf" and protobuf_to_dict:
+            try:
+                import importlib
+                module = importlib.import_module(protobuf_message)
+                self.protobuf_class = getattr(module, protobuf_message)
+
+            except Exception as e:
+                logging.error(f"Protobuf Import Error: {e}")
+                print(f"Protobuf Import Error: {e}")
+                self.protobuf_class = None
+
+        self.start()
+
+    def myparser(self, message):
+        """
+        Parses a message based on the specified format.
+
+        Args:
+            message: The message to parse.
+
+        Returns:
+            The parsed message, or None if parsing fails.
+        """
+        try:
+            if self.parsetype is None or self.parsetype.lower() == "json":
+                return json.loads(message)
+            elif self.parsetype.lower() == "pickle":
+                return pickle.loads(message)
+            elif self.parsetype.lower() == "xml" and xmltodict:
+                return xmltodict.parse(message)["root"]
+            elif self.parsetype.lower() == "protobuf" and protobuf_to_dict:
+                if self.protobuf_class:
+                    dynamic_message = self.protobuf_class()
+                    dynamic_message.ParseFromString(message)
+                    return protobuf_to_dict(dynamic_message)
+            elif self.parsetype.lower() == "avro" and avro:
+                decoder = BinaryDecoder(io.BytesIO(message))
+                return self.reader.read(decoder)
+        except Exception as e:
+            logging.error(f"Parsing Error ({self.parsetype}): {e}")
+            print(f"Parsing Error ({self.parsetype}): {e}")
+        return None
+
+    def process_message(self, msg):
+        """
+        Processes a single message from Kafka. This includes parsing, calculating latency,
+        and adding the message to the data queue.
+        """
+        receive_time = time.time()
+        try:
+            # Apply random sampling if configured
+            if self.random_sampling and self.random_sampling > random.randint(0, 100):
+                return
+            
+            message = msg.value().decode(self.decode)
+            parsed_message = self.myparser(message)
+
+            # Calculate and record latency if send_time is in the message
+            if parsed_message and isinstance(parsed_message, dict) and 'send_time' in parsed_message:
+                self.received_count += 1
+                send_time = parsed_message['send_time']
+                latency = receive_time - send_time
+                self.latencies.append(latency)
+                parsed_message['latency'] = latency
+                parsed_message['receive_time'] = receive_time
+
+            # Add the parsed message to the queue if it's not full
+            if parsed_message and not self.data.full():
+                self.data.put(parsed_message, block=False)
+                # Notify the twapi_instance on the first message
+                if not self.first_message_sent and self.twapi_instance:
+                    logging.info("First message received, enabling apply button.")
+                    self.twapi_instance.enable_apply_button()
+                    self.first_message_sent = True
+            elif self.data.full():
+                logging.warning("Queue is full, dropping message.")
+
+            # Update Count-Min Sketch if configured
+            if isinstance(parsed_message, dict) and self.countmin_width and self.countmin_depth:
+                for key, value in parsed_message.items():
+                    self.cms.setdefault(key, CountMinSketch(width=self.countmin_width, depth=self.countmin_depth))
+                    self.cms[key].add(str(value))
+
+            self.size += 1
+        except Exception as e:
+            logging.error(f"Message Processing Error: {e}, Message: {message}")
+            print(f"Message Processing Error: {e}, Message: {message}")
+
+    def consumer_loop(self):
+        """
+        The main loop for the Kafka consumer. It polls for messages, processes them,
+        and handles errors.
+        """
+        logging.info(f"Starting consumer loop for topic '{self.topic}'")
+        consumer = Consumer(self.consumer_config)
+        consumer.subscribe([self.topic])
+
+        while not self._quit.is_set():
+            msg = consumer.poll(self.poll)
+            if msg and not msg.error():
+                self.process_message(msg)
+            elif msg and msg.error():
+                logging.error(f"Kafka Error: {msg.error()}")
+                print(f"Kafka Error: {msg.error()}")
+
+        consumer.close()
+        logging.info("Consumer loop stopped")
+
+    def run(self):
+        """
+        Starts the consumer threads and the main watcher loop.
+        """
+        logging.info(f"Starting {self.cluster_size} consumer threads")
+        threads = [threading.Thread(target=self.consumer_loop, daemon=True) for _ in range(self.cluster_size)]
+        for thread in threads:
+            thread.start()
+
+        while not self._quit.is_set():
+            # Observe the data queue with TensorWatch
+            if not self.data.empty():
+                self.watcher.observe(data=list(self.data.queue), size=self.size, cms=self.cms)
+            
+            # --- BENCHMARK REPORTING ---
+            current_time = time.time()
+            if current_time - self.last_report_time > 5.0: # Report every 5 seconds
+                if self.latencies:
+                    avg_latency = sum(self.latencies) / len(self.latencies)
+                    max_latency = max(self.latencies)
+                    min_latency = min(self.latencies)
+                    
+                    time_since_last_report = current_time - self.last_report_time
+                    throughput = self.received_count / time_since_last_report if time_since_last_report > 0 else 0
+                    
+                    stats_str = (f"Recv Throughput: {throughput:.2f} msgs/s | "
+                                 f"Send-Recv Latency (ms): "
+                                 f"Avg: {avg_latency*1000:.2f}, "
+                                 f"Min: {min_latency*1000:.2f}, "
+                                 f"Max: {max_latency*1000:.2f}")
+                    logging.info(f"Benchmark stats: {stats_str}")
+                    print(stats_str)
+
+                    # Update the TensorWatch API with the latest metrics
+                    if self.twapi_instance:
+                        self.twapi_instance.update_metrics(stats_str)
+
+                    # Reset stats for the next interval
+                    self.latencies = []
+                    self.received_count = 0
+                self.last_report_time = current_time
+
+            time.sleep(0.4)

twapi/pykafka_connector.py

@@ -0,0 +1,293 @@
+from pykafka import KafkaClient
+from pykafka.common import OffsetType
+from tensorwatch import Watcher
+import threading
+from queue import Queue
+import json
+import pickle
+import time
+import random
+import logging
+import io
+
+# Optional Parsers
+try:
+    import xmltodict
+except ImportError:
+    xmltodict = None
+
+try:
+    import avro.schema
+    from avro.io import DatumReader, BinaryDecoder
+except ImportError:
+    avro = None
+
+try:
+    from protobuf_to_dict import protobuf_to_dict
+except ImportError:
+    protobuf_to_dict = None
+
+from probables import CountMinSketch
+
+class pykafka_connector(threading.Thread):
+    """
+    A Kafka consumer that uses the pykafka library to consume messages from a Kafka topic.
+    It runs in a separate thread and supports various message formats.
+    """
+    def __init__(self, hosts: str = None, topic: str = None, parsetype: str = None, queue_length: int = None, cluster_size: int = 1,
+                 consumer_group: bytes = b'default', auto_offset_reset: OffsetType = OffsetType.EARLIEST,
+                 fetch_message_max_bytes: int = 1024 * 1024, num_consumer_fetchers: int = 1,
+                 auto_commit_enable: bool = False, auto_commit_interval_ms: int = 1000,
+                 queued_max_messages: int = 2000, fetch_min_bytes: int = 1,
+                 consumer_timeout_ms: int = -1, decode: str = "utf-8",
+                 scema_path: str = None, random_sampling: int = None, countmin_width: int = None,
+                 countmin_depth: int = None, twapi_instance=None, parser_extra=None, probuf_message=None, zookeeper_hosts:str='127.0.0.1:2181'):
+        """
+        Initializes the pykafka_connector.
+
+        Args:
+            hosts (str): Comma-separated list of Kafka brokers.
+            topic (str): The Kafka topic to consume from.
+            parsetype (str): The format of the messages (e.g., "json", "pickle", "xml", "avro", "protobuf").
+            queue_length (int): The maximum number of messages to store in the internal queue.
+            cluster_size (int): The number of consumer threads to run.
+            consumer_group (bytes): The consumer group ID.
+            auto_offset_reset (OffsetType): The offset reset policy.
+            fetch_message_max_bytes (int): The maximum size of a message to fetch.
+            num_consumer_fetchers (int): The number of fetcher threads.
+            auto_commit_enable (bool): Whether to enable auto-commit.
+            auto_commit_interval_ms (int): The auto-commit interval in milliseconds.
+            queued_max_messages (int): The maximum number of messages to queue.
+            fetch_min_bytes (int): The minimum number of bytes to fetch.
+            consumer_timeout_ms (int): The consumer timeout in milliseconds.
+            decode (str): The encoding to use for decoding messages.
+            scema_path (str): The path to the Avro or Protobuf schema.
+            random_sampling (int): The percentage of messages to sample (0-100).
+            countmin_width (int): The width of the Count-Min Sketch.
+            countmin_depth (int): The depth of the Count-Min Sketch.
+            twapi_instance: An instance of the TensorWatch API for updating metrics.
+            parser_extra (str): Extra information for the parser (e.g., Avro schema, Protobuf module).
+            probuf_message (str): The name of the Protobuf message class.
+            zookeeper_hosts (str): Comma-separated list of Zookeeper hosts.
+        """
+        super().__init__()
+        self.hosts = hosts or "127.0.0.1:9092"
+        self.topic = topic
+        self.cluster_size = cluster_size
+        self.decode = decode
+        self.parsetype = parsetype
+        self.scema_path = scema_path
+        self.random_sampling = random_sampling
+        self.parser_extra = parser_extra
+        self.probuf_message = probuf_message
+        self.queue_length = queue_length
+        self.data = Queue(maxsize=queue_length or 50000)
+        self._quit = threading.Event()
+        self.size = 0
+        self.watcher = Watcher()
+        self.cms = {}
+        self.countmin_depth = countmin_depth
+        self.countmin_width = countmin_width
+
+        # pykafka specific settings
+        self.consumer_group = consumer_group
+        self.auto_offset_reset = auto_offset_reset
+        self.fetch_message_max_bytes = fetch_message_max_bytes
+        self.num_consumer_fetchers = num_consumer_fetchers
+        self.auto_commit_enable = auto_commit_enable
+        self.auto_commit_interval_ms = auto_commit_interval_ms
+        self.queued_max_messages = queued_max_messages
+        self.fetch_min_bytes = fetch_min_bytes
+        self.consumer_timeout_ms = consumer_timeout_ms
+        self.zookeeper_hosts = zookeeper_hosts
+
+        # twapi integration
+        self.twapi_instance = twapi_instance
+        self.latencies = []
+        self.received_count = 0
+        self.last_report_time = time.time()
+        self.first_message_sent = False
+
+        # Parsers initialization
+        self.reader = None
+        self.mymodule = None
+        if self.parsetype:
+            if self.parsetype.lower() == 'avro' and avro:
+                try:
+                    schema = avro.schema.parse(parser_extra)
+                    self.reader = DatumReader(schema)
+                except Exception as e:
+                    print(f"Avro schema error or avro not installed: {e}")
+            elif self.parsetype.lower() == 'protobuf' and protobuf_to_dict:
+                try:
+                    import sys
+                    import importlib
+                    if scema_path:
+                        sys.path.append(scema_path)
+                    mymodule = importlib.import_module(parser_extra)
+                    method_to_call = getattr(mymodule, probuf_message)
+                    self.mymodule = method_to_call
+                except Exception as e:
+                    print(f"Error importing protobuf: {e}")
+
+        self.start()
+
+    def myparser(self, message):
+        """
+        Parses a message based on the specified format.
+
+        Args:
+            message: The message to parse.
+
+        Returns:
+            The parsed message, or None if parsing fails.
+        """
+        try:
+            if self.parsetype is None or self.parsetype.lower() == 'json':
+                return json.loads(message)
+            elif self.parsetype.lower() == 'pickle':
+                return pickle.loads(message)
+            elif self.parsetype.lower() == 'xml' and xmltodict:
+                return xmltodict.parse(message).get("root")
+            elif self.parsetype.lower() == 'protobuf' and self.mymodule:
+                dynamic_message = self.mymodule()
+                dynamic_message.ParseFromString(message)
+                return protobuf_to_dict(dynamic_message)
+            elif self.parsetype.lower() == 'avro' and self.reader:
+                message_bytes = io.BytesIO(message)
+                decoder = BinaryDecoder(message_bytes)
+                return self.reader.read(decoder)
+        except Exception as e:
+            logging.error(f"Parsing Error ({self.parsetype}): {e}")
+        return None
+
+    def process_message(self, message_bytes):
+        """
+        Processes a single message from Kafka. This includes parsing, calculating latency,
+        and adding the message to the data queue.
+        """
+        receive_time = time.time()
+        try:
+            # Apply random sampling if configured
+            if self.random_sampling and self.random_sampling > random.randint(0, 100):
+                return
+
+            parsed_message = self.myparser(message_bytes)
+            if parsed_message is None:
+                return
+
+            # Calculate and record latency if send_time is in the message
+            if isinstance(parsed_message, dict) and 'send_time' in parsed_message:
+                self.received_count += 1
+                send_time = parsed_message['send_time']
+                latency = receive_time - send_time
+                self.latencies.append(latency)
+                parsed_message['latency'] = latency
+                parsed_message['receive_time'] = receive_time
+
+            # Add the parsed message to the queue if it's not full
+            if not self.data.full():
+                self.data.put(parsed_message, block=False)
+                # Notify the twapi_instance on the first message
+                if not self.first_message_sent and self.twapi_instance:
+                    logging.info("First message received, enabling apply button.")
+                    self.twapi_instance.enable_apply_button()
+                    self.first_message_sent = True
+            else:
+                logging.warning("Queue is full, dropping message.")
+
+            # Update Count-Min Sketch if configured
+            if isinstance(parsed_message, dict) and self.countmin_width and self.countmin_depth:
+                for key, value in parsed_message.items():
+                    self.cms.setdefault(key, CountMinSketch(width=self.countmin_width, depth=self.countmin_depth))
+                    self.cms[key].add(str(value))
+
+            self.size += 1
+        except Exception as e:
+            logging.error(f"Message Processing Error: {e}")
+
+    def consumer_loop(self):
+        """
+        The main loop for the Kafka consumer. It creates a Kafka client and consumes messages
+        from the specified topic.
+        """
+        logging.info(f"Starting pykafka consumer loop for topic '{self.topic}'")
+        client = KafkaClient(hosts=self.hosts)
+        topic = client.topics[self.topic]
+
+        # Use a balanced consumer if cluster_size is greater than 1
+        if self.cluster_size > 1:
+            consumer = topic.get_balanced_consumer(
+                consumer_group=self.consumer_group,
+                auto_commit_enable=self.auto_commit_enable,
+                auto_offset_reset=self.auto_offset_reset,
+                num_consumer_fetchers=self.num_consumer_fetchers,
+                auto_commit_interval_ms=self.auto_commit_interval_ms,
+                queued_max_messages=self.queued_max_messages,
+                fetch_min_bytes=self.fetch_min_bytes,
+                zookeeper_connect=self.zookeeper_hosts
+            )
+        else:
+            consumer = topic.get_simple_consumer(
+                auto_offset_reset=self.auto_offset_reset,
+                consumer_timeout_ms=self.consumer_timeout_ms,
+                fetch_message_max_bytes=self.fetch_message_max_bytes,
+                auto_commit_enable=self.auto_commit_enable,
+                auto_commit_interval_ms=self.auto_commit_interval_ms,
+                queued_max_messages=self.queued_max_messages,
+                fetch_min_bytes=self.fetch_min_bytes
+            )
+
+        for message in consumer:
+            if self._quit.is_set():
+                break
+            if message is not None:
+                self.process_message(message.value)
+        
+        consumer.stop()
+        logging.info("Consumer loop stopped")
+
+    def run(self):
+        """
+        Starts the consumer threads and the main watcher loop.
+        """
+        logging.info(f"Starting {self.cluster_size} pykafka consumer threads")
+        threads = [threading.Thread(target=self.consumer_loop, daemon=True) for _ in range(self.cluster_size)]
+        for thread in threads:
+            thread.start()
+
+        while not self._quit.is_set():
+            # Observe the data queue with TensorWatch
+            if not self.data.empty():
+                self.watcher.observe(data=list(self.data.queue), size=self.size, cms=self.cms)
+
+            # --- BENCHMARK REPORTING ---
+            current_time = time.time()
+            if current_time - self.last_report_time > 5.0: # Report every 5 seconds
+                if self.latencies:
+                    avg_latency = sum(self.latencies) / len(self.latencies)
+                    max_latency = max(self.latencies)
+                    min_latency = min(self.latencies)
+                    time_since_last_report = current_time - self.last_report_time
+                    throughput = self.received_count / time_since_last_report if time_since_last_report > 0 else 0
+                    
+                    stats_str = (f"Recv Throughput: {throughput:.2f} msgs/s | "
+                                 f"Send-Recv Latency (ms): "
+                                 f"Avg: {avg_latency*1000:.2f}, "
+                                 f"Min: {min_latency*1000:.2f}, "
+                                 f"Max: {max_latency*1000:.2f}")
+                    
+                    # Update the TensorWatch API with the latest metrics
+                    if self.twapi_instance:
+                        self.twapi_instance.update_metrics(stats_str)
+
+                    # Reset stats for the next interval
+                    self.latencies = []
+                    self.received_count = 0
+                self.last_report_time = current_time
+
+            time.sleep(0.4)
+
+    def quit(self):
+        """Stops the consumer thread."""
+        self._quit.set()

twapi/twapi.py

@@ -0,0 +1,197 @@
+import tensorwatch as tw
+from . import kafka_connector as kc
+from . import pykafka_connector as pyc
+from IPython.display import display
+from ipywidgets import widgets
+import asyncio
+import time
+import logging
+import matplotlib.pyplot as plt
+
+class twapi:
+    """TensorWatch API Wrapper for Kafka Streaming and Visualization"""
+
+    def __init__(self):
+        """Initializes the twapi class, setting up the UI widgets and event handlers."""
+        self.default_value = 10
+        self.visualizer = None  # Initialize visualizer as None
+        self.client = tw.WatcherClient()
+        self.out = widgets.Output(layout={})
+
+        # Initialize UI widgets
+        self.update_interval = 0.5  # Delay in seconds
+        self.my_slider = widgets.IntSlider(value=self.default_value, min=1, max=100, step=1, description="Window Size:")
+        self.my_slider2 = widgets.IntSlider(value=self.default_value, min=1, max=100, step=1, description="Window Width:")
+        self.datebutton = widgets.Checkbox(value=False, description="Date")
+        self.offsetbutton = widgets.Checkbox(value=False, description="Use Offset")
+        self.dimhistorybutton = widgets.Checkbox(value=True, description="Dim History")
+        self.colorpicker = widgets.ColorPicker(value="blue", description="Pick a Color")
+        
+        self.button_reset = widgets.Button(description="Reset", tooltip="Reset stream settings")        
+        self.button_apply = widgets.Button(description="Please wait", tooltip="Apply changes to the visualization", disabled=True)
+
+        # Group widgets for a cleaner UI
+        left_box = widgets.VBox([self.my_slider, self.my_slider2, self.colorpicker])
+        right_box = widgets.VBox([self.offsetbutton, self.dimhistorybutton, self.datebutton])
+        self.options_box = widgets.HBox([left_box, right_box])
+        self.accordion = widgets.Accordion(children=[self.options_box])
+        self.accordion.set_title(0, 'Visualization Options')
+
+        # Event handlers
+        self._last_update = time.time()
+        self.button_reset.on_click(self.reset)
+        self.button_apply.on_click(self.apply_with_debounce)
+        self.metrics_label = widgets.Label(value="")
+
+        # Observe widget changes directly
+        self.my_slider.observe(self.apply_with_debounce, names='value')
+        self.my_slider2.observe(self.apply_with_debounce, names='value')
+        self.colorpicker.observe(self.apply_with_debounce, names='value')
+
+    def stream(self, expr):
+        """Creates a TensorWatch stream from an expression."""
+        self.expr = expr
+        try:
+            self.streamdata = self.client.create_stream(expr=expr)
+            logging.debug("Stream created successfully")
+        except Exception as e:
+            logging.error(f"Error creating stream: {e}")
+            print(f"Error creating stream: {e}")
+        return self
+
+    def apply_with_debounce(self, _=None):
+        """Debounced apply function to prevent too frequent updates."""
+        now = time.time()
+        if now - self._last_update > self.update_interval:
+            self.update_visualizer()
+            self._last_update = now
+            if self.button_apply.description == "Start":
+                self.button_apply.description = "Apply Changes"
+
+    def update_visualizer(self, _=None):
+        """Updates the TensorWatch visualizer with the latest widget values."""
+        if not hasattr(self, 'streamdata') or not self.streamdata:
+            self.out.clear_output(wait=True)
+            with self.out:
+                print("Stream data not available or empty yet. Please wait for data.")
+            return
+
+        try:
+            # Always clear output before drawing
+            self.out.clear_output(wait=True)
+
+            # Close previous visualizer if it exists to free resources
+            if self.visualizer:
+                self.visualizer.close()
+                plt.close('all') # Also close any lingering matplotlib figures
+
+            # Create a new visualizer with the current settings
+            self.visualizer = tw.Visualizer(
+                self.streamdata,
+                vis_type="line",
+                window_width=self.my_slider2.value,
+                window_size=self.my_slider.value,
+                Date=self.datebutton.value,
+                useOffset=self.offsetbutton.value,
+                dim_history=self.dimhistorybutton.value,
+                color=self.colorpicker.value,
+            )
+            with self.out:
+                self.visualizer.show()
+                
+        except Exception as e:
+            self.out.clear_output(wait=True)
+            with self.out:
+                print(f"Error updating visualizer: {e}")
+
+    def enable_apply_button(self):
+        """Enables the apply button and changes its description to 'Start'."""
+        logging.debug("Enabling apply button.")
+        self.button_apply.disabled = False
+        self.button_apply.description = "Start"
+
+    def reset(self, _=None):
+        """Resets all widget values to their defaults and clears the visualization."""
+        self.my_slider.value = self.default_value
+        self.my_slider2.value = self.default_value
+        self.datebutton.value = False
+        self.offsetbutton.value = False
+        self.dimhistorybutton.value = True
+        self.colorpicker.value = "blue"
+        
+        # Clear the output and close the visualizer
+        self.out.clear_output()
+        if self.visualizer:
+            self.visualizer.close()
+        plt.close('all')
+        self.visualizer = None
+
+    def draw(self):
+        """Displays the UI for controlling the visualization."""
+        ui = widgets.VBox([
+            widgets.HBox([self.button_reset, self.button_apply]),
+            self.accordion,
+            self.out
+        ])
+        display(ui)
+
+    def draw_with_metrics(self):
+        """Displays the UI for controlling the visualization with a metrics label."""
+        ui = widgets.VBox([
+            self.metrics_label,
+            widgets.HBox([self.button_reset, self.button_apply]),
+            self.accordion,
+            self.out
+        ])
+        display(ui)
+
+    def update_metrics(self, metrics):
+        """Updates the metrics label with the provided text."""
+        self.metrics_label.value = metrics
+
+    def connector(self, topic, host, parsetype="json", cluster_size=1, conn_type="kafka", queue_length=50000,
+                  group_id="mygroup", avro_schema=None, schema_path=None, protobuf_message=None, parser_extra=None,
+                  random_sampling=None, countmin_width=None, countmin_depth=None):
+        """
+        Creates and returns a Kafka or PyKafka connector.
+
+        Args:
+            topic (str): The Kafka topic to consume from.
+            host (str): The Kafka broker host.
+            parsetype (str): The message format (e.g., 'json', 'pickle', 'avro').
+            cluster_size (int): The number of consumer threads.
+            conn_type (str): The type of connector to use ('kafka' or 'pykafka').
+            queue_length (int): The maximum size of the message queue.
+            group_id (str): The Kafka consumer group ID.
+            avro_schema (str): The Avro schema for 'kafka' connector.
+            schema_path (str): The path to the schema file.
+            protobuf_message (str): The name of the Protobuf message class.
+            parser_extra (str): Extra data for the parser (e.g., Avro schema for 'pykafka').
+            random_sampling (int): The percentage of messages to sample.
+            countmin_width (int): The width of the Count-Min Sketch.
+            countmin_depth (int): The depth of the Count-Min Sketch.
+
+        Returns:
+            A KafkaConnector or pykafka_connector instance.
+        """
+        if conn_type == "kafka":
+            return kc.KafkaConnector(
+                topic=topic, hosts=host, parsetype=parsetype, cluster_size=cluster_size,
+                twapi_instance=self, queue_length=queue_length, group_id=group_id,
+                avro_schema=avro_schema, schema_path=schema_path, protobuf_message=protobuf_message,
+                random_sampling=random_sampling, countmin_width=countmin_width,
+                countmin_depth=countmin_depth)
+        elif conn_type == "pykafka":
+            return pyc.pykafka_connector(
+                topic=topic, hosts=host, parsetype=parsetype, cluster_size=cluster_size,twapi_instance=self, 
+                queue_length=queue_length, consumer_group=bytes(group_id, 'utf-8'),
+                parser_extra=parser_extra, scema_path=schema_path, probuf_message=protobuf_message,
+                random_sampling=random_sampling, countmin_width=countmin_width,
+                countmin_depth=countmin_depth)
+        else:
+            raise ValueError("Invalid connector type. Choose 'kafka' or 'pykafka'.")
+
+    async def some_async_function(self):
+        """Example of an async function that can be called."""
+        await asyncio.sleep(1)
+        print("Async function completed")