trialpath/agent/setup.py

"""Parlant server and agent initialization."""

import os

from parlant.core.loggers import LogLevel
from parlant.sdk import Agent, NLPServices, Server

AGENT_NAME = "Patient Trial Copilot"
AGENT_DESCRIPTION = (
    "AI copilot helping NSCLC patients find matching clinical trials. "
    "Extracts patient profiles from medical documents, searches ClinicalTrials.gov, "
    "evaluates eligibility criterion-by-criterion, and identifies actionable gaps "
    "where additional data could unlock new trial options."
)

_LOG_LEVEL_MAP = {
    "trace": LogLevel.TRACE,
    "debug": LogLevel.DEBUG,
    "info": LogLevel.INFO,
    "warning": LogLevel.WARNING,
    "error": LogLevel.ERROR,
}


async def create_server_and_agent(
    host: str = "0.0.0.0",
    port: int = 8800,
) -> tuple[Server, Agent]:
    """Create and start Parlant server with the TrialPath agent."""
    log_level_str = os.environ.get("PARLANT_LOG_LEVEL", "info").lower()
    log_level = _LOG_LEVEL_MAP.get(log_level_str, LogLevel.INFO)

    server = Server(
        nlp_service=NLPServices.gemini,
        host=host,
        port=port,
        log_level=log_level,
    )
    await server.__aenter__()

    agent = await server.create_agent(
        name=AGENT_NAME,
        description=AGENT_DESCRIPTION,
        max_engine_iterations=10,
    )

    return server, agent