CopilotKit integration

from typing import AsyncGenerator

from blaxel.langgraph import bl_model, bl_tools
from langchain.tools import tool
from langchain_core.messages import AIMessageChunk
from langgraph.prebuilt import create_react_agent

@tool
def weather(city: str) -> str:
    """Get the weather in a given city"""
    return f"The weather in {city} is sunny"

async def langgraph_agent() -> AsyncGenerator[str, None]:
    prompt = (
        "You are a helpful assistant that can answer questions and help with tasks."
    )
    tools = await bl_tools(["blaxel-search"]) + [weather]
    model = await bl_model("sandbox-openai")
    agent = create_react_agent(model=model, tools=tools, prompt=prompt)
    return agent

async def agent(input: str) -> AsyncGenerator[str, None]:
    agent = await langgraph_agent()
    messages = {"messages": [("user", input)]}
    async for chunk in agent.astream(messages, stream_mode=["updates", "messages"]):
        type_, stream_chunk = chunk
        # This is to stream the response from the agent, filtering response from tools
        if (
            type_ == "messages"
            and len(stream_chunk) > 0
            and isinstance(stream_chunk[0], AIMessageChunk)
        ):
            msg = stream_chunk[0]
            if msg.content:
                if not msg.tool_calls:
                    yield msg.content
        # This to show a call has been made to a tool, usefull if you want to show the tool call in your interface
        if type_ == "updates":
            if "tools" in stream_chunk:
                for msg in stream_chunk["tools"]["messages"]:
                    yield f"Tool call: {msg.name}\n"

import os
from contextlib import asynccontextmanager
from logging import getLogger

from copilotkit import CopilotKitRemoteEndpoint, LangGraphAgent
from copilotkit.integrations.fastapi import add_fastapi_endpoint
from fastapi import FastAPI
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

from .agent import langgraph_agent
from .server.error import init_error_handlers
from .server.middleware import init_middleware
from .server.router import router

logger = getLogger(__name__)

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(f"Server running on port {os.getenv('BL_SERVER_PORT', 80)}")
    try:
        # Initialize the graph
        graph = await langgraph_agent()

        # Initialize the SDK
        sdk = await get_sdk(graph)

        # Store in app state
        app.state.sdk = sdk

        # Add CopilotKit endpoint
        add_fastapi_endpoint(app, sdk, "/copilotkit", use_thread_pool=False)

        yield
        logger.info("Server shutting down")
    except Exception as e:
        logger.error(f"Error during startup: {str(e)}", exc_info=True)
        raise

# Create the SDK after the graph is initialized
async def get_sdk(graph):
    sdk = CopilotKitRemoteEndpoint(
        agents=[
            LangGraphAgent(
                name="sample_agent",
                description="An agent that can provide weather information and handle other conversational tasks",
                graph=graph,
            )
        ],
    )
    return sdk

app = FastAPI(lifespan=lifespan)
init_error_handlers(app)
init_middleware(app)
app.include_router(router)

FastAPIInstrumentor.instrument_app(app, exclude_spans=["receive", "send"])

POST https://run.blaxel.ai/{YOUR-WORKSPACE}/agents/{YOUR-AGENT}