Haystack Example

Multi-call chains using Haystack

This example is an extension of Haystack's "Building a Chat Application with Function Calling" tutorial. To learn more see the original here. For the sake of brevity we will not be explaining Haystack components and concepts here, please see the original tutorial for more information regarding these.

Development Environment Setup

First we will setup the development environment and install haystack, context and sentence-transformers.

pip install context-py
pip install haystack-ai
pip install "sentence-transformers>=2.2.0"
Imports

We will import all the Haystack components required for this tutorial, setup our OpenAI key and import traceable and dynamic_traceable from context.

from getcontext.tracing import dynamic_traceable, traceable
import os

from haystack import Pipeline, Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.writers import DocumentWriter
from haystack.components.embedders import (
    SentenceTransformersTextEmbedder,
    SentenceTransformersDocumentEmbedder,
)
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator

os.environ["OPENAI_API_KEY"] = "OPENAI_KEY"

To begin we will create a function to setup a document store. create_document_store defines a list of Documents and a Haystack Pipeline to process these.

Notice the @traceable decorator, this allows context to trace this entire function from a unit test. We also override the indexing_pipeline.run method using dynamic_traceable and provide a custom name "indexing_pipeline_run". This provides the same tracing functionality on a function we do not have access to (defined in Haystack's library).

@traceable
def create_document_store():
    documents = [
        Document(content="My name is Jean and I live in Paris."),
        Document(content="My name is Mark and I live in Berlin."),
    ]
    document_store, indexing_pipeline = InMemoryDocumentStore(), Pipeline()
    
    # populate indexing_pipeline with embedder and writer...

    # override run method with dynamic_traceable wrapped instance
    indexing_pipeline.run = dynamic_traceable(
        indexing_pipeline.run,
        run_type="chain",
        name="indexing_pipeline_run")
    indexing_pipeline.run({"doc_embedder": {"documents": documents}})
    
    return document_store

dynamic_traceable is useful when you do not have access to the underlying function definition (e.g. when you want to trace a function from an imported library). You can also use dynamic_traceable to assign custom names individual spans at runtime.

Similar to create_document_store we define another function build_rag_pipeline which builds our document store and builds another RAG pipeline.

Once again we use a @traceable decorator around the function and use dynamic_traceable to add traces to the OpenAI model and the pipeline run methods. For the generator we override the create function and set the run_type to "llm". We can attach evaluators to LLM spans in unit tests.

@traceable
def build_rag_pipeline():
    document_store, rag_pipe = create_document_store(), Pipeline()

    # add an embedder, retriever and prompt builder to rag_pipe...

    # override OpenAI's chat.completions.create with an LLM dynamic_traceable
    generator = OpenAIGenerator(model="gpt-3.5-turbo")
    generator.client.chat.completions.create = dynamic_traceable(
        generator.client.chat.completions.create,
        run_type="llm",
        name="openai_completion_create")
    rag_pipe.add_component("llm", generator)

    # again, override pipeline run method with
    rag_pipe.run = dynamic_traceable(rag_pipe.run, run_type="chain", name="rag_pipe_run")

    return rag_pipe

Finally, we define an ask_question function which accepts a query which it passes to the RAG pipeline.

@traceable
def ask_question(query: str):
    rag_pipe = build_rag_pipeline()
    return rag_pipe.run(
        {"embedder": {"text": query}, "prompt_builder": {"question": query}}
    )

It is now time to test our new pipeline to make sure it works as expected! We will be using Python's unittest library for this.

We start off with a simple test which simply captures a trace and send it to context ai. We have optionally set the logger to logging.INFO to see what is happening under the hood.

import unittest
import logging
from question_answer import ask_question
from getcontext.tracing import capture_trace, Evaluator

class TestQuestionAnswer(unittest.TestCase):
    def setUp(self):
        logging.basicConfig()
        logging.getLogger('context-ai').setLevel(logging.INFO)

    def test_ask_question_mark(self):
        trace = capture_trace(ask_question, "Where does Mark live?", trace_name='Mark')

To make things a bit more interesting, let's add some evaluators and a test for a person not included in the document list (and which we would expect the LLM to not respond).

def test_ask_question_mark(self):
    trace = capture_trace(ask_question, "Where does Mark live?", trace_name='Mark')

    evaluator = Evaluator(
        evaluator='semantic_match',
        options={'semantic_match_string': 'Mark lives in Berlin'}
    )
    trace.add_evaluator('openai_completion_create', evaluator)
    trace.evaluate()

def test_ask_question_diana(self):
    trace = capture_trace(ask_question, "Where does Diana live?", trace_name='Diana')

    evaluator = Evaluator(
        evaluator='semantic_match',
        options={'semantic_match_string': "I don't know where Diana lives"}
    )
    trace.add_evaluator('openai_completion_create', evaluator)
    trace.evaluate()

We can see below that both test cases passed 🎉. The LLM was able to give information about Mark's location but replied saying that it did not know where Diana lived.

If we login in to the Context.ai platform we can then inspect the trace.

If you have any feedback, please let us know by emailing henry@context.ai

Last updated