Integrations

LlamaIndex

llama-index-wity lets you use Wity inside LlamaIndex. Ask typed questions in a workflow, let Wity pick the query engine for each query, and stop risky tool calls before they run.

It has three parts:

WityClassifier

classifier

Asks Wity typed questions from any LlamaIndex code. See In a workflow.

WitySelector

router selector

Wity picks which query engine or retriever answers each query. A query engine is LlamaIndex's question-answering object over your data. See Router.

WityToolGate

tool wrapper

Wity checks risky tool calls before they run. See Tool-risk gate.

Install#

pip install llama-index-wity

Needs Python 3.10 or newer.
It's built on the Python SDK, so answers have the same fields.
The package is called llama-index-wity, but in code you import it as llama_index.wity.

API key#

Get an API key from the keys page. Save it in your environment variables as WITY_API_KEY, and the classifier reads it by itself:

export WITY_API_KEY="wity_..."

To pass a key from somewhere else, like a secret manager, use WityClassifier(api_key=...).

This example reads a customer's message and asks Wity three questions about it: which team should handle it, whether it's urgent, and how much it hurts the customer. All three are answered in one request.

from llama_index.wity import WityClassifier, choice, noul, score
 
classifier = WityClassifier()   # reads your key from WITY_API_KEY
 
res = classifier.classify(
    # The text you want Wity to judge.
    "I was charged twice for my subscription this month. Please refund one of them today.",
    {
        # Pick one of these teams.
        "team": choice("Which team should handle this ticket?", {
            "billing": "Payments, charges and refunds",
            "technical": "Bugs, errors and outages",
            "other": "Anything else",
        }),
        # Yes or no.
        "urgent": noul(
            "Does the customer ask for something to happen soon?",
            yes="They ask for fast action",
            no="No time pressure",
        ),
        # A level on a scale, from lowest to highest.
        "severity": score("How much is this hurting the customer?", ["Not at all", "A little", "A lot"]),
    },
)
 
res.choices["team"].choice         # "billing"
res.choices["team"].probabilities  # {"billing": 0.99993..., "technical": 0.00004..., "other": 0.00001...}
res.nouls["urgent"].noul           # 0.99986...: the probability of yes
res.scores["severity"].score       # 1.97852...: the expected level, from 0 (Not at all) to 2 (A lot)

The first argument is the state: the text Wity judges. The second names each question, and you use that name to find its answer.
Answers are grouped by type in res.choices, res.nouls and res.scores. They're the Python SDK's answer objects, unchanged.
Use await classifier.aclassify(...) in async code.
Keep one classifier and reuse it, so later calls reuse its connection. Call close() when you're done.

Probabilities

Probabilities add up to 1, and a higher number means Wity is more sure. They aren't exact odds. Pick thresholds that suit your data. See Probabilities & confidence.

In a workflow#

Give the classifier its questions when you create it. Then each call takes only the state. State can be text, a dict, a list, or LlamaIndex ChatMessages, so Wity can judge a whole chat:

from llama_index.core.workflow import StartEvent, StopEvent, Workflow, step
from llama_index.wity import WityClassifier, choice
 
# Questions set once, so each call takes only the chat.
team = WityClassifier(
    questions={
        "team": choice("Which team should handle this chat?", {
            "billing": "Payments, charges and refunds",
            "technical": "Bugs, errors and outages",
        })
    },
    reasoning="off",
)
 
class Triage(Workflow):
    @step
    async def route(self, ev: StartEvent) -> StopEvent:
        res = await team.aclassify(ev.chat)
        answer = res.choices["team"]
        if answer.confidence < 0.7:
            return StopEvent(result="Sent to a person: Wity isn't sure.")
        return StopEvent(result=f"Sent to {answer.choice}.")
 
await Triage().run(chat=chat)  # "Sent to technical."

Give questions to the constructor or to each call, not both.
Images, audio and files in messages are sent by type only, like {"type": "image"}.

Router#

WitySelector plugs into RouterQueryEngine and RouterRetriever. Each tool's description is one option, and Wity picks one for each query.

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.tools import QueryEngineTool
from llama_index.wity import WitySelector
 
