Creating an assistant for RAG with query rephrasing
In AI Assistant API, you can enable an additional rephrasing model to make search index queries more accurate. Rephrasing will correct errors and typos and help you out if the response requires massive volumes of indexed data and heavily relies on the conversation context.
Getting started
To use the examples:
-
Create a service account and assign the
ai.assistants.editor
andai.languageModels.user
roles to it. -
Get the service account API key and save it.
The following examples use API key authentication. Yandex Cloud ML SDK also supports IAM token and OAuth token authentication. For more information, see Authentication in Yandex Cloud ML SDK.
-
Use the pip
package manager to install the ML SDK library:pip install yandex-cloud-ml-sdk
Create an assistant
In this example, we are going to use the tour information scenario to show how to create an assistant with a search index and an active rephrasing model.
-
Download and unpack the archive
with examples of files that will be used as an additional source of information. The files contain advertising texts for tours to Bali and Kazakhstan generated by YandexGPT Pro. -
Create a file named
rephraser.py
and paste the following code into it providing the folder ID and authentication credentials:#!/usr/bin/env python3 from __future__ import annotations import pathlib from yandex_cloud_ml_sdk import YCloudML LABEL_KEY = 'yc-ml-sdk-example' PATH = pathlib.Path(__file__) NAME = f'example-{PATH.parent.name}-{PATH.name}' def local_path(path: str) -> pathlib.Path: return pathlib.Path(__file__).parent / path def get_search_index(sdk): # Function for retrieving or creating a search_index object for search_index in sdk.search_indexes.list(): if search_index.labels and search_index.labels.get(LABEL_KEY) == NAME: print(f'using {search_index=}') break else: print('no search indexes found, creating new one') file = sdk.files.upload( local_path('turkey_example.txt') ) operation = sdk.search_indexes.create_deferred(file, labels={LABEL_KEY: NAME}) search_index = operation.wait() print(f'new {search_index=}') file.delete() return search_index def delete_labeled_entities(iterator): # Cleanup function for deleting all iterator objects with the .labels attribute for entity in iterator: if entity.labels and entity.labels.get(LABEL_KEY) == NAME: print(f'deleting {entity.__class__.__name__} with id={entity.id!r}') entity.delete() def main() -> None: sdk = YCloudML( folder_id="<folder_ID>", auth="<API_key>", ) sdk.setup_default_logging(log_level='WARNING') search_index = get_search_index(sdk) labels = {LABEL_KEY: NAME} # Search index tool tool = sdk.tools.search_index(search_index) # Rephraser tool rephraser = sdk.tools.rephraser('rephraser', model_version='latest') tool_with_rephraser = sdk.tools.search_index(search_index, rephraser=rephraser) # Creating two assistants: one with rephraser and one without it assistant_wo_rephraser = sdk.assistants.create('yandexgpt', labels=labels, tools=[tool]) assistant_with_rephraser = sdk.assistants.create('yandexgpt', labels=labels, tools=[tool_with_rephraser]) # The following code sends queries with and without rephrasing # to demonstrate what it does. # You can use other queries. thread = sdk.threads.create(labels=labels) def run(query, rephrase: bool) -> None: assistant = assistant_with_rephraser if rephrase else assistant_wo_rephraser thread.write(query) run = assistant.run(thread) result = run.wait() print(f"Question: {query}") preposition = 'with' if rephrase else 'without' print(f"Answer {preposition} rephraser:\n {result.text!r}") print() # Let's take two valid and two invalid queries: run('Where to escape to?', rephrase=False) # 1 run('Wher to cuutt looooose?', rephrase=False) # 2 run('Where to go for a break?', rephrase=False) # 3 run('Where to escape to?', rephrase=False) # 4 # The first two queries require rephrasing: run('Where to escape to?', rephrase=True) run('Wher to cuutt looooose?', rephrase=True) # Let's delete all queries and threads # but keep the search index because it takes a long time to create one. # You can reuse the search index you created. delete_labeled_entities(sdk.assistants.list()) delete_labeled_entities(sdk.threads.list()) if __name__ == '__main__': main()
Where:
<folder_ID>
: ID of the folder the service account was created in.<API_key>
: Service account API key required for authentication in the API.
-
Run the file you created:
python3 rephraser.py
Result
Question: Where to escape to? Answer without rephraser: 'The question of "where to escape to?" may be interpreted differently depending on context. If considering a vacation or relocation, the choice of place may depend on your preferences, interests, and budget. Here are some ideas for you:\n\n1. **Nature spots:** If you are a nature lover, consider visiting national parks, nature reserves, or mountain resorts.\n\n2. **Historical cities:** If you are a history and culture geek, explore rich historical heritage sites like Rome, Paris, Istanbul, etc.\n\n3. **Beach getaways:** If you prefer a seaside vacation, consider beach resorts around the world in countries like Thailand, Spain, Greece, etc.\n\n4. **Creative hubs:** If you are a creative person, you might enjoy cities with a vibrant cultural life like New York, London, Berlin, etc.\n\n5. **Exotic destinations:** If looking for something offbeat and exotic, consider visiting countries with unique culture and nature such as India, Peru, Indonesia and the like.\n\nRemember that your "escape" destination depends on your personal interests and preferences. Before you make any travel or relocation decisions, it is a good idea to research the place, its culture, security, and other aspects.\n\nIf your question has a different context, please clarify it, and I will try to give a more specific answer.' Question: Wher to cuutt looooose? Answer without rephraser: 'It looks like your question contains an error or a typo. Please rephrase it and I will try to help you.' Question: Where to go for a break? Answer without rephraser: 'Consider a break in Bali. This Indonesian island is famous for its beautiful beaches, unique culture, and warm hospitality of its people. Enjoy its natural beauty, try local cuisine, and make new friends.' Question: Where to escape to? Answer without rephraser: 'The question of "where to escape to?" can be answered differently depending on your goals and preferences. Here are some ideas for you:\n\n1. **Countryside destinations:** If looking for a break from the urban bustle, consider a country house, summer cottage, or campsite. It can be a great way to enjoy nature and fresh air.\n\n2. **Exploring Russia:** Russia is a large country with diverse regions, each with unique sights and natural landmarks. Explore historical cities like St. Petersburg or Kazan, or embark on a trip around the Golden Ring of Russia.\n\n3. **Traveling abroad:** If you are looking for something more exotic, consider an international destination. Choose a place resonating with your interests – be it a beach stay, mountain ski resort, or a cultural landmark.\n\n4. **Outdoor recreation:** If you enjoy outdoor activities, consider hiking, fishing, or a picnic in nature. It can be the perfect way to spend time in fresh air and enjoy the beauty of nature.\n\n5. **Spa retreat or resort:** For those looking for something peaceful and relaxing, a spa retreat or resort could be a great option. Here you can relax, improve your health, and enjoy spa treatment.\n\nRemember that the choice of destination depends on your personal preferences and goals. Before you make your decision, research the place to make sure it meets your expectations. Question: Where to escape to? Answer with rephraser: 'Your question of "where to escape to?" may imply looking for a temporary retreat or a long-term relocation. The choice depends on your goals, interests, and resources. Here are some ideas for you:\n\n1. **Peaceful natural places:** If looking for solitude in nature, consider visiting national parks, nature reserves, or remote islands. It may prove a great place to unwind from the city bustle and recharge.\n\n2. **Historical cities:** If you are a history and culture geek, explore rich heritage sites like St. Petersburg, Prague, Rome, or Tokyo. Here you can dive into the past by visiting museums, castles, and other historical attractions.\n\n3. **Beach vacation:** If you prefer hitting the beach, consider tropical islands or coastlines with beautiful beaches and a warm climate. The options include Seychelles, Hawaii, the Maldives, etc.\n\n4. **Outdoor activities:** There are plenty of options for outdoor enthusiasts including mountain skiing, surfing, rafting, and other sports. Choose a place equipped for your favorite activities.\n\n5. **Moving to another country:** If you are considering a long-term relocation, explore countries offering favorable conditions for life, work, and education. Factor in things like climate, cost of living, safety, and local culture.\n\nRemember that the choice of destination depends on your personal preferences and goals. Before you make your decision, research the place to make sure it meets your expectations. Question: Wher to cuutt looooose? Answer with rephraser: 'It looks like your question contains an error or a typo. Please rephrase it and I will try to help you. If you meant "where to go for a break," I can suggest a few ideas:\n\n* **Nature reserves and national parks**: For nature and active lifestyle enthusiasts.\n* **Historical cities**: For culture and history geeks.\n* **Beach vacation**: For those who like to chill in the sun and enjoy sea air.\n* **Ski resorts**: For winter sports enthusiasts.\n\nThe choice of place for a break depends on your interests and preferences.
After rephrasing, the model managed to respond to
Wher to cuutt looooose?
and gave a more detailed answer toWhere to escape to?
.