cfpb-assistant / agent.py
david-oplatka's picture
Add Assistant Files
b6fadc7
raw
history blame
4.14 kB
import os
from typing import Optional
from pydantic import Field, BaseModel
from omegaconf import OmegaConf
from llama_index.core.utilities.sql_wrapper import SQLDatabase
from sqlalchemy import create_engine, text
from dotenv import load_dotenv
load_dotenv(override=True)
from vectara_agentic.agent import Agent
from vectara_agentic.tools import ToolsFactory, VectaraToolFactory
def create_assistant_tools(cfg):
class QueryCFPBComplaints(BaseModel):
query: str = Field(description="The user query.")
vec_factory = VectaraToolFactory(vectara_api_key=cfg.api_keys,
vectara_customer_id=cfg.customer_id,
vectara_corpus_id=cfg.corpus_ids)
summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni'
ask_complaints = vec_factory.create_rag_tool(
tool_name = "ask_complaints",
tool_description = """
Given a user query,
returns a response to a user question about customer complaints about bank services.
""",
tool_args_schema = QueryCFPBComplaints,
reranker = "multilingual_reranker_v1", rerank_k = 100,
n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
summary_num_results = 5,
vectara_summarizer = summarizer,
include_citations = False,
)
tools_factory = ToolsFactory()
db_tools = tools_factory.database_tools(
tool_name_prefix = "cfpb",
content_description = 'Customer complaints about five banks (Bank of America, Wells Fargo, Capital One, Chase, and CITI Bank)',
sql_database = SQLDatabase(create_engine('sqlite:///cfpb_database.db')),
)
return (tools_factory.standard_tools() +
tools_factory.guardrail_tools() +
db_tools +
[ask_complaints]
)
def initialize_agent(_cfg, update_func=None):
cfpb_complaints_bot_instructions = """
- You are a helpful research assistant, with expertise in complaints from the Consumer Financial Protection Bureau, in conversation with a user.
- Before answering any user query, use cfpb_describe_tables to understand schema of each table, and use get_sample_data
to get sample data from each table in the database, so that you can understand NULL and unique values for each column.
- For a query with multiple sub-questions, break down the query into the sub-questions,
and make separate calls to the ask_complaints tool to answer each sub-question,
then combine the answers to provide a complete response.
- Use the database tools (cfpb_load_data, cfpb_describe_tables and cfpb_list_tables) to answer analytical queries.
- IMPORTANT: When using database_tools, always call the ev_load_sample_data tool with the table you want to query
to understand the table structure, column naming, and values in the table. Never call the cfpb_load_data tool for a query until you have called cfpb_load_sample_data.
- When providing links, try to put the name of the website or source of information for the displayed text. Don't just say 'Source'.
- Never discuss politics, and always respond politely.
"""
agent = Agent(
tools=create_assistant_tools(_cfg),
topic="Customer complaints from the Consumer Financial Protection Bureau (CFPB)",
custom_instructions=cfpb_complaints_bot_instructions,
update_func=update_func
)
agent.report()
return agent
def get_agent_config() -> OmegaConf:
cfg = OmegaConf.create({
'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
'corpus_ids': str(os.environ['VECTARA_CORPUS_IDS']),
'api_keys': str(os.environ['VECTARA_API_KEYS']),
'examples': os.environ.get('QUERY_EXAMPLES', None),
'demo_name': "cfpb-assistant",
'demo_welcome': "Welcome to the CFPB Customer Complaints demo.",
'demo_description': "This assistant can help you gain insights into customer complaints to banks recorded by the Consumer Financial Protection Bureau.",
})
return cfg