03_authentic_api_workflow.py

from llama_index.core.workflow import Workflow, step, Event, Context
from llama_index.core.workflow.events import (
    StartEvent,
    StopEvent,
    InputRequiredEvent,
    HumanResponseEvent,
)
from llama_index.core.workflow.handler import WorkflowHandler
from llama_index.llms.ollama import Ollama
from llama_index.core.llms import ChatMessage
from llama_index.core.prompts import PromptTemplate
import requests


# We'll create a workflow that will fetch data from authenticated API endpoint

model_name = "qwen2.5:0.5b"
llm = Ollama(model_name, request_timeout=120)

prompt_template = PromptTemplate(
    "You are an AI agent. Here is the user profile:\n{profile_str}\n\nUser Question:\n{question}"
)


class ProgressEvent(Event):
    def __init__(self, msg: str):
        super().__init__()
        self.msg = msg


class LoginEvent(InputRequiredEvent):
    pass


class CredentialSubmitEvent(HumanResponseEvent):
    credential: dict


class DataFetchEvent(Event):
    pass


class FinalResultEvent(Event):
    pass


class HumanInteractionWorkflow(Workflow):
    # 1st step, when user asked for the data
    # We'll check the session token is available or not
    # If token is not available then ask user to input it
    @step
    async def ask_question(
        self, ctx: Context, ev: StartEvent
    ) -> LoginEvent | DataFetchEvent:
        query = ev.query
        await ctx.set("original_query", query)

        session_token = await ctx.get("session_token", None)
        if session_token is None:
            ctx.write_event_to_stream(ProgressEvent(msg="user has no session token"))
            return LoginEvent(prefix="")
        return DataFetchEvent(prefix="", query=query)

    @step
    async def data_fetch(
        self, ctx: Context, ev: DataFetchEvent
    ) -> FinalResultEvent | LoginEvent:
        ctx.write_event_to_stream(ProgressEvent(msg="Fetching data from API"))
        token = await ctx.get("session_token")
        query = await ctx.get("original_query")
        try:
            headers = {"Authorization": f"Bearer {token}"}
            api_res = requests.get("https://dummyjson.com/auth/me", headers=headers)
            data = api_res.json()
            profie_str = "\n".join([f"{k}: {v}" for k, v in data.items()])

            formatted_prompt = prompt_template.format(
                profile_str=profie_str, question=query
            )
            payload = llm.chat([ChatMessage(role="user", content=formatted_prompt)])
            return FinalResultEvent(result=payload)
        except Exception as e:
            return FinalResultEvent(result=e.__str__())

    @step
    async def login(self, ctx: Context, ev: CredentialSubmitEvent) -> DataFetchEvent:
        ctx.write_event_to_stream(
            ProgressEvent(msg="Credential provided, trying to login")
        )
        credential = ev.credential

        try:
            api_res = requests.post("https://dummyjson.com/auth/login", credential)
            data = api_res.json()
            token = data["accessToken"]
            await ctx.set("username", credential["username"])
            await ctx.set("session_token", token)
            return DataFetchEvent()
        except Exception as e:
            # TODO handle re-try mechanism
            StopEvent(result=e)

        # validate input then call api

    @step
    async def final_result(self, ctx: Context, ev: FinalResultEvent) -> StopEvent:
        ctx.write_event_to_stream(ProgressEvent(msg="final result is preparing"))
        return StopEvent(result=ev.result)


async def executor(query: str):
    workflow = HumanInteractionWorkflow(timeout=None, verbose=False)
    try:
        handler: WorkflowHandler = workflow.run(query=query)

        async for event in handler.stream_events():
            if isinstance(event, ProgressEvent):
                print("[Progress]", event.msg)
            elif isinstance(event, LoginEvent):
                # query = event.query
                print("[System]: You are not logged in, please provide your credential")
                username = input("[You] Username: ")
                password = input("[You] Password: ")
                credential = {"username": username, "password": password}
                handler.ctx.send_event(
                    CredentialSubmitEvent(
                        prefix="", response="credential", credential=credential
                    )
                )

        final_result = await handler
        print("[Final Result]", final_result)
    except Exception as e:
        print("[Error]", e)


import asyncio


def main():
    query = input("[Ask Something]: ")
    asyncio.run(executor(query=query))


if __name__ == "__main__":
    main()


# Run: python hitl-workflow/03_authentic_api_workflow.py
# Then ask for your age, name, email, phone, hair color, city, address
# AI will ask for the credential then provide the credential
# AI will try to use the credential for login and then fetch your profile, based on the profile AI will response


# Data format of API
# ```json
# {
#   "username": "emilys",
#   "password": "emilyspass"
# }
# ```

# ```json
# {
#   "id": 1,
#   "firstName": "Emily",
#   "lastName": "Johnson",
#   "maidenName": "Smith",
#   "age": 28,
#   "gender": "female",
#   "email": "emily.johnson@x.dummyjson.com",
#   "phone": "+81 965-431-3024",
#   "username": "emilys",
#   "password": "emilyspass",
#   "birthDate": "1996-5-30",
#   "image": "https://dummyjson.com/icon/emilys/128",
#   "bloodGroup": "O-",
#   "height": 193.24,
#   "weight": 63.16,
#   "eyeColor": "Green",
#   "hair": {
#     "color": "Brown",
#     "type": "Curly"
#   },
#   "ip": "42.48.100.32",
#   "address": {
#     "address": "626 Main Street",
#     "city": "Phoenix",
#     "state": "Mississippi",
#     "stateCode": "MS",
#     "postalCode": "29112",
#     "coordinates": {
#       "lat": -77.16213,
#       "lng": -92.084824
#     },
#     "country": "United States"
#   },
#   "macAddress": "47:fa:41:18:ec:eb",
#   "university": "University of Wisconsin--Madison",
#   "bank": {
#     "cardExpire": "03/26",
#     "cardNumber": "9289760655481815",
#     "cardType": "Elo",
#     "currency": "CNY",
#     "iban": "YPUXISOBI7TTHPK2BR3HAIXL"
#   },
#   "company": {
#     "department": "Engineering",
#     "name": "Dooley, Kozey and Cronin",
#     "title": "Sales Manager",
#     "address": {
#       "address": "263 Tenth Street",
#       "city": "San Francisco",
#       "state": "Wisconsin",
#       "stateCode": "WI",
#       "postalCode": "37657",
#       "coordinates": {
#         "lat": 71.814525,
#         "lng": -161.150263
#       },
#       "country": "United States"
#     }
#   },
#   "ein": "977-175",
#   "ssn": "900-590-289",
#   "userAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.93 Safari/537.36",
#   "crypto": {
#     "coin": "Bitcoin",
#     "wallet": "0xb9fc2fe63b2a6c003f1c324c3bfa53259162181a",
#     "network": "Ethereum (ERC20)"
#   },
#   "role": "admin"
# }
# ```