Friday, March 1, 2024

RAG techniques: Function calling for more structured retrieval

Retrieval Augmented Generation (RAG) is a popular technique to get LLMs to provide answers that are grounded in a data source. When we use RAG, we use the user's question to search a knowledge base (like Azure AI Search), then pass along both the question and the relevant content to the LLM (gpt-3.5-turbo or gpt-4), with a directive to answer only according to the sources. In psuedo-code:

user_query = "what's in the Northwind Plus plan?"
user_query_vector = create_embedding(user_query, "ada-002")
results = search(user_query, user_query_vector)
response = create_chat_completion(system_prompt, user_query, results)

If the search function can find the right results in the index (assuming the answer is somewhere in the index), then the LLM can typically do a pretty good job of synthesizing the answer from the sources.

Unstructured queries

This simple RAG approach works best for "unstructured queries", like:

  • What's in the Northwind Plus plan?
  • What are the expectations of a product manager?
  • What benefits are provided by the company?

When using Azure AI Search as the knowledge base, the search call will perform both a vector and keyword search, finding all the relevant document chunks that match the keywords and concepts in the query.

Structured queries

But you may find that users are instead asking more "structured" queries, like:

  • Summarize the document called "perksplus.pdf"
  • What are the topics in documents by Pamela Fox?
  • Key points in most recent uploaded documents

We can think of them as structured queries, because they're trying to filter on specific metadata about a document. You could imagine a world where you used a syntax to specify that metadata filtering, like:

  • Summarize the document title:perksplus.pdf
  • Topics in documents author:PamelaFox
  • Key points time:2weeks

We don't want to actually introduce a query syntax to a a RAG chat application if we don't need to, since only power users tend to use specialized query syntax, and we'd ideally have our RAG just do the right thing in that situation.

Using function calling in RAG

Fortunately, we can use the OpenAI function-calling feature to recognize that a user's query would benefit from a more structured search, and perform that search instead.

If you've never used function calling before, it's an alternative way of asking an OpenAI GPT model to respond to a chat completion request. In addition to sending our usual system prompt, chat history, and user message, we also send along a list of possible functions that could be called to answer the question. We can define those in JSON or as a Pydantic model dumped to JSON. Then, when the response comes back from the model, we can see what function it decided to call, and with what parameters. At that point, we can actually call that function, if it exists, or just use that information in our code in some other way.

To use function calling in RAG, we first need to introduce an LLM pre-processing step to handle user queries, as I described in my previous blog post. That will give us an opportunity to intercept the query before we even perform the search step of RAG.

For that pre-processing step, we can start off with a function to handle the general case of unstructured queries:

tools: List[ChatCompletionToolParam] = [
    {
        "type": "function",
        "function": {
            "name": "search_sources",
            "description": "Retrieve sources from the Azure AI Search index",
            "parameters": {
                "type": "object",
                "properties": {
                    "search_query": {
                        "type": "string",
                        "description": "Query string to retrieve documents from azure search eg: 'Health care plan'",
                    }
                },
                "required": ["search_query"],
            },
        },
    }
]

Then we send off a request to the chat completion API, letting it know it can use that function.

chat_completion: ChatCompletion = self.openai_client.chat.completions.create(
    messages=messages,
    model=model,
    temperature=0.0,
    max_tokens=100,
    n=1,
    tools=tools,
    tool_choice="auto",
)

When the response comes back, we process it to see if the model decided to call the function, and extract the search_query parameter if so.

response_message = chat_completion.choices[0].message

if response_message.tool_calls:
    for tool in response_message.tool_calls:
        if tool.type != "function":
            continue
        function = tool.function
        if function.name == "search_sources":
            arg = json.loads(function.arguments)
            search_query = arg.get("search_query", self.NO_RESPONSE)

If the model didn't include the function call in its response, that's not a big deal as we just fall back to using the user's original query as the search query. We proceed with the rest of the RAG flow as usual, sending the original question with whatever results came back in our final LLM call.

