Agent for cleaning up named entities in YouTube video transcripts.

demo-enrichment-agent.py

...
from langchain_core.runnables import Runnable
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, SystemMessage

...

class DemoEnrichmentAgent:
    extractor_llm: Runnable

    def __init__(self):
        self.extractor_llm = init_chat_model(
            "gpt-4o-mini",
            model_provider="openai"
        )

    async def extractor_node(self, state: AgentState):
        sys_prompt = SystemMessage(
            content=(
                "You are a master at data extraction. "
                "You will receive the transcript of a YouTube video. "
                "Extract all entities named in the video and "
                "provide some context for each one."
            )
        )
        user_message = HumanMessage(
            content="TRANSCRIPT_TEXT:\n\n" + state["transcript_text"]
        )
        resp_msg = await self.extractor_llm.ainvoke([sys_prompt, user_message])
        print(f"extractor LLM response: {resp_msg}")