# billing_engine and api_engine are your own query engines.
router = RouterQueryEngine(
    selector=WitySelector(),
    query_engine_tools=[
        QueryEngineTool.from_defaults(billing_engine, description="Plans, invoices, charges and refunds"),
        QueryEngineTool.from_defaults(api_engine, description="API endpoints, parameters and errors"),
    ],
    llm=llm,
)
 
response = router.query("What does a 429 from the API mean?")  # answered by api_engine
response.metadata["selector_result"].reason
# "Wity picked this choice with probability 0.999333297708751."

To change how Wity picks, pass WitySelector(instructions="Which docs can answer this question?").
It picks exactly one. Merging answers from several engines isn't supported.
With only one tool, it's picked without a call to Wity.
If Wity fails, the error is raised. The selector never guesses.
Pass RouterQueryEngine an llm anyway, even though Wity does the picking. Without one, LlamaIndex falls back to its default, OpenAI.

Tool-risk gate#

WityToolGate checks a tool before it runs. Wrap each risky tool with gate.guard(), and give the wrapped tools to your agent. Wity sees the proposed call and the conversation, and answers one yes-or-no question: should this call be stopped?

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.wity import WityToolGate
 
# read_page, send_email and delete_file are your own tools.
# Only the wrapped tools are checked.
gate = WityToolGate()
agent = FunctionAgent(tools=[read_page, gate.guard(send_email), gate.guard(delete_file)], llm=llm)

It stops calls the user didn't ask for, calls that do more than the user asked, and calls that expose private data.
A call that does exactly what the user asked runs, even one like deleting the file they named.
Only the user's own messages can approve a call. Text in a web page or file can't, even if it pretends to give orders.
A stopped tool doesn't run. The agent gets an error result instead, so it can tell the user. The same happens if Wity fails.
The wrapped tool keeps its name, description and arguments, so the LLM sees no difference.
Call wrapped tools through a LlamaIndex agent, like FunctionAgent or ReActAgent. The gate reads the conversation from the agent.

Tuning the gate

A call is stopped when the probability of risky is at or above threshold, 0.5 by default. Try the gate on some of your own calls, safe and risky. Then set threshold between where they land:

from llama_index.wity import NoulCriteria
 
# Block only when Wity is more sure the call is risky.
gate = WityToolGate(threshold=0.8)
 
# Or say what risky and safe mean for your tools.
gate = WityToolGate(
    criteria=NoulCriteria(
        true="The call writes, deletes, publishes or changes access.",
        false="The call only reads public or user-provided data.",
    ),
)

Reasoning#

Wity answers most questions in one quick pass. On close calls it can think first, which usually takes 1 to 4 seconds. You choose when with reasoning:

"auto"

default

Think only when the question needs it. The classifier, the selector and the gate all use this.

"off"

fastest

Never think.

"always"

slowest

Think before every answer.

classifier = WityClassifier(reasoning="off", max_latency_ms=2000)
 
# The selector and the gate take a classifier of their own, without questions.
selector = WitySelector(classifier=WityClassifier(reasoning="off"))
gate = WityToolGate(classifier=WityClassifier(reasoning="always"))

max_latency_ms caps how long one request may take, from 200 to 120000 ms. More in Reasoning & auto mode.

Options#

Everything WityClassifier takes:

questions

dict

Fixed questions. Then each call takes only the state.

reasoning

off | auto | always

When Wity may think first. Default auto.

max_latency_ms

200 to 120000

Time limit per request.

api_key

str

Your key. Default: WITY_API_KEY.

base_url

str

The API address. Default: WITY_BASE_URL, then Wity's address.

timeout

seconds

Per attempt. Default 60.

max_retries

int

Retries on 408, 429, 5xx and network failures. Default 2.

Errors#

Errors are the Python SDK errors, like RateLimitError. You can import them from llama_index.wity, and all of them are a WityError.
The classifier already retries. Set max_retries=0 before you add retries of your own.
Each call is a span in LlamaIndex's tracing, so it shows up in tools like Arize Phoenix and Langfuse.
Your key never appears in errors, logs or traces.

The package is llama-index-wity on PyPI. Questions and bugs: info@alphanimble.com. It's licensed Apache-2.0.

← LangChain

Vercel AI SDK →