Chris4K commited on
Commit
4066fdd
·
verified ·
1 Parent(s): 126ce84

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +47 -55
app.py CHANGED
@@ -87,84 +87,90 @@ import uvicorn
87
  # FastAPI App
88
  app = FastAPI()
89
 
 
 
 
 
 
 
 
 
 
90
  class AIStateManager:
91
  def __init__(self):
92
  self.state = "awake"
93
- self.msg_queue = queue.Queue()
94
  self.heartbeat_count = 0
95
- self.lock = threading.Lock()
96
- self.clients = set()
97
 
98
- # Research Task List
99
- self.research_tasks = ["Explore AI Ethics", "Find latest AI models", "Investigate quantum computing"]
100
- self.current_task = None
101
-
102
- def set_state(self, new_state):
103
- with self.lock:
104
  print(f"[STATE CHANGE] {self.state} → {new_state}")
105
  self.state = new_state
106
- self.add_message("system", f"State changed to {new_state}")
107
 
108
- def receive_message(self, sender, message):
109
  """Adds messages to queue and resets heartbeat if human sends input."""
110
- with self.lock:
111
- self.msg_queue.put((sender, message))
112
  if sender == "human":
113
  self.heartbeat_count = 0
114
  if self.state != "awake":
115
- self.set_state("awake")
116
 
117
- def add_message(self, sender, message):
118
  """AI or System can add messages to queue."""
119
- with self.lock:
120
- self.msg_queue.put((sender, message))
121
 
122
- def process_messages(self):
123
  """Processes messages in queue."""
124
  while not self.msg_queue.empty():
125
- sender, message = self.msg_queue.get()
126
  print(f"[{sender.upper()}] {message}")
127
- asyncio.create_task(self.broadcast(f"[{sender.upper()}] {message}"))
128
 
129
  async def broadcast(self, message):
130
  """Sends message to all connected WebSocket clients."""
131
- print("sending broadcast message")
132
- for client in ai_manager.clients:
133
- await client.send_text(f"Received: {data}")
134
- #for ws in self.clients:
135
- # await ws.send_text(message)
 
 
136
 
137
  async def heartbeat(self):
138
  """Basic heartbeat loop."""
139
  while True:
140
  await asyncio.sleep(1) # One 'beat'
141
- with self.lock:
142
  self.heartbeat_count += 1
143
- self.process_messages()
144
 
145
  async def consciousness(self):
146
  """Controls research/sleep cycle."""
147
  while True:
148
- print("one two three")
149
- self.broadcast("one, two three...")
150
  await asyncio.sleep(2) # Consciousness checks every 2 sec
151
- with self.lock:
152
  if self.state == "awake":
153
  if self.heartbeat_count >= 5:
154
- self.set_state("research")
155
  asyncio.create_task(self.run_research())
156
 
157
  elif self.state == "research":
158
  if not self.research_tasks: # If no tasks left, move to sleep
159
- self.set_state("sleeping")
160
  asyncio.create_task(self.run_sleeping())
161
 
162
  elif self.state == "sleeping":
163
  if self.heartbeat_count >= 20:
164
- self.set_state("research") # Restart research after sleep
165
  asyncio.create_task(self.run_research())
166
  print(self.state)
167
-
168
  async def run_research(self):
169
  """Runs research tasks in order."""
170
  while self.state == "research" and self.research_tasks:
@@ -202,33 +208,19 @@ threading.Thread(target=ai_manager.modify_research_tasks, daemon=True).start()
202
 
203
  @app.websocket("/ws")
204
  async def websocket_endpoint(websocket: WebSocket):
205
- print("WebSocket endpoint called")
206
  await websocket.accept()
207
- print(f"WebSocket connection accepted from {websocket.client}")
208
-
 
209
  try:
210
- # Send initial message explicitly
211
- await websocket.send_text("Connected to AI State Machine")
212
- print("Initial message sent to WebSocket")
213
-
214
- # Directly broadcast the message
215
- #for client in ai_manager.clients:
216
- # await client.send_text(f"Received: {data}")
217
-
218
-
219
  while True:
220
  data = await websocket.receive_text()
221
  print(f"BACKEND RECEIVED: {data}")
222
-
223
- # Echo the message back
224
- response = f"Echo: {data}"
225
- await websocket.send_text(response)
226
- print(f"BACKEND SENT: {response}")
227
  except Exception as e:
