/examples/flows/chat/chat-with-wikipedia-maf/workflow.py

import asyncio
import os
import re
import random
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from functools import partial

import bs4
import requests
from dotenv import load_dotenv
from typing_extensions import Never

from agent_framework import Agent, Executor, WorkflowBuilder, WorkflowContext, handler
from agent_framework.openai import OpenAIChatClient

load_dotenv()

# ---------------------------------------------------------------------------
# Prompt templates (from .jinja2 files)
# ---------------------------------------------------------------------------

EXTRACT_QUERY_INSTRUCTIONS = """\
You are an AI assistant reading the transcript of a conversation between an AI and a human. \
Given an input question and conversation history, infer user real intent.

The conversation history is provided just in case of a context \
(e.g. "What is this?" where "this" is defined in previous conversation).

Return the output as query used for next round user message."""

AUGMENTED_CHAT_INSTRUCTIONS = """\
You are a chatbot having a conversation with a human.
Given the following extracted parts of a long document and a question, \
create a final answer with references ("SOURCES").
If you don't know the answer, just say that you don't know. Don't try to make up an answer.
ALWAYS return a "SOURCES" part in your answer."""

# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class ChatInput:
    question: str
    chat_history: list = field(default_factory=list)


@dataclass
class QueryWithContext:
    question: str
    chat_history: list
    contexts: str


# ---------------------------------------------------------------------------
# Wikipedia helper functions (from get_wiki_url.py, search_result_from_url.py,
# process_search_result.py)
# ---------------------------------------------------------------------------

_session = requests.Session()


def _decode_str(string):
    return string.encode().decode("unicode-escape").encode("latin1").decode("utf-8")


def _remove_nested_parentheses(string):
    pattern = r"\([^()]+\)"
    while re.search(pattern, string):
        string = re.sub(pattern, "", string)
    return string


def get_wiki_url(entity: str, count: int = 2) -> list[str]:
    url = f"https://en.wikipedia.org/w/index.php?search={entity}"
    url_list: list[str] = []
    try:
        headers = {
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                "(KHTML, like Gecko) Chrome/113.0.0.0 Safari/537.36 Edg/113.0.1774.35"
            )
        }
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            soup = bs4.BeautifulSoup(response.text, "html.parser")
            mw_divs = soup.find_all("div", {"class": "mw-search-result-heading"})
            if mw_divs:
                result_titles = [_decode_str(div.get_text().strip()) for div in mw_divs]
                result_titles = [_remove_nested_parentheses(t) for t in result_titles]
                url_list.extend(
                    f"https://en.wikipedia.org/w/index.php?search={t}" for t in result_titles
                )
            else:
                page_content = [p.get_text().strip() for p in soup.find_all("p") + soup.find_all("ul")]
                if any("may refer to:" in p for p in page_content):
                    url_list.extend(get_wiki_url("[" + entity + "]"))
                else:
                    url_list.append(url)
        return url_list[:count]
    except Exception as e:
        print(f"Get url failed with error: {e}")
        return url_list


def _get_page_sentence(page: str, count: int = 10) -> str:
    paragraphs = [p.strip() for p in page.split("\n") if p.strip()]
    sentences: list[str] = []
    for p in paragraphs:
        sentences += p.split(". ")
    sentences = [s.strip() + "." for s in sentences if s.strip()]
    return " ".join(sentences[:count])


def _fetch_text_content_from_url(url: str, count: int = 10) -> tuple[str, str]:
    try:
        headers = {
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
                "(KHTML, like Gecko) Chrome/113.0.0.0 Safari/537.36 Edg/113.0.1774.35"
            )
        }
        delay = random.uniform(0, 0.5)
        time.sleep(delay)
        response = _session.get(url, headers=headers)
        if response.status_code == 200:
            soup = bs4.BeautifulSoup(response.text, "html.parser")
            page_content = [p.get_text().strip() for p in soup.find_all("p") + soup.find_all("ul")]
            page = ""
            for content in page_content:
                if len(content.split(" ")) > 2:
                    page += _decode_str(content)
                if not content.endswith("\n"):
                    page += "\n"
            return (url, _get_page_sentence(page, count=count))
        return (url, "No available content")
    except Exception as e:
        print(f"Get url failed with error: {e}")
        return (url, "No available content")


