guardrails icon indicating copy to clipboard operation
guardrails copied to clipboard

feat: add GuardrailsEngine for llama index

Open kaushikb11 opened this issue 1 year ago • 4 comments

Updated.

Implement GuardrailsEngine for LlamaIndex integration

This PR introduces the GuardrailsEngine class, which integrates Guardrails validation with LlamaIndex's query and chat engines. The GuardrailsEngine provides a unified interface for applying guardrails to both query and chat functionalities while maintaining compatibility with LlamaIndex's expected interfaces.

  1. Inherits from BaseQueryEngine for compatibility with LlamaIndex components.
  2. Supports both query and chat functionalities by wrapping either a BaseQueryEngine or a BaseChatEngine.

QueryEngine and ChatEngine Integration:

The GuardrailsEngine is designed to work seamlessly with both LlamaIndex's QueryEngine and ChatEngine.

This dual functionality allows users to apply guardrails to both one-off queries and multi-turn conversations, enhancing the safety and reliability of LLM responses in various use cases.

Example

from guardrails import Guard
from guardrails.hub import ToxicLanguage, CompetitorCheck
from guardrails.integrations.llama_index import GuardrailsEngine

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core import Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model

documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

guard = Guard().use(
    ToxicLanguage(on_fail="exception"),
    CompetitorCheck(competitors=["Apple", "Google", "Microsoft"], on_fail="exception")
)

# For query engine
query_engine = index.as_query_engine()
guardrails_engine = GuardrailsEngine(engine=query_engine, guard=guard)

try:
    response = query_engine.query("What are the main products of our company?")
    print(response)
except ValueError as e:
    print(f"Validation error: {str(e)}")
    
# For chat engine
chat_engine = index.as_chat_engine()
guardrails_engine = GuardrailsEngine(engine=chat_engine, guard=guard)

try:
	response = guardrails_engine.chat("Tell me about our product compared to Apple's.")
	print(response)
except ValueError as e:
	print(f"Validation error: {str(e)}")

kaushikb11 avatar Aug 11 '24 23:08 kaushikb11

Interesting design! A few followups

  1. Why implement this as a query engine instead of something else?
  2. Guardrails already natively has support for input and output validation at the guard level. You can do Guard().use(Validator, on='messages') to run validation on input
  3. The flow here is interesting, and I wonder if you've thought of inverting it. i.e. in teh current code, I assume something happens like input validation, call the LLM through llamaindex's call (not guardrails), and then output parse. What if we passed in a llamaindex module as the LLM call itself to an existing guard. What do we lose with that design?

zsimjee avatar Aug 13 '24 14:08 zsimjee

Why implement this as a query engine instead of something else?

  • Query engines in LlamaIndex are end-to-end pipelines that allow asking questions over data. They handle the retrieval of relevant context and passing it to the LLM along with the query.
  • By implementing GuardrailsEngine as a query engine, it integrates seamlessly with LlamaIndex's existing infrastructure, allowing users to easily wrap it on top of other query engines.
  • Query engines in LlamaIndex are designed to work with both one-off queries and multi-turn conversations (via chat engines), making the GuardrailsEngine versatile for various use cases.

Guardrails already natively has support for input and output validation at the guard level. You can do Guard().use(Validator, on='messages') to run validation on input.

You're correct. Removed the need of defining the input and output validators separately. Current design should account for it. It will be highlighted in the documentation.

The flow here is interesting, and I wonder if you've thought of inverting it. i.e. in the current code, ...... What if we passed in a llamaindex module as the LLM call itself to an existing guard. What do we lose with that design?

Let's consider the implications of inverting the flow: Current flow: Input validation -> LlamaIndex LLM call -> Output parsing Proposed flow: Pass LlamaIndex module as LLM to Guardrails Guard

Potential benefits of the inverted design:

  • It would leverage more of Guardrails' native functionality, potentially simplifying the implementation.
  • It might offer more flexibility in applying different guardrails to different LlamaIndex components.

Potential drawbacks or limitations:

  • We might lose some fine-grained control over the LlamaIndex query process, as it would be abstracted within the Guard's LLM call.
  • It could be more challenging to access and utilize LlamaIndex-specific features and metadata.
  • Gets messier to deal with the LLM API responses, could pass a wrapper but doesn't seem clean from developer's perspective.
  • The integration might be less intuitive for users already familiar with LlamaIndex's query and chat engine paradigm. (v important for developer experience).

kaushikb11 avatar Aug 23 '24 15:08 kaushikb11

Could you add some unit and integration tests to assert the code is acting as expected?

yup, of course. it would be better for us to align on the RFC first. Will be fixing the typing issues too.

kaushikb11 avatar Aug 23 '24 15:08 kaushikb11

This PR is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 14 days.

github-actions[bot] avatar Sep 29 '24 03:09 github-actions[bot]

This PR was closed because it has been stalled for 14 days with no activity.

github-actions[bot] avatar Oct 14 '24 03:10 github-actions[bot]

Closing in favor of branch on guardrails repo: https://github.com/guardrails-ai/guardrails/tree/feat/llama-index

CalebCourier avatar Nov 12 '24 15:11 CalebCourier

See: https://github.com/guardrails-ai/guardrails/pull/1160

CalebCourier avatar Nov 12 '24 15:11 CalebCourier