Adding more functions for structured queries

Now that we've introduced one function into the RAG flow, we can more easily add additional functions to recognize structured queries. For example, this function recognizes when a user wants to search by a particular filename:

{
    "type": "function",
    "function": {
        "name": "search_by_filename",
        "description": "Retrieve a specific filename from the Azure AI Search index",
        "parameters": {
            "type": "object",
            "properties": {
                "filename": {
                    "type": "string",
                    "description": "The filename, like 'PerksPlus.pdf'",
                }
            },
            "required": ["filename"],
        },
    },
},

We need to extend the function parsing code to extract the filename argument:

if function.name == "search_by_filename":
    arg = json.loads(function.arguments)
    filename = arg.get("filename", "")
    filename_filter = filename

Then we can decide how to use that filename filter. In the case of Azure AI search, I build a filter that checks that a particular index field matches the filename argument, and pass that to my search call. If using a relational database, it'd become an additional WHERE clause.

Simply by adding that function, I was able to get much better answers to questions in my RAG app like 'Summarize the document called "perksplus.pdf"', since my search results were truly limited to chunks from that file. You can see my full code changes to add this function to our RAG starter app repo in this PR.

Considerations

This can be a very powerful technique, but as with all things LLM, there are gotchas:

  • Function definitions add to your prompt token count, increasing cost.
  • There may be times where the LLM doesn't decide to return the function call, even when you thought it should have.
  • The more functions you add, the more likely the LLM will get confused about which one to pick, especially if functions are similar to each other. You can try to make it more clear to the LLM by prompt engineering the function name and description, or even providing few shots.

Here are additional approaches you can try:

  • Add filters to the UI of your RAG chat application, as part of the chat box or a sidebar of settings.
  • Add metadata as fields in the search index, and append those to the content sent to the LLM. For example, you could put "Last modified: 2 weeks ago" in each chunk sent to the LLM, if you were trying to help it's ability to answer questions about recency.
  • Use fine-tuning on a model to help it realize when it should call particular functions or respond a certain way. You could even teach it to use a structured query syntax, and remove the functions entirely from your call. This is a last resort, however, since fine-tuning is costly and time-consuming.

Friday, February 16, 2024

RAG techniques: Cleaning user questions with an LLM

📺 You can also watch the video version of this blog post.

When I introduce app developers to the concept of RAG (Retrieval Augmented Generation), I often present a diagram like this:

Diagram of RAG flow, user question to data source to LLM

The app receives a user question, uses the user question to search a knowledge base, then sends the question and matching bits of information to the LLM, instructing the LLM to adhere to the sources.

That's the most straightforward RAG approach, but as it turns out, it's not what quite what we do in our most popular open-source RAG solution, azure-search-openai-demo.

The flow instead looks like this:

diagram of extendex RAG flow, user question to LLM to data source to LLM

After the app receives a user question, it makes an initial call to an LLM to turn that user question into a more appropriate search query for Azure AI search. More generally, you can think of this step as turning the user query into a datastore-aware query. This additional step tends to improve the search results, and is a (relatively) quick task for an LLM. It also cheap in terms of output token usage.

I'll break down the particular approach our solution uses for this step, but I encourage you to think more generally about how you might make your user queries more datastore-aware for whatever datastore you may be using in your RAG chat apps.

Converting user questions for Azure AI search

Here is our system prompt:

Below is a history of the conversation so far, and a new question asked by
the user that needs to be answered by searching in a knowledge base.
You have access to Azure AI Search index with 100's of documents.
Generate a search query based on the conversation and the new question.
Do not include cited source filenames and document names e.g info.txt or doc.pdf in the search query terms.
Do not include any text inside [] or <<>> in the search query terms.
Do not include any special characters like '+'.
If the question is not in English, translate the question to English
before generating the search query.
If you cannot generate a search query, return just the number 0.

