browser.py

"""
Simple browser tool.

# Usage

Please start the backend browser server according to the instructions in the README.
"""

import re
from dataclasses import dataclass
from pprint import pprint

import requests
import streamlit as st

from .config import BROWSER_SERVER_URL
from .interface import ToolObservation


QUOTE_REGEX = re.compile(r"\[(\d+)†(.+?)\]")


@dataclass
class Quote:
    title: str
    url: str


# Quotes for displaying reference
if "quotes" not in st.session_state:
    st.session_state.quotes = {}

quotes: dict[str, Quote] = st.session_state.quotes


def map_response(response: dict) -> ToolObservation:
    # Save quotes for reference
    print("===BROWSER_RESPONSE===")
    pprint(response)
    role_metadata = response.get("roleMetadata")
    metadata = response.get("metadata")

    if role_metadata.split()[0] == "quote_result" and metadata:
        quote_id = QUOTE_REGEX.search(role_metadata.split()[1]).group(1)
        quote: dict[str, str] = metadata["metadata_list"][0]
        quotes[quote_id] = Quote(quote["title"], quote["url"])
    elif role_metadata == "browser_result" and metadata:
        for i, quote in enumerate(metadata["metadata_list"]):
            quotes[str(i)] = Quote(quote["title"], quote["url"])

    return ToolObservation(
        content_type=response.get("contentType"),
        text=response.get("result"),
        role_metadata=role_metadata,
        metadata=metadata,
    )


def tool_call(code: str, session_id: str) -> list[ToolObservation]:
    request = {
        "session_id": session_id,
        "action": code,
    }
    response = requests.post(BROWSER_SERVER_URL, json=request).json()
    return list(map(map_response, response))