165 lines
4.7 KiB
Python
165 lines
4.7 KiB
Python
import contextlib
|
|
import logging
|
|
import threading
|
|
|
|
import zmq
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from zmq.asyncio import Context
|
|
|
|
# BDI agents
|
|
from control_backend.agents.bdi import (
|
|
BDIBeliefCollectorAgent,
|
|
BDICoreAgent,
|
|
TextBeliefExtractorAgent,
|
|
)
|
|
|
|
# Communication agents
|
|
from control_backend.agents.communication import RICommunicationAgent
|
|
|
|
# Emotional Agents
|
|
# LLM Agents
|
|
from control_backend.agents.llm import LLMAgent
|
|
|
|
# Perceive agents
|
|
from control_backend.agents.perception import VADAgent
|
|
|
|
# Other backend imports
|
|
from control_backend.api.v1.router import api_router
|
|
from control_backend.core.config import settings
|
|
from control_backend.logging import setup_logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def setup_sockets():
|
|
context = Context.instance()
|
|
|
|
internal_pub_socket = context.socket(zmq.XPUB)
|
|
internal_pub_socket.bind(settings.zmq_settings.internal_sub_address)
|
|
logger.debug("Internal publishing socket bound to %s", internal_pub_socket)
|
|
|
|
internal_sub_socket = context.socket(zmq.XSUB)
|
|
internal_sub_socket.bind(settings.zmq_settings.internal_pub_address)
|
|
logger.debug("Internal subscribing socket bound to %s", internal_sub_socket)
|
|
try:
|
|
zmq.proxy(internal_sub_socket, internal_pub_socket)
|
|
except zmq.ZMQError:
|
|
logger.warning("Error while handling PUB/SUB proxy. Closing sockets.")
|
|
finally:
|
|
internal_pub_socket.close()
|
|
internal_sub_socket.close()
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""
|
|
Application lifespan context manager to handle startup and shutdown events.
|
|
"""
|
|
# --- APPLICATION STARTUP ---
|
|
setup_logging()
|
|
logger.info("%s is starting up.", app.title)
|
|
|
|
# Initiate sockets
|
|
proxy_thread = threading.Thread(target=setup_sockets)
|
|
proxy_thread.daemon = True
|
|
proxy_thread.start()
|
|
|
|
context = Context.instance()
|
|
|
|
endpoints_pub_socket = context.socket(zmq.PUB)
|
|
endpoints_pub_socket.connect(settings.zmq_settings.internal_pub_address)
|
|
app.state.endpoints_pub_socket = endpoints_pub_socket
|
|
|
|
# --- Initialize Agents ---
|
|
logger.info("Initializing and starting agents.")
|
|
|
|
agents_to_start = {
|
|
"RICommunicationAgent": (
|
|
RICommunicationAgent,
|
|
{
|
|
"name": settings.agent_settings.ri_communication_name,
|
|
"address": settings.zmq_settings.ri_communication_address,
|
|
"bind": True,
|
|
},
|
|
),
|
|
"LLMAgent": (
|
|
LLMAgent,
|
|
{
|
|
"name": settings.agent_settings.llm_name,
|
|
},
|
|
),
|
|
"BDICoreAgent": (
|
|
BDICoreAgent,
|
|
{
|
|
"name": settings.agent_settings.bdi_core_name,
|
|
"asl": "src/control_backend/agents/bdi/bdi_core_agent/rules.asl",
|
|
},
|
|
),
|
|
"BeliefCollectorAgent": (
|
|
BDIBeliefCollectorAgent,
|
|
{
|
|
"name": settings.agent_settings.bdi_belief_collector_name,
|
|
},
|
|
),
|
|
"TextBeliefExtractorAgent": (
|
|
TextBeliefExtractorAgent,
|
|
{
|
|
"name": settings.agent_settings.text_belief_extractor_name,
|
|
},
|
|
),
|
|
"VADAgent": (
|
|
VADAgent,
|
|
{"audio_in_address": settings.zmq_settings.vad_agent_address, "audio_in_bind": False},
|
|
),
|
|
}
|
|
|
|
agents = []
|
|
|
|
vad_agent = None
|
|
|
|
for name, (agent_class, kwargs) in agents_to_start.items():
|
|
try:
|
|
logger.debug("Starting agent: %s", name)
|
|
agent_instance = agent_class(**kwargs)
|
|
await agent_instance.start()
|
|
if isinstance(agent_instance, VADAgent):
|
|
vad_agent = agent_instance
|
|
agents.append(agent_instance)
|
|
logger.info("Agent '%s' started successfully.", name)
|
|
except Exception as e:
|
|
logger.error("Failed to start agent '%s': %s", name, e, exc_info=True)
|
|
raise
|
|
|
|
assert vad_agent is not None
|
|
await vad_agent.reset_stream()
|
|
|
|
logger.info("Application startup complete.")
|
|
|
|
yield
|
|
|
|
# --- APPLICATION SHUTDOWN ---
|
|
logger.info("%s is shutting down.", app.title)
|
|
|
|
# Potential shutdown logic goes here
|
|
|
|
logger.info("Application shutdown complete.")
|
|
|
|
|
|
# if __name__ == "__main__":
|
|
app = FastAPI(title=settings.app_title, lifespan=lifespan)
|
|
|
|
# This middleware allows other origins to communicate with us
|
|
app.add_middleware(
|
|
CORSMiddleware, # https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/CORS
|
|
allow_origins=[settings.ui_url], # address of our UI application
|
|
allow_methods=["*"], # GET, POST, etc.
|
|
)
|
|
|
|
app.include_router(api_router, prefix="") # TODO: make prefix /api/v1
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return {"status": "ok"}
|