LlamaIndex integration

In this notebook, we show how it’s easy to build ChainML agents that leverage the power of LlamaIndex to integrate data for your agents.

Setup

Integration with LlamaIndex is easy and straightforward. To use LlamaIndex with the ChainML framework, you will need to install “Llama-Index” via pip.

Example

$ pip install Llama-Index

# Load environment variables
import dotenv
import os

dotenv.load_dotenv()
print(os.getenv("OPENAI_API_KEY", None) is not None )

True

Example of using Llama Index to retrieve relevant information from documents - here we use the book, the Great Gatbsy

# download Great Gatsby example from Llama Index

import os
import requests

url = "https://github.com/jerryjliu/llama_index/blob/main/examples/gatsby/gatsby_full.txt"
filename = url.split("/")[-1]

os.makedirs("gatsby_download", exist_ok=True)

response = requests.get(url)
with open(os.path.join("gatsby_download", filename), "wb") as f:
    f.write(response.content)

from llama_index import VectorStoreIndex, SimpleDirectoryReader

# build index of book / This step could be slow
documents = SimpleDirectoryReader("gatsby_download").load_data()
index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine()

# check the index is working
response = query_engine.query("Where do Gatsby and Daisy meet?")
print(response)

Gatsby and Daisy meet at Gatsby's house on Long Island.

Examples

Chain Using LlamaIndex

Let’s create a chain that uses LlamaIndex to retrieve relevant context for a user’s query for a simple chatbot that can query the data. We will look up context then feed it into a prompt template for an LLMSkill to respond.

We define a LlamaIndexSkill that uses a query engine to look up indexed data.

from council.runners import Budget
from council.contexts import SkillContext, ChatMessage
from council.skills import SkillBase
from llama_index.indices.query.base import BaseQueryEngine


class LlamaIndexSkill(SkillBase):
    queryEngine: BaseQueryEngine

    def __init__(self, query_engine: BaseQueryEngine):
        SkillBase.__init__(self, "llama index skill")
        self.queryEngine = query_engine

    def execute(self, context: SkillContext, budget: Budget) -> ChatMessage:
        prompt = context.chat_history.try_last_user_message.unwrap("no user message").message
        print(prompt)
        response = self.queryEngine.query(prompt)
        return self.build_success_message(response)

from council.agents import Agent
from council.chains import Chain

# wrap into a trivial agent that just answers document queries
index_skill = LlamaIndexSkill(query_engine)
agent = Agent.from_skill(index_skill, "document index")

# message="Whose eyes are on the billboard?"
# message="What are the personalities of Tom and Daisy?"
# message="What era does the book take place in?"
message="Who falls in love with Daisy?"

result = agent.execute_from_user_message(message)
print(result.best_message.message)

Who falls in love with Daisy?

Gatsby falls in love with Daisy.

Here’s an example of integrating LlamaIndex into a more capable chain - this chain answers a user request for information about the book by: 1. formulating a query for information to retrieve from the book using an LLM (here GPT-3.5-Turbo) 2. passing that query to LlamaIndex to retrieve relevant passages from the book

3. passing those passages as context to the LLM along with the original query to generate a summary response

from council.utils import Option
from council.evaluators import BasicEvaluator
from council.controllers import BasicController
from council.prompt import PromptBuilder
from council.skills import PromptToMessages, LLMSkill
from council.llm import OpenAILLMConfiguration, OpenAILLM

# agent to use index to provide context for more complex answers
import dotenv


dotenv.load_dotenv()
config = OpenAILLMConfiguration.from_env()
config.model = Option("gpt-3.5-turbo")
llm = OpenAILLM(config)

context_prompt = PromptToMessages(
    PromptBuilder("Please identify query terms to respond to the following user request {{chat_history.last_message}}")
)
context_query_skill = LLMSkill(
    llm,
    "You are an expert in the Great Gatbsy. Identify relevant query terms to search for context in the book.",
    context_messages=context_prompt.to_user_message,
)

index_skill = LlamaIndexSkill(query_engine)
index_prompt = PromptToMessages(
    PromptBuilder(
        "Here are relevant quotes from the book: {{chain_history.last_message}} \nUse this to respond to the following user request {{chat_history.last_message}}"
    )
)
response_skill = LLMSkill(
    llm,
    "You are an expert in the Great Gatbsy. Provide a helpful response to the user's question",
    context_messages=index_prompt.to_user_message,
)

chain = Chain("docindex", "document index", [context_query_skill, index_skill, response_skill])
agent = Agent(BasicController(), [chain], BasicEvaluator())

import logging

logging.basicConfig(
    format="[%(asctime)s %(levelname)s %(threadName)s %(name)s:%(funcName)s:%(lineno)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S%z",
)
## uncomment me to see the engine logs
logging.getLogger("council").setLevel(logging.WARN)

# message="Whose eyes are on the billboard?")
# message="What are the personalities of Tom and Daisy?"
# message="What era does the book take place in?"
# message="What is the significance of the eyes on the billboard?"
message="What are the key plot events in the book?"

result = agent.execute_from_user_message(message=message)
print(result.best_message.message)

What are the key plot events in the book?
The key plot events in the book "The Great Gatsby" include Gatsby's funeral, the man with owl-eyed glasses attending the funeral, Tom and Jordan taking Nick away from the funeral, Nick's realization that it is his 30th birthday, the discovery of George Wilson's wife locked in the garage, and the death car accident.