Notice that it describes the kind of data source, indicates that the conversation history should be considered, and describes a lot of things that the LLM should not do.

We also provide a few examples (also known as "few-shot prompting"):

query_prompt_few_shots = [
    {"role": "user", "content": "How did crypto do last year?"},
    {"role": "assistant", "content": "Summarize Cryptocurrency Market Dynamics from last year"},
    {"role": "user", "content": "What are my health plans?"},
    {"role": "assistant", "content": "Show available health plans"},
]

Developers use our RAG solution for many domains, so we encourage them to customize few-shots like this to improve results for their domain.

We then combine the system prompts, few shots, and user question with as much conversation history as we can fit inside the context window.

messages = self.get_messages_from_history(
   system_prompt=self.query_prompt_template,
   few_shots=self.query_prompt_few_shots,
   history=history,
   user_content="Generate search query for: " + original_user_query,
   model_id=self.chatgpt_model,
   max_tokens=self.chatgpt_token_limit - len(user_query_request), 
)

We send all of that off to GPT-3.5 in a chat completion request, specifying a temperature of 0 to reduce creativity and a max tokens of 100 to avoid overly long queries:

chat_completion = await self.openai_client.chat.completions.create(
    messages=messages,
    model=self.chatgpt_model,
    temperature=0.0,
    max_tokens=100,
    n=1
)

Once the search query comes back, we use that to search Azure AI search, doing a hybrid search using both the text version of the query and the embedding of the query, in order to optimize the relevance of the results.

Using chat completion tools to request the query conversion

What I just described is actually the approach we used months ago. Once the OpenAI chat completion API added support for tools (also known as "function calling"), we decided to use that feature in order to further increase the reliability of the query conversion result.

We define our tool, a single function search_sources that takes a search_query parameter:

tools = [
  {
    "type": "function",
    "function": {
      "name": "search_sources",
      "description": "Retrieve sources from the Azure AI Search index",
      "parameters": {
        "type": "object",
        "properties": {
          "search_query": {
            "type": "string",
            "description": "Query string to retrieve documents from
                            Azure search eg: 'Health care plan'",
          }
        },
        "required": ["search_query"],
      },
    },
  }
]

Then, when we make the call (using the same messages as described earlier), we also tell the OpenAI model that it can use that tool:

chat_completion = await self.openai_client.chat.completions.create(
    messages=messages,
    model=self.chatgpt_model,
    temperature=0.0,
    max_tokens=100,
    n=1,
    tools=tools,
    tool_choice="auto",
)

Now the response that comes back may contain a function_call with a name of search_sources and an argument called search_query. We parse back the response to look for that call, and extract the value of the query parameter if so. If not provided, then we fallback to assuming the converted query is in the usual content field. That extraction looks like:

def get_search_query(self, chat_completion: ChatCompletion, user_query: str):
    response_message = chat_completion.choices[0].message

    if response_message.tool_calls:
        for tool in response_message.tool_calls:
            if tool.type != "function":
                continue
            function = tool.function
            if function.name == "search_sources":
                arg = json.loads(function.arguments)
                search_query = arg.get("search_query", self.NO_RESPONSE)
                if search_query != self.NO_RESPONSE:
                    return search_query
    elif query_text := response_message.content:
        if query_text.strip() != self.NO_RESPONSE:
            return query_text
    return user_query

This is admittedly a lot of work, but we have seen much improved results in result relevance since making the change. It's also very helpful to have an initial step that uses tools, since that's a place where we could also bring in other tools, such as escalating the conversation to a human operator or retrieving data from other data sources.

To see the full code, check out chatreadretrieveread.py.

When to use query cleaning

We currently only use this technique for the multi-turn "Chat" tab, where it can be particularly helpful if the user is referencing terms from earlier in the chat. For example, consider the conversation below where the user's first question specified the full name of the plan, and the follow-up question used a nickname - the cleanup process brings back the full term.

Screenshot of a multi-turn conversation with final question 'what else is in plus?'

