ev-assistant / agent.py
ofermend's picture
initial
49b0a2d
raw
history blame
4.9 kB
from omegaconf import OmegaConf
import os
from typing import Optional
from pydantic import Field, BaseModel
from dotenv import load_dotenv
load_dotenv(override=True)
from vectara_agent.agent import Agent
from vectara_agent.tools import ToolsFactory, VectaraToolFactory
def create_assistant_tools(cfg):
class QueryElectricCars(BaseModel):
query: str = Field(description="The user query.")
vec_factory_1 = VectaraToolFactory(vectara_api_key=cfg.api_keys[0],
vectara_customer_id=cfg.customer_id,
vectara_corpus_id=cfg.corpus_ids[0])
ask_vehicles = vec_factory_1.create_rag_tool(
tool_name = "ask_vehicles",
tool_description = """
Given a user query,
returns a response (str) to a user question about electric vehicles based on online resources.
You can ask this tool any question about electric cars, including the different types of EVs, how they work, the pros and cons of different models, the environmental impact, and more.
""",
tool_args_schema = QueryElectricCars,
reranker = "multilingual_reranker_v1", rerank_k = 100,
n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
summary_num_results = 10,
vectara_summarizer = 'vectara-summary-ext-24-05-sml',
include_citations = False,
)
vec_factory_2 = VectaraToolFactory(vectara_api_key=cfg.api_keys[1],
vectara_customer_id=cfg.customer_id,
vectara_corpus_id=cfg.corpus_ids[1])
class QueryEVLaws(BaseModel):
query: str = Field(description="The user query")
state: Optional[str] = Field(default=None,
description="The two digit state code. Optional.",
examples=['CA', 'US', 'WA'])
type: Optional[str] = Field(default=None,
description="The type of policy. Optional",
examples = ['Laws and Regulations', 'State Incentives', 'Incentives', 'Utility / Private Incentives', 'Programs'])
ask_policies = vec_factory_2.create_rag_tool(
tool_name = "ask_policies",
tool_description = """
Given a user query,
returns a response (str) to a user question about incentives and regulations about electric vehicles in the United States.
You can ask this tool any question about laws passed by states or the federal government related to electric vehicles.
""",
tool_args_schema = QueryEVLaws,
reranker = "multilingual_reranker_v1", rerank_k = 100,
n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005,
summary_num_results = 10,
vectara_summarizer = 'vectara-summary-ext-24-05-sml',
include_citations = False,
)
tools_factory = ToolsFactory()
return (tools_factory.standard_tools() +
tools_factory.guardrail_tools() +
tools_factory.database_tools(
content_description = 'Electric Vehicles',
scheme = 'postgresql',
host = 'localhost', port = '5432',
user = 'ofer',
password = 'noanoa',
dbname = 'ev_database'
) +
[ask_vehicles, ask_policies]
)
def initialize_agent(_cfg, update_func):
electric_vehicle_bot_instructions = """
- You are a helpful research assistant, with expertise in electric vehicles, in conversation with a user.
- For a query with multiple sub-questions, break down the query into the sub-questions,
and make separate calls to the ask_vehicles or ask_policies tool to answer each sub-question,
then combine the answers to provide a complete response.
- Never discuss politics, and always respond politely.
"""
agent = Agent(
tools=create_assistant_tools(_cfg),
topic="Electric vehicles in the United States",
custom_instructions=electric_vehicle_bot_instructions,
update_func=update_func
)
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']).split(','),
'api_keys': str(os.environ['VECTARA_API_KEYS']).split(','),
'examples': os.environ.get('QUERY_EXAMPLES', None),
'title': "Electric Vehicles in the United States",
'demo_welcome': "Welcome to the EV Assistant demo.",
'demo_description': "This assistant can help you learn about electric vehicles in the United States, including how they work, the advantages of purchasing them, and reviews on the top choices.",
})
return cfg