Coda's Help Desk with and without RAG

%pip install -U autoevals braintrust requests openai lancedb markdownify

QA_GEN_MODEL = "gpt-3.5-turbo"
QA_ANSWER_MODEL = "gpt-3.5-turbo"
QA_GRADING_MODEL = "gpt-4"
RELEVANCE_MODEL = "gpt-3.5-turbo"
 
NUM_SECTIONS = 20
NUM_QA_PAIRS = 20  # Increase this number to test at a larger scale

import asyncio
import os
import re
import time
 
import autoevals
import braintrust
import markdownify
import openai
import requests
 
 
data = requests.get(
    "https://gist.githubusercontent.com/wong-codaio/b8ea0e087f800971ca5ec9eef617273e/raw/39f8bd2ebdecee485021e20f2c1d40fd649a4c77/articles.json"
).json()
 
markdown_docs = [{"id": row["id"], "markdown": markdownify.markdownify(row["body"])} for row in data]
 
i = 0
markdown_sections = []
for markdown_doc in markdown_docs:
    sections = re.split(r"(.*\n=+\n)", markdown_doc["markdown"])
    current_section = ""
    for section in sections:
        if not section.strip():
            continue
 
        if re.match(r".*\n=+\n", section):
            current_section = section
        else:
            section = current_section + section
            markdown_sections.append({"doc_id": markdown_doc["id"], "section_id": i, "markdown": section.strip()})
            current_section = ""
            i += 1
 
print(f"Downloaded {len(markdown_sections)} Markdown sections. Here are the first 3:")
markdown_sections[:3]

client = braintrust.wrap_openai(openai.AsyncOpenAI(
    base_url="https://braintrustproxy.com/v1",
    default_headers={"x-bt-use-cache": "always"},
    api_key=os.environ.get("OPENAI_API_KEY", "Your OPENAI_API_KEY here"),
))

import json
from typing import List
from pydantic import BaseModel, Field
 
 
class QAPair(BaseModel):
    questions: List[str] = Field(
        ..., description="List of questions, all with the same meaning but worded differently"
    )
    answer: str = Field(..., description="Answer")
 
 
class QAPairs(BaseModel):
    pairs: List[QAPair] = Field(..., description="List of question/answer pairs")
 
 
async def produce_candidate_questions(row):
    response = await client.chat.completions.create(
        model=QA_GEN_MODEL,
        messages=[{"role": "user", "content": f"""\
Please generate 8 question/answer pairs from the following text. For each question, suggest
2 different ways of phrasing the question, and provide a unique answer.
 
Content:
 
{row['markdown']}
""",
        }],
        functions=[
            {
                "name": "propose_qa_pairs",
                "description": "Propose some question/answer pairs for a given document",
                "parameters": QAPairs.schema(),
            }
        ],
    )
 
    pairs = QAPairs(**json.loads(response.choices[0].message.function_call.arguments))
    return pairs.pairs

all_candidates_tasks = [
    asyncio.create_task(produce_candidate_questions(a)) for a in markdown_sections[:NUM_SECTIONS]
]
all_candidates = [await f for f in all_candidates_tasks]
 
data = []
row_id = 0
for row, doc_qa in zip(markdown_sections[:NUM_SECTIONS], all_candidates):
    for i, qa in enumerate(doc_qa):
        for j, q in enumerate(qa.questions):
            data.append(
                {
                    "input": q,
                    "expected": qa.answer,
                    "metadata": {
                        "document_id": row["doc_id"],
                        "section_id": row["section_id"],
                        "question_idx": i,
                        "answer_idx": j,
                        "id": row_id,
                        "split": "test" if j == len(qa.questions) - 1 and j > 0 else "train",
                    },
                }
            )
            row_id += 1
 
print(f"Generated {len(data)} QA pairs. Here are the first 10...")
for x in data[:10]:
    print(x)

