Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add GuardrailsEngine for llama index #1005

Closed
wants to merge 10 commits into from

Conversation

kaushikb11
Copy link
Contributor

@kaushikb11 kaushikb11 commented Aug 11, 2024

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)}")

@zsimjee
Copy link
Collaborator

zsimjee commented Aug 13, 2024

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?

Copy link
Collaborator

@CalebCourier CalebCourier left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have similar questions as @zsimjee. WRT to principal of least surprise I would expect to see something like query_engine.query passed as the llm_api and the existing forms of input and output validation used on a single Guard.

In addition to this I have doubts about the behaviour of the current implementation since the types don't line up. Could you add some unit and integration tests to assert the code is acting as expected?

@kaushikb11 kaushikb11 changed the title feat: add GuardrailsQueryEngine for llama index feat: add GuardrailsEngine for llama index Aug 23, 2024
@kaushikb11
Copy link
Contributor Author

kaushikb11 commented Aug 23, 2024

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
Copy link
Contributor Author

kaushikb11 commented Aug 23, 2024

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 kaushikb11 marked this pull request as ready for review August 27, 2024 12:50
Copy link

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 github-actions bot added the Stale label Sep 29, 2024
Copy link

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

@github-actions github-actions bot closed this Oct 14, 2024
@CalebCourier CalebCourier reopened this Nov 12, 2024
@CalebCourier
Copy link
Collaborator

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

@CalebCourier
Copy link
Collaborator

See: #1160

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants