server.py

"""
Nutricore Food Solution - ChatKit Server
"""

from collections.abc import AsyncIterator
from chatkit.server import ChatKitServer
from chatkit.types import ThreadMetadata, ThreadStreamEvent, UserMessageItem
from chatkit.agents import AgentContext, simple_to_agent_input, stream_agent_response
from agents import Runner

from agent import nutri_agent
from store import InMemoryStore


class NutricoreServer(ChatKitServer[dict]):

    def __init__(self, store: InMemoryStore):
        super().__init__(store)

    async def respond(
        self,
        thread: ThreadMetadata,
        input_user_message: UserMessageItem | None,
        context: dict,
    ) -> AsyncIterator[ThreadStreamEvent]:
        items_page = await self.store.load_thread_items(
            thread.id, after=None, limit=20, order="asc", context=context
        )
        agent_input = await simple_to_agent_input(items_page.data)
        agent_context = AgentContext(thread=thread, store=self.store, request_context=context)
        result = Runner.run_streamed(nutri_agent, agent_input, context=agent_context)
        async for event in stream_agent_response(agent_context, result):
            yield event


store = InMemoryStore()
nutricore_server = NutricoreServer(store=store)