Agent for cleaning up named entities in YouTube video transcripts.

demo-enrichment-agent.py

...
from pydantic import BaseModel, Field

class NamedEntity(BaseModel):
    entity_name: str = Field(description="A short name for the entity.")
    entity_context: str = Field(
        description="A **brief** phrase (≤ 12 words) that is suitable for web search."
    )

class NamedEntities(BaseModel):
    named_entities: list[NamedEntity]

...

class DemoEnrichmentAgent:
    extractor_llm: Runnable

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