def search_result_from_url(url_list: list[str], count: int = 10) -> list[tuple[str, str]]:
    results: list[tuple[str, str]] = []
    fn = partial(_fetch_text_content_from_url, count=count)
    with ThreadPoolExecutor(max_workers=5) as executor:
        for r in executor.map(fn, url_list):
            results.append(r)
    return results


def process_search_result(search_result: list[tuple[str, str]]) -> str:
    context = []
    for url, content in search_result:
        context.append({"Content": content, "Source": url})
    return "\n\n".join(f"Content: {c['Content']}\nSource: {c['Source']}" for c in context)


# ---------------------------------------------------------------------------
# Executors
# ---------------------------------------------------------------------------


class InputExecutor(Executor):
    """Passes the ChatInput through to the next stage."""

    @handler
    async def receive(self, chat_input: ChatInput, ctx: WorkflowContext[ChatInput]) -> None:
        await ctx.send_message(chat_input)


class ExtractQueryExecutor(Executor):
    """LLM node: extracts a search query from the user question + chat history."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        client = OpenAIChatClient(
            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
            model=os.environ.get("AZURE_OPENAI_DEPLOYMENT", "gpt-35-turbo"),
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
        )
        self._agent = Agent(
            client=client,
            name="ExtractQueryAgent",
            instructions=EXTRACT_QUERY_INSTRUCTIONS,
        )

    @handler
    async def extract(self, chat_input: ChatInput, ctx: WorkflowContext[ChatInput]) -> None:
        # Build the prompt with history
        parts: list[str] = []
        for turn in chat_input.chat_history:
            parts.append(f"Human: {turn['inputs']['question']}")
            parts.append(f"AI: {turn['outputs']['answer']}")
        parts.append(f"Human: {chat_input.question}")
        parts.append("\nOutput:")

        response = await self._agent.run("\n".join(parts))
        # Replace the question with the extracted query for downstream
        await ctx.send_message(
            ChatInput(
                question=response.text.strip(),
                chat_history=chat_input.chat_history,
            )
        )


class WikiSearchExecutor(Executor):
    """Python node: searches Wikipedia and retrieves page content."""

    @handler
    async def search(self, chat_input: ChatInput, ctx: WorkflowContext[QueryWithContext]) -> None:
        entity = chat_input.question
        urls = get_wiki_url(entity, count=2)
        search_results = search_result_from_url(urls, count=10)
        contexts = process_search_result(search_results)
        await ctx.send_message(
            QueryWithContext(
                question=chat_input.question,
                chat_history=chat_input.chat_history,
                contexts=contexts,
            )
        )


class AugmentedChatExecutor(Executor):
    """LLM node: answers the question using Wikipedia context."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        client = OpenAIChatClient(
            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
            model=os.environ.get("AZURE_OPENAI_DEPLOYMENT", "gpt-35-turbo"),
            api_key=os.environ["AZURE_OPENAI_API_KEY"],
        )
        self._agent = Agent(
            client=client,
            name="AugmentedChatAgent",
            instructions=AUGMENTED_CHAT_INSTRUCTIONS,
        )

    @handler
    async def answer(self, qc: QueryWithContext, ctx: WorkflowContext[Never, str]) -> None:
        parts: list[str] = [qc.contexts, ""]
        for turn in qc.chat_history:
            parts.append(f"User: {turn['inputs']['question']}")
            parts.append(f"Assistant: {turn['outputs']['answer']}")
        parts.append(qc.question)

        response = await self._agent.run("\n".join(parts))
        await ctx.yield_output(response.text)


# ---------------------------------------------------------------------------
# Workflow: input → extract_query → wiki_search → augmented_chat → output
# ---------------------------------------------------------------------------

def create_workflow():
    """Create a fresh workflow instance.

    MAF workflows do not support concurrent execution, so each
    concurrent caller needs its own workflow instance.
    """
    _input = InputExecutor(id="input")
    _extract = ExtractQueryExecutor(id="extract_query")
    _wiki = WikiSearchExecutor(id="wiki_search")
    _chat = AugmentedChatExecutor(id="augmented_chat")
    return (
        WorkflowBuilder(name="ChatWithWikipediaWorkflow", start_executor=_input)
        .add_edge(_input, _extract)
        .add_edge(_extract, _wiki)
        .add_edge(_wiki, _chat)
        .build()
    )


async def main():
    workflow = create_workflow()
    result = await workflow.run(ChatInput(question="What is ChatGPT?"))
    print("Answer:", result.get_outputs()[0])


if __name__ == "__main__":
    asyncio.run(main())