💬 Chat with Webpage

from burr.core import action, State, ApplicationBuilder
from scrapegraph_py import ScrapeGraphAI, MarkdownFormatConfig
import lancedb
from lancedb.pydantic import LanceModel, Vector
import openai
import tiktoken
from typing import List, Optional

# Schema for storing text chunks
class TextDocument(LanceModel):
    url: str
    position: int
    text: str
    vector: Vector(dim=1536)  # OpenAI embedding dimensions

# Action to fetch and convert webpage to markdown
@action(reads=[], writes=["markdown_content"])
def fetch_webpage(state: State, webpage_url: str) -> State:
    sgai = ScrapeGraphAI()  # reads SGAI_API_KEY from env
    res = sgai.scrape(webpage_url, formats=[MarkdownFormatConfig()])
    markdown = res.data.results["markdown"]["data"][0] if res.status == "success" else ""
    return state.update(markdown_content=markdown)

# Action to embed and store chunks
@action(reads=["markdown_content"], writes=[])
def embed_and_store(state: State, webpage_url: str) -> State:
    chunks = get_text_chunks(state["markdown_content"])
    con = lancedb.connect("./webpages")
    table = con.create_table("chunks", schema=TextDocument)
    table.add([{
        "text": chunk,
        "url": webpage_url,
        "position": i
    } for i, chunk in enumerate(chunks)])
    return state

# Action to answer questions
@action(reads=[], writes=["llm_answer"])
def ask_question(state: State, user_query: str) -> State:
    chunks_table = lancedb.connect("./webpages").open_table("chunks")
    relevant_chunks = chunks_table.search(user_query).limit(3).to_list()
    
    response = openai.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": f"Answer based on: {relevant_chunks}"},
            {"role": "user", "content": user_query}
        ]
    )
    return state.update(llm_answer=response.choices[0].message.content)

{
    "question": "Who are the founders of ScrapeGraphAI?",
    "answer": "The founders of ScrapeGraphAI are:\n\n1. Marco Perini - Founder & Technical Lead\n2. Marco Vinciguerra - Founder & Software Engineer\n3. Lorenzo Padoan - Founder & Product Engineer"
}