We do not use this for our single-turn "Ask" tab. It could still be useful, particularly for other datastores that benefit from additional formatting, but we opted to use the simpler RAG flow for that approach.

Depending on your app and datastore, your answer quality may benefit from this approach. Try it out, do some evaluations, and discover for yourself!

Sunday, January 28, 2024

Converting HTML pages to PDFs with Playwright

In this post, I'll share a fairly easy way to convert HTML pages to PDF files using the Playwright E2E testing library.

Background: I am working on a RAG chat app solution that has a PDF ingestion pipeline. For a conference demo, I needed it to ingest HTML webpages instead. I could have written my own HTML parser or tried to integrate the LlamaIndex reader, but since I was pressed for time, I decided to just convert the webpages to PDF.

My first idea was to use dedicated PDF export libraries like pdfkit and wkhtml2pdf but kept running into issues trying to get them working. But then I discovered that my new favorite package for E2E testing, Playwright, has a PDF saving function. 🎉 Here’s my setup for conversion.

Step 1: Prepare a list of URLs

For this script, I use the requests package to fetch the HTML for the main page of the website. Then I use the BeautifulSoup scraping library to grab all the links from the table of contents. I process each URL, turning it back into an absolute URL, and add it to the list.

urls = set()
response = requests.get(url, timeout=10)
soup = BeautifulSoup(response.text, "html.parser")
links = soup.find("section", {"id": "flask-sqlalchemy"}).find_all("a")
for link in links:
    if "href" not in link.attrs:
        continue
    # strip off the hash and add back the domain
    link_url = link["href"].split("#")[0]
    if not link_url.startswith("https://"):
        link_url = url + link_url
    if link_url not in urls:
        urls.add(link_url)

See the full code here

Save each URL as PDF

For this script, I import the asynchronous version of the Playwright library. That allows my script to support concurrency when processing the list of URLs, which can speed up the conversion.

from playwright.async_api import BrowserContext, async_playwright

Then I define a function to save a single URL as a PDF. It uses Playwright to goto() the URL, decides on an appropriate filename for that URL, and saves the file with a call to pdf().

async def convert_to_pdf(context: BrowserContext, url: str):
    try:
        page = await context.new_page()
        await page.goto(url)
        filename = url.split("https://flask-sqlalchemy.palletsprojects.com/en/3.1.x/")[1].replace("/", "_") + ".pdf"
        filepath = "pdfs/" / Path(filename)
        await page.pdf(path=filepath)
    except Exception as e:
        logging.error(f"An error occurred while converting {url} to PDF: {e}")

Next I define a function to process the whole list. It starts up a new Playwright browser process, creates an asyncio.TaskGroup() (new in 3.11), and adds a task to convert each URL using the first function.

async def convert_many_to_pdf():
    async with async_playwright() as playwright:
        chromium = playwright.chromium
        browser = await chromium.launch()
        context = await browser.new_context()

        urls = []
        with open("urls.txt") as file:
            urls = [line.strip() for line in file]

        async with asyncio.TaskGroup() as task_group:
            for url in urls:
                task_group.create_task(convert_to_pdf(context, url))
        await browser.close()

Finally, I call that convert-many-to-pdf function using asyncio.run():

asyncio.run(convert_many_to_pdf())

See the full code here

Considerations

Here are some things to think about when using this approach:

  • How will you get all the URLs for the website, while avoiding external URLs? A sitemap.xml would be an ideal way, but not all websites create those.
  • Whats an appropriate filename for a URL? I wanted filenames that I could convert back to URLs later, so I converted / to _ but that only worked because those URLs had no underscores in them.
  • Do you want to visit the webpage at full screen or mobile sized? Playwright can open at any resolution, and you might want to convert the mobile version of your site for whatever reason.

Tuesday, January 16, 2024

Evaluating a RAG chat app: Approach, SDKs, and Tools

