You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
93 lines
2.7 KiB
93 lines
2.7 KiB
"""
|
|
main.py — Agent worker entry point. Runs all agents concurrently.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import asyncpg
|
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
|
|
|
|
class AgentSettings(BaseSettings):
|
|
model_config = SettingsConfigDict(env_file='.env', extra='ignore')
|
|
database_url: str = 'postgresql://brain:brain@postgres:5432/second_brain'
|
|
ollama_url: str = 'http://ollama:11434'
|
|
chat_model: str = 'mistral'
|
|
log_level: str = 'INFO'
|
|
ingestion_poll: int = 15
|
|
linking_poll: int = 30
|
|
tagging_poll: int = 60
|
|
summarization_poll: int = 120
|
|
maintenance_poll: int = 3600
|
|
|
|
|
|
def setup_logging(level: str) -> None:
|
|
logging.basicConfig(
|
|
level=getattr(logging, level.upper(), logging.INFO),
|
|
format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
|
|
datefmt='%Y-%m-%dT%H:%M:%S',
|
|
stream=sys.stdout,
|
|
)
|
|
|
|
|
|
async def main() -> None:
|
|
settings = AgentSettings()
|
|
setup_logging(settings.log_level)
|
|
logger = logging.getLogger('agents')
|
|
logger.info('Starting agent workers...')
|
|
|
|
# Add parent dirs to path for cross-service imports
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / 'ingestion-worker'))
|
|
|
|
pool = await asyncpg.create_pool(settings.database_url, min_size=2, max_size=10)
|
|
|
|
# Import agents after path setup
|
|
from ingestion.agent import IngestionAgent
|
|
from linking.agent import LinkingAgent
|
|
from tagging.agent import TaggingAgent
|
|
from summarization.agent import SummarizationAgent
|
|
from maintenance.agent import MaintenanceAgent
|
|
|
|
agents_tasks = [
|
|
asyncio.create_task(
|
|
IngestionAgent(pool, settings).run_forever(settings.ingestion_poll)
|
|
),
|
|
asyncio.create_task(
|
|
LinkingAgent(pool, settings).run_forever(settings.linking_poll)
|
|
),
|
|
asyncio.create_task(
|
|
TaggingAgent(pool, settings).run_forever(settings.tagging_poll)
|
|
),
|
|
asyncio.create_task(
|
|
SummarizationAgent(pool, settings).run_forever(settings.summarization_poll)
|
|
),
|
|
asyncio.create_task(
|
|
MaintenanceAgent(pool, settings).run_forever(settings.maintenance_poll)
|
|
),
|
|
]
|
|
|
|
logger.info('All agents running.')
|
|
|
|
try:
|
|
await asyncio.gather(*agents_tasks)
|
|
except asyncio.CancelledError:
|
|
pass
|
|
finally:
|
|
for task in agents_tasks:
|
|
task.cancel()
|
|
await pool.close()
|
|
logger.info('Agent workers stopped.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
pass
|