tiantian-paris/FRM_Study_chatbot

1from __future__ import annotations as _annotations2 3import json4import os5from dataclasses import dataclass6from typing import Any7 8import gradio as gr9from dotenv import load_dotenv10from httpx import AsyncClient11from pydantic_ai import Agent, ModelRetry, RunContext12from pydantic_ai.messages import ModelStructuredResponse, ModelTextResponse, ToolReturn13 14load_dotenv()15 16 17@dataclass18class Deps:19    client: AsyncClient20    weather_api_key: str | None21    geo_api_key: str | None22 23 24weather_agent = Agent(25    "openai:gpt-4o",26    system_prompt="You are an expert packer. A user will ask you for help packing for a trip given a destination. Use your weather tools to provide a concise and effective packing list. Also ask follow up questions if neccessary.",27    deps_type=Deps,28    retries=2,29)30 31 32@weather_agent.tool33async def get_lat_lng(34    ctx: RunContext[Deps], location_description: str35) -> dict[str, float]:36    """Get the latitude and longitude of a location.37    Args:38        ctx: The context.39        location_description: A description of a location.40    """41    if ctx.deps.geo_api_key is None:42        # if no API key is provided, return a dummy response (London)43        return {"lat": 51.1, "lng": -0.1}44 45    params = {46        "q": location_description,47        "api_key": ctx.deps.geo_api_key,48    }49    r = await ctx.deps.client.get("https://geocode.maps.co/search", params=params)50    r.raise_for_status()51    data = r.json()52 53    if data:54        return {"lat": data[0]["lat"], "lng": data[0]["lon"]}55    else:56        raise ModelRetry("Could not find the location")57 58 59@weather_agent.tool60async def get_weather(ctx: RunContext[Deps], lat: float, lng: float) -> dict[str, Any]:61    """Get the weather at a location.62    Args:63        ctx: The context.64        lat: Latitude of the location.65        lng: Longitude of the location.66    """67    if ctx.deps.weather_api_key is None:68        # if no API key is provided, return a dummy response69        return {"temperature": "21 °C", "description": "Sunny"}70 71    params = {72        "apikey": ctx.deps.weather_api_key,73        "location": f"{lat},{lng}",74        "units": "metric",75    }76    r = await ctx.deps.client.get(77        "https://api.tomorrow.io/v4/weather/realtime", params=params78    )79    r.raise_for_status()80    data = r.json()81 82    values = data["data"]["values"]83    # https://docs.tomorrow.io/reference/data-layers-weather-codes84    code_lookup = {85        1000: "Clear, Sunny",86        1100: "Mostly Clear",87        1101: "Partly Cloudy",88        1102: "Mostly Cloudy",89        1001: "Cloudy",90        2000: "Fog",91        2100: "Light Fog",92        4000: "Drizzle",93        4001: "Rain",94        4200: "Light Rain",95        4201: "Heavy Rain",96        5000: "Snow",97        5001: "Flurries",98        5100: "Light Snow",99        5101: "Heavy Snow",100        6000: "Freezing Drizzle",101        6001: "Freezing Rain",102        6200: "Light Freezing Rain",103        6201: "Heavy Freezing Rain",104        7000: "Ice Pellets",105        7101: "Heavy Ice Pellets",106        7102: "Light Ice Pellets",107        8000: "Thunderstorm",108    }109    return {110        "temperature": f'{values["temperatureApparent"]:0.0f}°C',111        "description": code_lookup.get(values["weatherCode"], "Unknown"),112    }113 114 115TOOL_TO_DISPLAY_NAME = {"get_lat_lng": "Geocoding API", "get_weather": "Weather API"}116 117client = AsyncClient()118weather_api_key = os.getenv("WEATHER_API_KEY")119# create a free API key at https://geocode.maps.co/120geo_api_key = os.getenv("GEO_API_KEY")121deps = Deps(client=client, weather_api_key=weather_api_key, geo_api_key=geo_api_key)122 123 124async def stream_from_agent(prompt: str, chatbot: list[dict], past_messages: list):125    chatbot.append({"role": "user", "content": prompt})126    yield gr.Textbox(interactive=False, value=""), chatbot, gr.skip()127    async with weather_agent.run_stream(128        prompt, deps=deps, message_history=past_messages129    ) as result:130        for message in result.new_messages():131            past_messages.append(message)132            if isinstance(message, ModelStructuredResponse):133                for call in message.calls:134                    gr_message = {135                        "role": "assistant",136                        "content": "",137                        "metadata": {138                            "title": f"### 🛠️ Using {TOOL_TO_DISPLAY_NAME[call.tool_name]}",139                            "id": call.tool_id,140                        },141                    }142                    chatbot.append(gr_message)143            if isinstance(message, ToolReturn):144                for gr_message in chatbot:145                    if gr_message.get("metadata", {}).get("id", "") == message.tool_id:146                        gr_message["content"] = f"Output: {json.dumps(message.content)}"147            yield gr.skip(), chatbot, gr.skip()148        chatbot.append({"role": "assistant", "content": ""})149        async for message in result.stream_text():150            chatbot[-1]["content"] = message151            yield gr.skip(), chatbot, gr.skip()152        data = await result.get_data()153        past_messages.append(ModelTextResponse(content=data))154        yield gr.Textbox(interactive=True), gr.skip(), past_messages155 156 157async def handle_retry(chatbot, past_messages: list, retry_data: gr.RetryData):158    new_history = chatbot[: retry_data.index]159    previous_prompt = chatbot[retry_data.index]["content"]160    past_messages = past_messages[: retry_data.index]161    async for update in stream_from_agent(previous_prompt, new_history, past_messages):162        yield update163 164 165def undo(chatbot, past_messages: list, undo_data: gr.UndoData):166    new_history = chatbot[: undo_data.index]167    past_messages = past_messages[: undo_data.index]168    return chatbot[undo_data.index]["content"], new_history, past_messages169 170 171def select_data(message: gr.SelectData) -> str:172    return message.value["text"]173 174 175with gr.Blocks() as demo:176    gr.HTML(177        """178<div style="display: flex; justify-content: center; align-items: center; gap: 2rem; padding: 1rem; width: 100%">179    <img src="https://ai.pydantic.dev/img/logo-white.svg" style="max-width: 200px; height: auto">180    <div>181        <h1 style="margin: 0 0 1rem 0">Vacation Packing Assistant</h1>182        <h3 style="margin: 0 0 0.5rem 0">183            This assistant will help you pack for your vacation. Enter your destination and it will provide you with a concise packing list based on the weather forecast.184        </h3>185        <h3 style="margin: 0">186            Feel free to ask for help with any other questions you have about your trip!187        </h3>188    </div>189</div>190"""191    )192    past_messages = gr.State([])193    chatbot = gr.Chatbot(194        label="Packing Assistant",195        type="messages",196        avatar_images=(None, "https://ai.pydantic.dev/img/logo-white.svg"),197        examples=[198            {"text": "I am going to Paris for the holidays, what should I pack?"},199            {"text": "I am going to Tokyo this week."},200        ],201    )202    with gr.Row():203        prompt = gr.Textbox(204            lines=1,205            show_label=False,206            placeholder="I am planning a trip to Miami, what should I pack?",207        )208    generation = prompt.submit(209        stream_from_agent,210        inputs=[prompt, chatbot, past_messages],211        outputs=[prompt, chatbot, past_messages],212    )213    chatbot.example_select(select_data, None, [prompt])214    chatbot.retry(215        handle_retry, [chatbot, past_messages], [prompt, chatbot, past_messages]216    )217    chatbot.undo(undo, [chatbot, past_messages], [prompt, chatbot, past_messages])218 219 220if __name__ == "__main__":221    demo.launch()