When we’re programming user-facing experiences, we want to feel confident that we’re creating a functional user experience - not a broken one! How do we do that? We write tests, like unit tests, integration tests, smoke tests, accessibility tests, loadtests, property-based tests. We can’t automate all forms of testing, so we test what we can, and hire humans to audit what we can’t.

But when we’re building RAG chat apps built on LLMs, we need to introduce an entirely new form of testing to give us confidence that our LLM responses are coherent, grounded, and well-formed.

We call this form of testing “evaluation”, and we can now automate it with the help of the most powerful LLM in town: GPT-4.

How to evaluate a RAG chat app

The general approach is:

  1. Generate a set of “ground truth” data- at least 200 question-answer pairs. We can use an LLM to generate that data, but it’s best to have humans review it and update continually based on real usage examples.
  2. For each question, pose the question to your chat app and record the answer and context (data chunks used).
  3. Send the ground truth data with the newly recorded data to GPT-4 and prompt it to evaluate its quality, rating answers on 1-5 scales for each metric. This step involves careful prompt engineering and experimentation.
  4. Record the ratings for each question, compute average ratings and overall pass rates, and compare to previous runs.
  5. If your statistics are better or equal to previous runs, then you can feel fairly confident that your chat experience has not regressed.

Evaluate using the Azure AI Generative SDK

A team of ML experts at Azure have put together an SDK to run evaluations on chat apps, in the azure-ai-generative Python package. The key functions are:

Start with this evaluation project template

Since I've been spending a lot of time maintaining our most popular RAG chat app solution, I wanted to make it easy to test changes to that app's base configuration - but also make it easy for any developers to test changes to their own RAG chat apps. So I've put together ai-rag-chat-evaluator, a repository with command-line tools for generating data, evaluating apps (local or deployed), and reviewing the results.

For example, after configuring an OpenAI connection and Azure AI Search connection, generate data with this command:

python3 -m scripts generate --output=example_input/qa.jsonl --numquestions=200

To run an evaluation against ground truth data, run this command:

python3 -m scripts evaluate --config=example_config.json

You'll then be able to view a summary of results with the summary tool:

Screenshot of summary tool which shows GPT metrics for each run

You'll also be able to easily compare answers across runs with the compare tool:

Screenshot of compare tool showing answers side by side with GPT metrics below

For more details on using the project, check the README and please file an issue with any questions, concerns, or bug reports.

When to run evaluation tests

This evaluation process isn’t like other automated testing that a CI would runs on every commit, as it is too time-intensive and costly.

Instead, RAG development teams should run an evaluation flow when something has changed about the RAG flow itself, like the system message, LLM parameters, or search parameters.

