Build your own customer support agent

import os
import time

from hydra_db import HydraDB
from hydra_db.errors import ConflictError
from hydra_db.helpers import build_string
from openai import OpenAI

client = HydraDB(token=os.environ["HYDRADB_API_KEY"])
llm = OpenAI()


def create_database(database, schema=None):
    """Create the database if it is new, then wait until it accepts context."""
    try:
        if schema:
            client.databases.create(database=database, database_metadata_schema=schema)
        else:
            client.databases.create(database=database)
    except ConflictError:
        pass  # The database already exists.
    while not client.databases.status(database=database).data.infra.ready_for_ingestion:
        time.sleep(5)


def source_ids(result):
    """Return the IDs an ingest accepted. Raise if it rejected any item."""
    for item in result.data.results:
        if item.error_code:
            raise RuntimeError(f"{item.filename or item.id}: {item.error_code} {item.error}")
    return [item.id for item in result.data.results]


def wait_indexed(database, collection, ids):
    """Block until every source has finished indexing. Indexing runs in the background."""
    while True:
        statuses = client.context.status(
            database=database, collection=collection, ids=ids
        ).data.statuses
        for status in statuses:
            if status.indexing_status == "errored":
                raise RuntimeError(f"{status.error_code}: {status.error_message}")
        if all(s.indexing_status == "completed" for s in statuses):
            return
        time.sleep(2)


def answer(question, result, instructions=""):
    """Send a query result to the model and return its answer."""
    completion = llm.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": "Answer using only the context provided. If the answer "
                "is not in the context, say you don't know. " + instructions,
            },
            {
                "role": "user",
                "content": f"Context:\n{build_string(result)}\n\nQuestion: {question}",
            },
        ],
    )
    return completion.choices[0].message.content

import json

from hydra import answer, client, create_database, source_ids, wait_indexed

DATABASE = "acme_support"
HELP_CENTER = "help_center"

create_database(DATABASE)


def customer_collection(customer_id):
    return f"customer_{customer_id}"

articles = {
    "refunds.md": "# Refunds\n\nRefunds are available within 30 days of purchase. "
                  "Annual plans are refunded pro rata.\n",
    "sso.md": "# Single sign-on\n\nSSO is available on the Business plan. "
              "Set it up under Settings, then Security.\n",
}
for name, body in articles.items():
    with open(name, "w") as f:
        f.write(body)

handles = [open(name, "rb") for name in articles]
result = client.context.ingest(
    type="knowledge",
    database=DATABASE,
    collection=HELP_CENTER,
    documents=[(name, h, "text/markdown") for name, h in zip(articles, handles)],
    document_metadata=json.dumps([{"id": name.removesuffix(".md")} for name in articles]),
)
for h in handles:
    h.close()
wait_indexed(DATABASE, HELP_CENTER, source_ids(result))

def remember(customer_id, turns):
    """turns is a list of (customer message, agent reply) pairs."""
    collection = customer_collection(customer_id)
    result = client.context.ingest(
        type="memory",
        database=DATABASE,
        collection=collection,
        memories=json.dumps([{
            "user_assistant_pairs": [{"user": u, "assistant": a} for u, a in turns],
            "infer": True,
        }]),
    )
    wait_indexed(DATABASE, collection, source_ids(result))


remember("4821", [
    ("We're on the annual Business plan. I asked for a refund on September 30.",
     "I've logged the request. Finance will confirm once the billing cycle closes."),
])

ROLES = {
    "triage": "Decide whether this is a billing or a technical question, and say why.",
    "billing": "Answer billing questions. State the policy, then what applies to this customer.",
    "technical": "Answer setup questions with numbered steps.",
}


def support_agent(role, customer_id, question):
    result = client.query(
        database=DATABASE,
        collections=[HELP_CENTER, customer_collection(customer_id)],
        query=question,
        type="all",
        mode="thinking",
    )
    return answer(question, result, ROLES[role])


question = "Any update on my refund, and how much will I get back?"
print(support_agent("triage", "4821", question))
print(support_agent("billing", "4821", question))