async def simple_qa(input):
    completion = await client.chat.completions.create(
        model=QA_ANSWER_MODEL,
        messages=[{"role": "user", "content": f"""\
Please answer the following question:
 
Question: {input}
""",
        }],
    )
    return completion.choices[0].message.content
 
 
await braintrust.Eval(
    name="Coda Help Desk Cookbook",
    experiment_name="No RAG",
    data=data[:NUM_QA_PAIRS],
    task=simple_qa,
    scores=[autoevals.Factuality(model=QA_GRADING_MODEL)],
)

import tempfile
import lancedb
 
 
tempdir = tempfile.TemporaryDirectory()
LANCE_DB_PATH = os.path.join(tempdir.name, "docs-lancedb")
 
 
@braintrust.traced
async def embed_text(text):
    params = dict(input=text, model="text-embedding-ada-002")
    response = await client.embeddings.create(**params)
    embedding = response.data[0].embedding
 
    braintrust.current_span().log(
        metrics={"tokens": response.usage.total_tokens, "prompt_tokens": response.usage.prompt_tokens},
        metadata={"model": response.model},
        input=text,
        output=embedding,
    )
    
    return embedding
 
 
embedding_tasks = [asyncio.create_task(embed_text(row["markdown"])) for row in markdown_sections[:NUM_SECTIONS]]
embeddings = [await f for f in embedding_tasks]
 
db = lancedb.connect(LANCE_DB_PATH)
 
try:
    db.drop_table("sections")
except:
    pass
 
table = db.create_table(
    "sections",
    data=[
        {"doc_id": row["doc_id"], "section_id": row["section_id"], "vector": embedding}
        for (row, embedding) in zip(markdown_sections[:NUM_SECTIONS], embeddings)
    ],
)

TOP_K = 2
 
 
@braintrust.traced
async def relevance_score(query, document):
    response = await client.chat.completions.create(
        model=RELEVANCE_MODEL,
        messages=[{"role": "user", "content": f"""\
Consider the following query and a document
 
Query:
{query}
 
Document:
{document}
 
 
Please score the relevance of the document to a query, on a scale of 0 to 1.
""",
        }],
        functions=[
            {
                "name": "has_relevance",
                "description": "Declare the relevance of a document to a query",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "score": {"type": "number"},
                    },
                },
            }
        ],
    )
 
    arguments = response.choices[0].message.function_call.arguments
    result = json.loads(arguments)
 
    braintrust.current_span().log(
        input={"query": query, "document": document},
        output=result,
    )
 
    return result["score"]
 
 
async def retrieval_qa(input):
    embedding = await embed_text(input)
 
    with braintrust.current_span().start_span(name="vector search", input=input) as span:
        result = table.search(embedding).limit(TOP_K + 3).to_arrow().to_pylist()
        docs = [markdown_sections[i["section_id"]]["markdown"] for i in result]
 
        relevance_scores = []
        for doc in docs:
            relevance_scores.append(await relevance_score(input, doc))
 
        span.log(
            output=[
                {"doc": markdown_sections[r["section_id"]]["markdown"], "distance": r["_distance"]} for r in result
            ],
            metadata={"top_k": TOP_K, "retrieval": result},
            scores={
                "avg_relevance": sum(relevance_scores) / len(relevance_scores),
                "min_relevance": min(relevance_scores),
                "max_relevance": max(relevance_scores),
            },
        )
 
    context = "\n------\n".join(docs[:TOP_K])
    completion = await client.chat.completions.create(
        model=QA_ANSWER_MODEL,
        messages=[{"role": "user", "content": f"""\
Given the following context
 
{context}
 
Please answer the following question:
 
Question: {input}
""",
        }],
    )
    
    return completion.choices[0].message.content

await braintrust.Eval(
    name="Coda Help Desk Cookbook",
    experiment_name=f"RAG TopK={TOP_K}",
    data=data[:NUM_QA_PAIRS],
    task=retrieval_qa,
    scores=[autoevals.Factuality(model=QA_GRADING_MODEL)],
)