|
import os |
|
from typing import Optional |
|
from pydantic import Field, BaseModel |
|
from omegaconf import OmegaConf |
|
|
|
from vectara_agentic.agent import Agent |
|
from vectara_agentic.tools import VectaraToolFactory |
|
|
|
from dotenv import load_dotenv |
|
load_dotenv(override=True) |
|
|
|
initial_prompt = "How can I help you today?" |
|
|
|
def create_assistant_tools(cfg): |
|
|
|
class QueryHMC(BaseModel): |
|
query: str = Field(description="The user query.") |
|
ticker: Optional[str] = Field( |
|
default=None, |
|
description="The company ticker.", |
|
examples=['GOOG', 'META'] |
|
) |
|
year: Optional[str] = Field( |
|
default=None, |
|
description="The year of the report.", |
|
examples=[2020, 2023] |
|
) |
|
quarter: Optional[int] = Field( |
|
default=None, |
|
description="The quarter of the report.", |
|
examples=[1, 2, 3, 4] |
|
) |
|
filing_type: Optional[str] = Field( |
|
default=None, |
|
description="The type of filing.", |
|
examples=['10K', '10Q'] |
|
) |
|
|
|
vec_factory = VectaraToolFactory( |
|
vectara_api_key=cfg.api_key, |
|
vectara_customer_id=cfg.customer_id, |
|
vectara_corpus_id=cfg.corpus_id |
|
) |
|
|
|
summarizer = 'vectara-experimental-summary-ext-2023-12-11-med-omni' |
|
|
|
ask_hmc = vec_factory.create_rag_tool( |
|
tool_name = "ask_hmc", |
|
tool_description = """ |
|
Given a user query, |
|
returns a response to a user question about fund management companies. |
|
""", |
|
tool_args_schema = QueryHMC, |
|
reranker = "chain", rerank_k = 100, |
|
rerank_chain = [ |
|
{ |
|
"type": "slingshot", |
|
"cutoff": 0.2 |
|
}, |
|
{ |
|
"type": "mmr", |
|
"diversity_bias": 0.05, |
|
"limit": 20 |
|
} |
|
], |
|
n_sentences_before = 2, n_sentences_after = 2, lambda_val = 0.005, |
|
vectara_summarizer = summarizer, |
|
summary_num_results = 10, |
|
include_citations = True, |
|
) |
|
return [ask_hmc] |
|
|
|
def initialize_agent(_cfg, agent_progress_callback=None): |
|
bot_instructions = """ |
|
- You are a helpful assistant, with expertise in management of public company stock portfolios. |
|
- Use the 'ask_hmc' tool to answer questions about public company performance, risks, and other financial metrics. |
|
If the tool responds with "I don't have enough information to answer", try rephrasing the question. |
|
- Use the year, quarter, filing_type and ticker arguments to the 'ask_hmc' tool to get more specific answers. |
|
- Note that 10Q reports exist for quarters 1, 2, 3 and for the 4th quarter there is a 10K report. |
|
""" |
|
|
|
agent = Agent( |
|
tools=create_assistant_tools(_cfg), |
|
topic="Endowment fund management", |
|
custom_instructions=bot_instructions, |
|
agent_progress_callback=agent_progress_callback, |
|
) |
|
agent.report() |
|
return agent |
|
|
|
|
|
def get_agent_config() -> OmegaConf: |
|
cfg = OmegaConf.create({ |
|
'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']), |
|
'corpus_id': str(os.environ['VECTARA_CORPUS_ID']), |
|
'api_key': str(os.environ['VECTARA_API_KEY']), |
|
'examples': os.environ.get('QUERY_EXAMPLES', None), |
|
'demo_name': "Harvard Management Company", |
|
'demo_welcome': "Harvard Management Company.", |
|
'demo_description': "AI Assistant.", |
|
}) |
|
return cfg |
|
|
