Building a Data Retrieval Workflow for AI with Structured Output Libraries like Marvin and Instructor

def query_db(query):
 conn = sqlite3.connect('telco.db')
  try:
    return pd.read_sql_query(query, conn).to_markdown()
  finally:
 conn.close()

import marvin


marvin.settings.openai.api_key=OPENAI_API_KEY

from pydantic import BaseModel


class CanAnswerQuestion(BaseModel):
 reasoning: str
 can_answer: bool

@marvin.fn
def check_if_question_can_be_answered(database_description: str, question: str) -> CanAnswerQuestion:
  """You are a database reading bot that can answer users' questions using information from a database.

 Given the user's question, decide whether the question can be answered using the information in the database.

 Examples:
 reasoning: I can find the average revenue of customers with tenure over 24 months by averaging the Total Revenue column in the Billing table filtered by Tenure in Months > 24, can_answer: true
 reasoning: I can find customers who signed up during the last 12 months using the Tenure in Months column in the Billing table, can_answer: true
 reasoning: I can't answer how many customers churned previous year because the Churn table doesn't contain a year, can_answer: false"""

check_if_question_can_be_answered(DB_DESCRIPTION, "How many married seniors have the revenue higher than the average value for all users?")

@marvin.fn
def write_sql_query(database_description: str, question: str, plan: str) -> str:
  """You are a database reading bot that can answer users' questions using information from a database.

 In the previous step, you have prepared a query plan.

 Return an SQL query with no preamble or explanation. Don't include any markdown characters or quotation marks around the query."""

write_sql_query(
    DB_DESCRIPTION,
    "How many married seniors have the revenue higher than the average value for all users?",
    "I can find how many married senior citizens have revenue higher than the average value for all users by filtering the Customer table for 'Married' equal to 'Yes' and 'Senior_Citizen' equal to 'Yes' and then comparing their 'Total_Revenue' from the Billing table to the average 'Total_Revenue' of all users."
)

query_db("SELECT COUNT(*) FROM Customer c JOIN Billing b ON c.Customer_ID = b.Customer_ID WHERE c.Married = 'Yes' AND c.Senior_Citizen = 'Yes' AND b.Total_Revenue > (SELECT AVG(Total_Revenue) FROM Billing)")

@marvin.fn
def write_answer(question: str, plan: str, sql_query: str, sql_result: str) -> str:
  """You are a database reading bot that can answer users' questions using information from a database.

 In the previous step, you have prepared a query plan, generated a query, and retrieved the data from the database.

 Return a text answering the user's question using the provided data."""

@marvin.fn
def explain_why_cannot_answer(question: str, problem_explanation: str) -> str:
  """You are a database reading bot that can answer users' questions using information from a database.

 You cannot answer the user's questions because of the problem with data.

 Explain the issue to the user and apologize for the inconvenience."""

def answer_user_question(question: str) -> str:
 can_answer = check_if_question_can_be_answered(DB_DESCRIPTION, question)

  if can_answer.can_answer:
 plan = can_answer.reasoning
 sql_query = write_sql_query(DB_DESCRIPTION, question, plan)
 query_result = query_db(sql_query)
    return write_answer(question, plan, sql_query, query_result)
  else:
    return explain_why_cannot_answer(question, can_answer.reasoning)

answer_user_question("Count customers by zip code. Return the 5 most common zip codes")

marvin.generate(n=5, instructions="Generate subjects for an email about using structured outputs in LLMs")

marvin.classify(
    "Revolutionize Your AI Workflow with Structured Output Implementation", labels=["clickbait", "not_clickbait"]
)
# clickbait

marvin.extract("Great hotel, nice and clean room, good location, and convenient parking. Staff were very nice, friendly and helpful too. We stayed here for a night and will definitely stay again.", instructions="get any amenities mentioned in the review")

# ['clean room', 'good location', 'convenient parking']

import instructor
from openai import OpenAI


openai_client = OpenAI(api_key=OPENAI_API_KEY)
client = instructor.from_openai(openai_client)

def check_if_question_can_be_answered(question: str) -> CanAnswerQuestion:
 prompt = f"""You are a database reading bot that can answer users' questions using information from a database. \n

 Given the user's question, decide whether the question can be answered using the information in the database. \n\n

    {DB_DESCRIPTION}

 Examples:

 reasoning: I can find the average revenue of customers with tenure over 24 months by averaging the Total Revenue column in the Billing table filtered by Tenure in Months > 24, can_answer: true
 reasoning: I can find customers who signed up during the last 12 month using the Tenure in Months column in the Billing table, can_answer: true
 reasoning: I can't answer how many customers churned last year because the Churn table doesn't contain a year, can_answer: false"""

  return client.chat.completions.create(
    model="gpt-3.5-turbo",
    response_model=CanAnswerQuestion,
    messages=[
 {"role": "system", "content": prompt},
 {"role": "user", "content": question}
 ],
)

check_if_question_can_be_answered("How many married seniors have the revenue higher than the average value for all users?")

def write_sql_query(question: str, plan: str) -> str:
 prompt = f"""You are a database reading bot that can answer users' questions using information from a database. \n

  {DB_DESCRIPTION}

 In the previous step, you have prepared a query plan.

 Return an SQL query with no preamble or explanation. Don't include any markdown characters or quotation marks around the query."""

  return client.chat.completions.create(
    model="gpt-3.5-turbo",
    response_model=str,
    messages=[
 {"role": "system", "content": prompt},
 {"role": "user", "content": "Query plan: " + plan}
 ],
 )