228
- print(f"WebSocket ERROR: {e}")
229
- # Log full traceback
230
- import traceback
231
- traceback.print_exc()
232
 
233
  @app.get("/")
234
  async def get():
 
87
  # FastAPI App
88
  app = FastAPI()
89
 
90
+ import asyncio
91
+ import queue
92
+ import threading
93
+ import random
94
+ import time
95
+ from fastapi import FastAPI, WebSocket
96
+ from fastapi.responses import HTMLResponse
97
+ import uvicorn
98
+
99
  class AIStateManager:
100
  def __init__(self):
101
  self.state = "awake"
102
+ self.msg_queue = asyncio.Queue() # Use asyncio.Queue for async operations
103
  self.heartbeat_count = 0
104
+ self.lock = asyncio.Lock() # Use asyncio.Lock for async locking
105
+ self.clients = set() # Set to store WebSocket clients
106
 
107
+ async def set_state(self, new_state):
108
+ async with self.lock:
 
 
 
 
109
  print(f"[STATE CHANGE] {self.state} → {new_state}")
110
  self.state = new_state
111
+ await self.add_message("system", f"State changed to {new_state}")
112
 
113
+ async def receive_message(self, sender, message):
114
  """Adds messages to queue and resets heartbeat if human sends input."""
115
+ async with self.lock:
116
+ await self.msg_queue.put((sender, message))
117
  if sender == "human":
118
  self.heartbeat_count = 0
119
  if self.state != "awake":
120
+ await self.set_state("awake")
121
 
122
+ async def add_message(self, sender, message):
123
  """AI or System can add messages to queue."""
124
+ await self.msg_queue.put((sender, message))
 
125
 
126
+ async def process_messages(self):
127
  """Processes messages in queue."""
128
  while not self.msg_queue.empty():
129
+ sender, message = await self.msg_queue.get()
130
  print(f"[{sender.upper()}] {message}")
131
+ await self.broadcast(f"[{sender.upper()}] {message}")
132
 
133
  async def broadcast(self, message):
134
  """Sends message to all connected WebSocket clients."""
135
+ print(f"Broadcasting: {message}")
136
+ for client in list(self.clients): # Create a copy to avoid runtime modification
137
+ try:
138
+ await client.send_text(message)
139
+ except Exception as e:
140
+ print(f"Error broadcasting to client: {e}")
141
+ self.clients.remove(client)
142
 
143
  async def heartbeat(self):
144
  """Basic heartbeat loop."""
145
  while True:
146
  await asyncio.sleep(1) # One 'beat'
147
+ async with self.lock:
148
  self.heartbeat_count += 1
149
+ await self.process_messages()
150
 
151
  async def consciousness(self):
152
  """Controls research/sleep cycle."""
153
  while True:
154
+ print("Consciousness cycle")
155
+ await self.broadcast("Consciousness cycle...")
156
  await asyncio.sleep(2) # Consciousness checks every 2 sec
157
+ async with self.lock:
158
  if self.state == "awake":
159
  if self.heartbeat_count >= 5:
160
+ await self.set_state("research")
161
  asyncio.create_task(self.run_research())
162
 
163
  elif self.state == "research":
164
  if not self.research_tasks: # If no tasks left, move to sleep
165
+ await self.set_state("sleeping")
166
  asyncio.create_task(self.run_sleeping())
167
 
168
  elif self.state == "sleeping":
169
  if self.heartbeat_count >= 20:
170
+ await self.set_state("research") # Restart research after sleep
171
  asyncio.create_task(self.run_research())
172
  print(self.state)
173
+
174
  async def run_research(self):
175
  """Runs research tasks in order."""
176
  while self.state == "research" and self.research_tasks:
 
208
 
209
  @app.websocket("/ws")
210
  async def websocket_endpoint(websocket: WebSocket):
 
211
  await websocket.accept()
212
+ print("Initial message sent to WebSocket")
213
+
214
+ ai_manager.clients.add(websocket) # Add client to set
215
  try:
 
 
 
 
 
 
 
 
 
216
  while True:
217
  data = await websocket.receive_text()
218
  print(f"BACKEND RECEIVED: {data}")
219
+ await ai_manager.receive_message("human", data)
 
 
 
 
220
  except Exception as e:
221
+ print(f"WebSocket error: {e}")
222
+ finally:
223
+ ai_manager.clients.remove(websocket) # Remove client when connection closes
 
224
 
225
  @app.get("/")
226
  async def get():