Example: Research Assistant

from beam import Bot, BotContext, BotLocation, Image
from pydantic import BaseModel

from dotenv import load_dotenv
import os

load_dotenv()

OPEN_AI_API_KEY = os.getenv("OPEN_AI_API_KEY")
SERPAPI_API_KEY = os.getenv("SERPAPI_API_KEY")
FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY")

NUMBER_OF_PRODUCT_REVIEWS_TO_SUMMARIZE = 3


# Define Locations (States)
class ProductName(BaseModel):
  product_name: str


class URL(BaseModel):
  url: str


class ReviewPage(BaseModel):
  review_page: str


# Create the Bot
bot = Bot(
  model="gpt-4o",
  api_key=OPEN_AI_API_KEY,
  locations=[
      BotLocation(marker=ProductName),
      BotLocation(marker=URL, expose=False),
      BotLocation(marker=ReviewPage, expose=False),
  ],
  description="This bot will take a product category as input (i.e. 'headphones') and search Google shopping for those products, lookup reviews for each of them, and then summarize the reviews of all products in a summary.",
)


# Transition 1: Retrieve 3 Google shopping URLs for each product
@bot.transition(
  inputs={ProductName: 1},
  outputs=[URL],
  description="Takes a product name and retrieves 5 Google shopping results",
  cpu=1,
  memory=128,
  image=Image(python_packages=["serpapi", "google-search-results", "python-dotenv"]),
)
def get_product_urls(context: BotContext, inputs):
  product_name = inputs[ProductName][0].product_name

  from serpapi import GoogleSearch

  params = {
      "engine": "google_shopping",
      "q": product_name,
      "api_key": SERPAPI_API_KEY,
  }

  search = GoogleSearch(params)
  results = search.get_dict()
  urls = results["shopping_results"][:NUMBER_OF_PRODUCT_REVIEWS_TO_SUMMARIZE]

  # Return a product url
  return {URL: [URL(url=url["product_link"]) for url in urls]}


# Transition 2: Scrape review page
@bot.transition(
  inputs={URL: 1},
  outputs=[ReviewPage],
  description="Scrapes the review page for each URL provided.",
  cpu=1,
  memory=128,
  image=Image(python_packages=["firecrawl-py", "python-dotenv"]),
  expose=False,
)
def scrape_reviews(context: BotContext, inputs):
  url = inputs[URL][0].url

  import json

  from firecrawl import FirecrawlApp

  app = FirecrawlApp(api_key=FIRECRAWL_API_KEY)

  # Scrape reviews from the product page
  scrape_result = app.scrape_url(url, params={"formats": ["markdown"]})
  print(scrape_result)

  return {ReviewPage: [ReviewPage(review_page=json.dumps(scrape_result))]}


# Transition 3: Summarize the product reviews
@bot.transition(
  inputs={ReviewPage: NUMBER_OF_PRODUCT_REVIEWS_TO_SUMMARIZE},
  outputs=[],
  description="Summarizes the reviews.",
  cpu=1,
  memory=128,
  image=Image().add_python_packages(["python-dotenv"]),
  expose=False,
)
def summarize_reviews(context: BotContext, inputs):
  try:
      all_review_pages = "\n".join(
          [input.review_page for input in inputs[ReviewPage]]
      )

      print(all_review_pages)

      prompt = f"""
          The following page contains markdown with a review for a product.
          Please highlight the key takeaways from all the reviews,
          and include 1-3 direct quotes from reviewers to support your points.
          In each quote, make sure to cite the name of the reviewer (if available).
          Make sure to include the name of the product, and a URL to buy it, in your response:
          {all_review_pages}
          """

      event = context.prompt(
          msg=prompt,
          timeout_seconds=30,
      )

      context.say("I've summarized product reviews like so: " + event.value)

      file_path = "/tmp/product-reviews.md"
      with open(file_path, "w") as f:
          f.write(event.value)

      if context.confirm(description="Do you want a sharable link to the summary?"):
          context.send_file(path=file_path, description="Summary of product reviews")

  except AttributeError:
      context.say("Review not found.")