Separate and tested web_socket_manager
This commit is contained in:
136
sia/__main__.py
136
sia/__main__.py
@@ -1,65 +1,105 @@
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from aiohttp import web
|
||||
from pathlib import Path
|
||||
import asyncio
|
||||
import mimetypes
|
||||
import time
|
||||
|
||||
from .llm_engine import LlmEngine
|
||||
from .system_metrics import SystemMetrics
|
||||
from .web_agent import WebAgent
|
||||
from .web_io_buffer import WebIOBuffer
|
||||
from .web_server import WebServer
|
||||
from .web_socket_manager import WebSocketManager
|
||||
from .working_memory import WorkingMemory
|
||||
from .xml_validator import XMLValidator
|
||||
from .response_parser import ResponseParser
|
||||
from .web_agent import WebAgent
|
||||
from .web_io_buffer import WebIOBuffer
|
||||
|
||||
async def run_web_server():
|
||||
"""Run the web server with default configuration."""
|
||||
base_dir = Path(__file__).parent.parent
|
||||
mimetypes.add_type("application/javascript", ".js")
|
||||
mimetypes.add_type("application/javascript", ".jsx")
|
||||
mimetypes.add_type("text/javascript", ".js")
|
||||
mimetypes.add_type("text/javascript", ".jsx")
|
||||
|
||||
# Load system prompt and action schema
|
||||
system_prompt = (base_dir / "system_prompt.md").read_text()
|
||||
action_schema = (base_dir / "action_schema.xsd").read_text()
|
||||
class TestLLM:
|
||||
def infer(self, prompt: str, context: str):
|
||||
yield "<reasoning>"
|
||||
time.sleep(2)
|
||||
yield "test reasoning"
|
||||
time.sleep(2)
|
||||
yield "</reasoning>"
|
||||
|
||||
# Initialize core components
|
||||
llm = LlmEngine("/root/model")
|
||||
io_buffer = WebIOBuffer()
|
||||
agent = WebAgent(
|
||||
system_prompt=system_prompt,
|
||||
action_schema=action_schema,
|
||||
working_memory=WorkingMemory(),
|
||||
system_metrics=SystemMetrics(),
|
||||
llm=llm,
|
||||
validator=XMLValidator(action_schema),
|
||||
parser=ResponseParser(io_buffer)
|
||||
)
|
||||
class Main:
|
||||
def __init__(self):
|
||||
self._base_dir = Path(__file__).parent.parent
|
||||
self._system_prompt = (self._base_dir / "system_prompt.md").read_text()
|
||||
self._action_schema = (self._base_dir / "action_schema.xsd").read_text()
|
||||
self._static_dir = self._base_dir / "static"
|
||||
|
||||
# Start the web server
|
||||
server = WebServer(agent, io_buffer, "0.0.0.0", 8080)
|
||||
await server.start()
|
||||
self._llm = LlmEngine("/root/model")
|
||||
#self._llm = TestLLM()
|
||||
self._io_buffer = WebIOBuffer()
|
||||
self._agent = WebAgent(
|
||||
system_prompt=self._system_prompt,
|
||||
action_schema=self._action_schema,
|
||||
working_memory=WorkingMemory(),
|
||||
system_metrics=SystemMetrics(),
|
||||
llm=self._llm,
|
||||
validator=XMLValidator(self._action_schema),
|
||||
parser=ResponseParser(self._io_buffer)
|
||||
)
|
||||
self._ws_manager = WebSocketManager(
|
||||
self._agent,
|
||||
self._io_buffer
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
if server._response_tasks:
|
||||
done, _ = await asyncio.wait(
|
||||
server._response_tasks,
|
||||
timeout=0.1,
|
||||
return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
for task in done:
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
print("Task was cancelled")
|
||||
except Exception as e:
|
||||
print(f"Task error: {e}")
|
||||
server._response_tasks.discard(task) # Use discard to prevent KeyError
|
||||
await asyncio.sleep(0.1)
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down...")
|
||||
finally:
|
||||
await server.stop()
|
||||
self._app = web.Application()
|
||||
self._init_routes()
|
||||
|
||||
def main():
|
||||
"""Main entry point for the SIA application."""
|
||||
asyncio.run(run_web_server())
|
||||
async def app(self):
|
||||
return self._app
|
||||
|
||||
def _init_routes(self):
|
||||
"""Initialize application routes."""
|
||||
self._app.middlewares.append(self._cors_middleware)
|
||||
self._app.router.add_get("/ws", self._ws_manager.handle_websocket)
|
||||
self._app.router.add_get("/", self._serve_index)
|
||||
self._app.router.add_static("/static/", self._static_dir, show_index=False)
|
||||
self._app.router.add_static("/assets/", self._static_dir / "assets", show_index=False)
|
||||
self._app.router.add_get("/{path:.*}", self._serve_index)
|
||||
|
||||
async def _serve_index(self, request: web.Request) -> web.Response:
|
||||
"""Serve the React application HTML for any unmatched routes."""
|
||||
index_path = self._static_dir / "index.html"
|
||||
if not index_path.exists():
|
||||
raise web.HTTPNotFound()
|
||||
|
||||
with open(index_path, "r") as f:
|
||||
html_content = f.read()
|
||||
|
||||
return web.Response(
|
||||
text=html_content,
|
||||
content_type="text/html"
|
||||
)
|
||||
|
||||
@web.middleware
|
||||
async def _cors_middleware(self, request: web.Request, handler):
|
||||
"""Handle CORS headers."""
|
||||
if request.method == "OPTIONS":
|
||||
response = web.Response()
|
||||
else:
|
||||
response = await handler(request)
|
||||
|
||||
response.headers.update({
|
||||
'Access-Control-Allow-Origin': '*',
|
||||
'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
|
||||
'Access-Control-Allow-Headers': 'Content-Type',
|
||||
'Access-Control-Max-Age': '86400',
|
||||
})
|
||||
return response
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main = Main()
|
||||
app = asyncio.run(main.app())
|
||||
print("Web server started at http://localhost:8080")
|
||||
web.run_app(app, port=8080)
|
||||
Reference in New Issue
Block a user