- tracing.py: init/shutdown, trace_span, get_langfuse_callback, prompt mgmt - main.py: init_langfuse at startup, shutdown on teardown - redis_consumer.py: trace_span around journey_start/message/agent_trigger - agent_runner.py: thread langfuse_handler through classify + processing LLM - journey.py: thread langfuse_handler through _call_llm_with_tools - llm.py: accept callbacks param, forward to LLM constructors - requirements.txt: add langfuse>=3.0.0
72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
"""Batch Agent Service — FastAPI application.
|
|
|
|
Owns: agent_runner (local directory + cloud connectors), journey builder,
|
|
filesystem_agent, integrations (Gmail, MS Graph).
|
|
|
|
Communicates with WS Gateway via Redis:
|
|
- Subscribes to batch:request:{user_id} (journey_start, journey_message)
|
|
- Publishes to ws:out:{user_id} (journey replies + tool calls)
|
|
- BRPOP on tool:result:{call_id} (tool-call round-trip, 30s timeout)
|
|
- SET+EX on journey:{user_id} (journey session state, TTL 1800s)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from app.redis_consumer import start_consumer
|
|
from app.routes import router
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
|
# Initialise Langfuse tracing (no-op if keys are missing)
|
|
from app.tracing import init_langfuse
|
|
init_langfuse()
|
|
|
|
logger.info("batch-agent: starting Redis consumer")
|
|
task = asyncio.create_task(start_consumer())
|
|
yield
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
from app.tracing import shutdown as shutdown_langfuse
|
|
shutdown_langfuse()
|
|
|
|
from shared.db import engine
|
|
await engine.dispose()
|
|
|
|
from shared.redis import redis_client
|
|
await redis_client.aclose()
|
|
|
|
logger.info("batch-agent: Redis consumer stopped")
|
|
|
|
|
|
app = FastAPI(title="Adiuva Batch Agent Service", lifespan=lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["GET", "POST"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
app.include_router(router)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health() -> dict[str, str]:
|
|
return {"status": "ok", "service": "batch-agent"}
|