Update app.py

app.py

@@ -87,84 +87,90 @@ import uvicorn
 # FastAPI App
 app = FastAPI()
 
+import asyncio
+import queue
+import threading
+import random
+import time
+from fastapi import FastAPI, WebSocket
+from fastapi.responses import HTMLResponse
+import uvicorn
+
 class AIStateManager:
     def __init__(self):
         self.state = "awake"
-        self.msg_queue = queue.Queue()
+        self.msg_queue = asyncio.Queue()  # Use asyncio.Queue for async operations
         self.heartbeat_count = 0
-        self.lock = threading.Lock()
-        self.clients = set()
+        self.lock = asyncio.Lock()  # Use asyncio.Lock for async locking
+        self.clients = set()  # Set to store WebSocket clients
 
-        # Research Task List
-        self.research_tasks = ["Explore AI Ethics", "Find latest AI models", "Investigate quantum computing"]
-        self.current_task = None
-
-    def set_state(self, new_state):
-        with self.lock:
+    async def set_state(self, new_state):
+        async with self.lock:
             print(f"[STATE CHANGE] {self.state} → {new_state}")
             self.state = new_state
-            self.add_message("system", f"State changed to {new_state}")
+            await self.add_message("system", f"State changed to {new_state}")
 
-    def receive_message(self, sender, message):
+    async def receive_message(self, sender, message):
         """Adds messages to queue and resets heartbeat if human sends input."""
-        with self.lock:
-            self.msg_queue.put((sender, message))
+        async with self.lock:
+            await self.msg_queue.put((sender, message))
             if sender == "human":
                 self.heartbeat_count = 0
                 if self.state != "awake":
-                    self.set_state("awake")
+                    await self.set_state("awake")
 
-    def add_message(self, sender, message):
+    async def add_message(self, sender, message):
         """AI or System can add messages to queue."""
-        with self.lock:
-            self.msg_queue.put((sender, message))
+        await self.msg_queue.put((sender, message))
 
-    def process_messages(self):
+    async def process_messages(self):
         """Processes messages in queue."""
         while not self.msg_queue.empty():
-            sender, message = self.msg_queue.get()
+            sender, message = await self.msg_queue.get()
             print(f"[{sender.upper()}] {message}")
-            asyncio.create_task(self.broadcast(f"[{sender.upper()}] {message}"))
+            await self.broadcast(f"[{sender.upper()}] {message}")
 
     async def broadcast(self, message):
         """Sends message to all connected WebSocket clients."""
-        print("sending broadcast message")
-        for client in ai_manager.clients:
-            await client.send_text(f"Received: {data}")
-        #for ws in self.clients:
-        #    await ws.send_text(message)
+        print(f"Broadcasting: {message}")
+        for client in list(self.clients):  # Create a copy to avoid runtime modification
+            try:
+                await client.send_text(message)
+            except Exception as e:
+                print(f"Error broadcasting to client: {e}")
+                self.clients.remove(client)
 
     async def heartbeat(self):
         """Basic heartbeat loop."""
         while True:
             await asyncio.sleep(1)  # One 'beat'
-            with self.lock:
+            async with self.lock:
                 self.heartbeat_count += 1
-                self.process_messages()
+                await self.process_messages()
 
     async def consciousness(self):
         """Controls research/sleep cycle."""
         while True:
-            print("one two three")
-            self.broadcast("one, two three...")
+            print("Consciousness cycle")
+            await self.broadcast("Consciousness cycle...")
             await asyncio.sleep(2)  # Consciousness checks every 2 sec
-            with self.lock:
+            async with self.lock:
                 if self.state == "awake":
                     if self.heartbeat_count >= 5:
-                        self.set_state("research")
+                        await self.set_state("research")
                         asyncio.create_task(self.run_research())
 
                 elif self.state == "research":
                     if not self.research_tasks:  # If no tasks left, move to sleep
-                        self.set_state("sleeping")
+                        await self.set_state("sleeping")
                         asyncio.create_task(self.run_sleeping())
 
                 elif self.state == "sleeping":
                     if self.heartbeat_count >= 20:
-                        self.set_state("research")  # Restart research after sleep
+                        await self.set_state("research")  # Restart research after sleep
                         asyncio.create_task(self.run_research())
                 print(self.state)
-
+                
     async def run_research(self):
         """Runs research tasks in order."""
         while self.state == "research" and self.research_tasks:
@@ -202,33 +208,19 @@ threading.Thread(target=ai_manager.modify_research_tasks, daemon=True).start()
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    print("WebSocket endpoint called")
     await websocket.accept()
-    print(f"WebSocket connection accepted from {websocket.client}")
-    
+    print("Initial message sent to WebSocket") 
+
+    ai_manager.clients.add(websocket)  # Add client to set
     try:
-        # Send initial message explicitly
-        await websocket.send_text("Connected to AI State Machine")
-        print("Initial message sent to WebSocket") 
-        
-            # Directly broadcast the message
-            #for client in ai_manager.clients:
-            #    await client.send_text(f"Received: {data}")
-            
-        
         while True:
             data = await websocket.receive_text()
             print(f"BACKEND RECEIVED: {data}")
-            
-            # Echo the message back
-            response = f"Echo: {data}"
-            await websocket.send_text(response)
-            print(f"BACKEND SENT: {response}")
+            await ai_manager.receive_message("human", data)
     except Exception as e:
-        print(f"WebSocket ERROR: {e}")
-        # Log full traceback
-        import traceback
-        traceback.print_exc()
+        print(f"WebSocket error: {e}")
+    finally:
+        ai_manager.clients.remove(websocket)  # Remove client when connection closes
 
 @app.get("/")
 async def get():