Here is one possible workflow:

  • A developer tests a modification of the RAG prompt and runs the evaluation on their local machine, against a locally running app, and compares to an evaluation for the previous state ("baseline").
  • That developer makes a PR to the app repository with their prompt change.
  • A CI action notices that the prompt has changed, and adds a comment requiring the developer to point to their evaluation results, or possibly copy them into the repo into a specified folder.
  • The CI action could confirm the evaluation results exceed or are equal to the current statistics, and mark the PR as mergeable. (It could also run the evaluation itself at this point, but I'm wary of recommending running expensive evaluations twice).
  • After any changes are merged, the development team could use an A/B or canary test alongside feedback buttons (thumbs up/down) to make sure that the chat app is working as well as expected.

I'd love to hear how RAG chat app development teams are running their evaluation flows, to see how we can help in providing reusable tools for all of you. Please let us know!

Wednesday, January 10, 2024

Developer relations & motherhood: Will they blend?

My very first job out of college was in developer relations at Google, and it was absolutely perfect for me; a way to combine my love for programming with my interest in teaching. I got to code, write blog posts, organize events, work tightly with eng teams, and do so much traveling, giving talks all over the world. I only left when Google started killing products left and right, including the one I was working on (Wave), and well, my heart was a little broken. (I'm now jaded enough to not loan my whole heart out to corporations)

12 years pass...

I'm back in developer relations, this time for Microsoft/Azure on the Python Advocacy team, and I once again am loving it. It's similar to my old role at Google, but involves more open source work (yay!) and more forms of virtual advocacy (both due to Pandemic and increasingly global audience).

There's a big difference for me this time though: I'm a mom of two kids, a 4 year and a 1 year old (born the week after I started the job). My littlest one is still very attached to me, both emotionally and physically, as she's still nursing and co-sleeping at night, so I essentially have no free time outside of 9-5. (For example, I am writing this a few inches away from here in our floor bed, and have already had to stop/start a few times).

Generally, developer advocacy has been fairly compatible with motherhood, and I'm hugely thankful to Microsoft for their parental leave program (5 months) and support for remote work, and to my manager for understanding my needs as a mother.

However, I've found it stressful to participate fully in all the kinds of events that used to fill my days in DevRel. I'll break down difficulties I've had in fitting events in with my new mom-of-two life, from least to most friction:

  • Live streams: Many advocates (and content creators, generally) will easily hop on a stream to show what they're working on, and it can be a really fun, casual way to connect with the community. I avoided casual streams for the first year of my baby's life, while I was still pumping, as I had to pump too often for it to be practical to be on camera. Now that I'm done pumping, I've had a great time jumping on streams on my colleague's channel. Thanks for the invites, Jay!
  • Virtual events: I'm the one that gets really excited when I hear a conference will be online, since then I can participate from the comfort of my own home. But after speaking at a number of virtual events, I've learnt to ask for more information about the exact timing before getting too excited. Specifically:
    • Is the event in my timezone? I'm in PT, and lots of events cater to audiences in Europe/Asia (rightly so), and their timing may not overlap my workday.
    • Is the event during the week? Lots of conferences are on the weekend, which means paying for childcare and potentially missing out on events with my kids.
    • Is the speaker rehearsal check-in at a convenient time? This is what keeps burning me: I'll happily get a slot speaking at 10AM PT, and then realize there's a speaker mic check at 7AM. I am usually awake at that time, but with a child draped over me who will wake up screaming if I jostle her, waking the rest of the house. Now, if I discover early mic checks, I either pay for my nanny to come early or I explain to them that I can connect but can't test my A/V yet.
  • Local events: I've attended a few Microsoft-sponsored events in SF that were pretty fun. I had to leave before the after parties, and even before the final keynote, in order to get home at a reasonable time for evening nursing, but I still got a lot of good interactions in from 10AM-4PM. There are some local meetups as well, but they tend to be on weeknights/weekends, so I generally avoid them due to the need for childcare. The hassle and added stress on the household often doesn't seem worth it.
  • Non-local events: I've managed to attend zero such events in my 1.5 years at Microsoft! My colleagues have attended events like PyCon and PyCascades, but I haven't felt like I could take an airplane ride with a nursing baby at home. Now that she's nearing two years old, I'm hoping to wean her soon, and a non-local event might become the forcing function for that. I'll be running a session in March at SIGCSE 2024 in Portland, Oregon, which is just a 2-hour plane ride from here, but I'd love to attend for a few days. I'll need to pay our nanny for the night, since she and I are the only two people who can get my little one to sleep, but hey, at least Microsoft pays me fairly well.

You may very well read through all my difficulties and think, "well, why doesn't she just wean the baby? or at least sleep train her?" Reader, I've tried. I'm trying. We're trying. It'll happen eventually.

Once both our kids are preschool aged, it should be much easier for me to participate more fully in events. I never see myself doing anywhere as much travel as I did back in my 20-something Google days, however. It wouldn't be fair to my never-traveling partner to constantly leave him with full parenting duties, and as the child of an always-traveling parent, it's not something I want to do to my kids either. Fortunately, the developer relations field is already much more focused on virtual forms of advocacy, so that is where I hope to hone my skills.

I hope this posts helps anyone else considering the combination of developer relations and motherhood (or more generally, parenting).

Wednesday, January 3, 2024

Using FastAPI for an OpenAI chat backend

When building web APIs that make calls to OpenAI servers, we really want a backend that supports concurrency, so that it can handle a new user request while waiting for the OpenAI server response. Since my apps have Python backends, I typically use either Quart, the asynchronous version of Flask, or FastAPI, the most popular asynchronous Python web framework.

In this post, I'm going to walk through a FastAPI backend that makes chat completion calls to OpenAI. Full code is available on GitHub: github.com/pamelafox/chatgpt-backend-fastapi/


Initializing the OpenAI client

In the new (>= 1.0) version of the openai Python package, the first step is to construct a client, using either OpenAI(), AsyncOpenAI(), AsyncOpenAI, or AzureAsyncOpenAI(). Since we're using FastAPI, we should use the Async* variants, either AsyncOpenAI() for openai.com accounts or AzureAsyncOpenAI() for Azure OpenAI accounts.

But when do we actually initialize that client? We could do it in every single request, but that would be doing unnecessary work. Ideally, we would do it once, when the app started up on a particular machine, and keep the client in memory for future requests. The way to do that in FastAPI is with lifespan events.

When constructing the FastAPI object, we must point the lifespan parameter at a function.

app = fastapi.FastAPI(docs_url="/", lifespan=lifespan)

That lifespan function must be wrapped with the @contextlib.asynccontextmanager decorator. The body of the function setups the OpenAI client, stores it as a global, issues a yield to signal it's done setting up, and then closes the client as part of shutdown.

from .globals import clients


@contextlib.asynccontextmanager
async def lifespan(app: fastapi.FastAPI):
    
    if os.getenv("OPENAI_KEY"):
         # openai.com OpenAI
        clients["openai"] = openai.AsyncOpenAI(
            api_key=os.getenv("OPENAI_KEY")
        )
    else:
        # Azure OpenAI: auth is more involved, see full code.
        clients["openai"] = openai.AsyncAzureOpenAI(
            api_version="2023-07-01-preview",
            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
            **client_args,
        )

    yield

    await clients["openai"].close()

See full __init__.py.

Unfortunately, FastAPI doesn't have a standard way of defining globals (like Flask/Quart with the g object), so I am storing the client in a dictionary from a shared module. There are some more sophisticated approaches to shared globals in this discussion.


Making chat completion API calls

Now that the client is setup, the next step is to create a route that processes a message from the user, sends it to OpenAI, and returns the OpenAI response as an HTTP response.

We start off by defining pydantic models that describe what a request looks like for our chat app. In our case, each HTTP request will contain JSON with two keys, a list of "messages" and a "stream" boolean:

class ChatRequest(pydantic.BaseModel):
    messages: list[Message]
    stream: bool = True

Each message contains a "role" and "content" key, where role defaults to "user". I chose to be consistent with the OpenAI API here, but you could of course define your own input format and do pre-processing as needed.

class Message(pydantic.BaseModel):
    content: str
    role: str = "user"

Then we can define a route that handles chat requests over POST and send backs a non-streaming response:

@router.post("/chat")
async def chat_handler(chat_request: ChatRequest):
	messages = [{"role": "system", "content": system_prompt}] + chat_request.messages
    response = await clients["openai"].chat.completions.create(
    	messages=messages,
    	stream=False,
    )
    return response.model_dump()

The auto-generated documentation shows the JSON response as expected:

Screenshot of FastAPI documentation with JSON response from OpenAI chat completion call


Sending back streamed responses

It gets more interesting when we add support for streamed responses, as we need to return a StreamingResponse object pointing at an asynchronous generator function.

We'll add this code inside the "/chat" route:

if chat_request.stream:
    async def response_stream():
        chat_coroutine = clients["openai"].chat.completions.create(
            messages=messages,
            stream=True,
        )
        async for event in await chat_coroutine:
            yield json.dumps(event.model_dump(), ensure_ascii=False) + "\n"

    return fastapi.responses.StreamingResponse(response_stream())

The response_stream() function is an asynchronous generator, since it is defined with async and has a yield inside it. It uses async for to loop through the asynchronous iterable results of the Chat Completion call. For each event it receives, it yields a JSON string with a newline after it. This sort of response is known as "json lines" or "ndjson" and is my preferred approach for streaming JSON over HTTP versus other protocols like server-sent events.

The auto-generated documentation doesn't natively understand streamed JSON lines, but it happily displays it anyways:

Screenshot of FastAPI server response with streamed JSON lines


All together now

You can see the full router code in chat.py. You may also be interested in the tests folder to see how I fully tested the app using pytest, extensive mocks of the OpenAI API, including Azure OpenAI variations, and snapshot testing.

Using llamafile for local dev for an OpenAI Python web app

We're seeing more and more LLMs that can be run locally on a laptop, especially those with GPUs and multiple cores. Open source projects are also making it easier to run LLMs locally, so that you don't have to be an ML engineer or C/C++ programmer to get started (Phew!).

One of those projects is llamafile, which provides a single executable that serves up an API and frontend to interact with a local LLM (defaulting to LLaVa). With just a few steps, I was able to get the llamafile server running. I then discovered that llamafile includes an OpenAI-compatible endpoint, so I can point my Azure OpenAI apps at the llamafile server for local development. That means I can save costs and also evaluate the quality difference between deployed models and local models. Amazing!

I'll step through the process and share my sample app, a FastAPI chat app backend.

Running the llamafile server

Follow the instructions in the quickstart to get the server running.

Test out the server by chatting with the LLM:

Screenshot of llama.cpp conversation about haikus

Using the OpenAI-compatible endpoint

The llamafile server includes an endpoint at "/v1" that behaves just like the OpenAI servers. Note that it mimics the OpenAI servers, not the *Azure* OpenAI servers, so it does not include additional properties like the content safety filters.

Test out that endpoint by running the curl command in the JSON API quickstart.

You can also test the Python code in that JSON quickstart to confirm it works as well.

Using llamafile with an existing OpenAI app

As the llamafile documentation shows, you can point an OpenAI Python client at a local server by overriding base_url and providing a bogus api_key.

client = AsyncOpenAI(
    base_url="http://localhost:8080/v1",
    api_key = "sk-no-key-required"
)

I tried that out with one of my Azure OpenAI samples, a FastAPI chat backend, and it worked for both streaming and non-streaming responses! 🎉

Screenshot of response from FastAPI generated documentation for a request to make a Haiku

Switching with environment variables

I wanted it to be easy to switch between Azure OpenAI and a local LLM without changing any code, so I made an environment variable for the local LLM endpoint. Now my client initialization code looks like this:

if os.getenv("LOCAL_OPENAI_ENDPOINT"):
    client = openai.AsyncOpenAI(
    	api_key="no-key-required",
        base_url=os.getenv("LOCAL_OPENAI_ENDPOINT")
    )
else:
    # Lots of Azure initialization code here...
    # See link below for full code.
    client = openai.AsyncAzureOpenAI(
      api_version="2023-07-01-preview",
      azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
      # plus additional args
    )

See full code in __init__.py.

Notice that I am using the Async version of the OpenAI clients in both cases, since this backend uses FastAPI and 100% asynchronous calls. For llamafile, I don't bother using the Azure version of the client, since llamafile is only trying to mimic the openai.com servers. That should be fine, as I typically code assuming the openai.com servers as a baseline, and just taking advantage of Azure extras (like content safety filters) when available.

I will likely try out llamafile for my other Azure OpenAI samples soon, and run some evaluations to see how llamafile compares in terms of quality. I don't have any plans to use non-OpenAI models in production, but I want to keep monitoring how well the local LLMs can perform and what use cases there may be for them.