acp.py.j2

"""
ACP handler for async LangGraph agent.

Uses the async ACP model with Redis streaming instead of HTTP yields.
"""

from dotenv import load_dotenv

load_dotenv()
import os

# LiteLLM proxy auth: copy LITELLM_API_KEY to OPENAI_API_KEY for OpenAI client compatibility
_litellm_key = os.environ.get("LITELLM_API_KEY")
if _litellm_key:
    os.environ["OPENAI_API_KEY"] = _litellm_key

import agentex.lib.adk as adk
from agentex.lib.adk import create_langgraph_tracing_handler, stream_langgraph_events
from agentex.lib.core.tracing.tracing_processor_manager import add_tracing_processor_config
from agentex.lib.sdk.fastacp.fastacp import FastACP
from agentex.protocol.acp import SendEventParams, CancelTaskParams, CreateTaskParams
from agentex.lib.types.fastacp import AsyncACPConfig
from agentex.lib.types.tracing import SGPTracingProcessorConfig
from agentex.lib.utils.logging import make_logger

from project.graph import create_graph

logger = make_logger(__name__)

add_tracing_processor_config(
    SGPTracingProcessorConfig(
        sgp_api_key=os.environ.get("SGP_API_KEY", ""),
        sgp_account_id=os.environ.get("SGP_ACCOUNT_ID", ""),
        sgp_base_url=os.environ.get("SGP_CLIENT_BASE_URL", ""),
    ))

acp = FastACP.create(
    acp_type="async",
    config=AsyncACPConfig(type="base"),
)

_graph = None


async def get_graph():
    global _graph
    if _graph is None:
        _graph = await create_graph()
    return _graph


@acp.on_task_event_send
async def handle_task_event_send(params: SendEventParams):
    """Handle incoming events, streaming tokens and tool calls via Redis."""
    graph = await get_graph()
    task_id = params.task.id
    user_message = params.event.content.content

    logger.info(f"Processing message for thread {task_id}")

    # Echo the user's message
    await adk.messages.create(task_id=task_id, content=params.event.content)

    async with adk.tracing.span(
        trace_id=task_id,
        name="message",
        input={"message": user_message},
        data={"__span_type__": "AGENT_WORKFLOW"},
    ) as turn_span:
        callback = create_langgraph_tracing_handler(
            trace_id=task_id,
            parent_span_id=turn_span.id if turn_span else None,
        )

        stream = graph.astream(
            {"messages": [{"role": "user", "content": user_message}]},
            config={
                "configurable": {"thread_id": task_id},
                "callbacks": [callback],
            },
            stream_mode=["messages", "updates"],
        )

        final_output = await stream_langgraph_events(stream, task_id)

        if turn_span:
            turn_span.output = {"final_output": final_output}


@acp.on_task_create
async def handle_task_create(params: CreateTaskParams):
    logger.info(f"Task created: {params.task.id}")


@acp.on_task_cancel
async def handle_task_canceled(params: CancelTaskParams):
    logger.info(f"Task canceled: {params.task.id}")