import json
import uuid
from typing import Optional, Dict, Any, List
from openai import OpenAI
import yfinance as yf
from dotenv import load_dotenv
import os
from helicone_helpers import HeliconeManualLogger
# Load environment variables
load_dotenv()
class StockInfoAgent:
def __init__(self):
# Initialize OpenAI client with Helicone
self.client = OpenAI(
api_key=os.getenv('OPENAI_API_KEY'),
base_url="https://oai.helicone.ai/v1",
default_headers={
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}"
}
)
# Initialize Helicone manual logger for tool calls
self.helicone_logger = HeliconeManualLogger(
api_key=os.getenv('HELICONE_API_KEY'),
headers={
"Helicone-Property-Type": "Stock-Info-Agent",
}
)
self.conversation_history = []
self.session_id = None
self.session_headers = {}
def start_new_session(self):
"""Initialize a new session for tracking."""
self.session_id = str(uuid.uuid4())
self.session_headers = {
"Helicone-Session-Id": self.session_id,
"Helicone-Session-Name": "Stock Information Chat",
"Helicone-Session-Path": "/stock-chat",
"Helicone-Property-Environment": "production"
}
print(f"Started new session: {self.session_id}")
def get_stock_price(self, ticker_symbol: str) -> Optional[str]:
"""Fetches the current stock price for the given ticker_symbol with Helicone logging."""
def price_operation(result_recorder):
try:
stock = yf.Ticker(ticker_symbol.upper())
info = stock.info
current_price = info.get('currentPrice') or info.get('regularMarketPrice')
if current_price:
result = f"{current_price:.2f} USD"
result_recorder.append_results({
"ticker": ticker_symbol.upper(),
"price": current_price,
"formatted_price": result,
"status": "success"
})
return result
else:
result_recorder.append_results({
"ticker": ticker_symbol.upper(),
"error": "Price not found",
"status": "error"
})
return None
except Exception as e:
result_recorder.append_results({
"ticker": ticker_symbol.upper(),
"error": str(e),
"status": "error"
})
print(f"Error fetching stock price: {e}")
return None
# Log the tool call with Helicone
return self.helicone_logger.log_request(
provider=None,
request={
"_type": "tool",
"toolName": "get_stock_price",
"input": {"ticker_symbol": ticker_symbol},
"metadata": {
"source": "yfinance",
"operation": "get_current_price"
}
},
operation=price_operation,
additional_headers={
**self.session_headers,
"Helicone-Session-Path": f"/stock-chat/price/{ticker_symbol.lower()}"
}
)
def get_company_ceo(self, ticker_symbol: str) -> Optional[str]:
"""Fetches the name of the CEO for the company with Helicone logging."""
def ceo_operation(result_recorder):
try:
stock = yf.Ticker(ticker_symbol.upper())
info = stock.info
# Look for CEO in various possible fields
ceo = None
for field in ['companyOfficers', 'officers']:
if field in info:
officers = info[field]
if isinstance(officers, list):
for officer in officers:
if isinstance(officer, dict):
title = officer.get('title', '').lower()
if 'ceo' in title or 'chief executive' in title:
ceo = officer.get('name')
break
result_recorder.append_results({
"ticker": ticker_symbol.upper(),
"ceo": ceo,
"status": "success" if ceo else "not_found"
})
return ceo
except Exception as e:
result_recorder.append_results({
"ticker": ticker_symbol.upper(),
"error": str(e),
"status": "error"
})
print(f"Error fetching CEO info: {e}")
return None
return self.helicone_logger.log_request(
provider=None,
request={
"_type": "tool",
"toolName": "get_company_ceo",
"input": {"ticker_symbol": ticker_symbol},
"metadata": {
"source": "yfinance",
"operation": "get_company_officers"
}
},
operation=ceo_operation,
additional_headers={
**self.session_headers,
"Helicone-Session-Path": f"/stock-chat/ceo/{ticker_symbol.lower()}"
}
)
def find_ticker_symbol(self, company_name: str) -> Optional[str]:
"""Tries to identify the stock ticker symbol with Helicone logging."""
def ticker_search_operation(result_recorder):
try:
# Use yfinance Lookup to search for the company
lookup = yf.Lookup(company_name)
stock_results = lookup.get_stock(count=5)
if not stock_results.empty:
ticker = stock_results.index[0]
result_recorder.append_results({
"company_name": company_name,
"ticker": ticker,
"search_type": "stock",
"results_count": len(stock_results),
"status": "success"
})
return ticker
# If no stocks found, try all instruments
all_results = lookup.get_all(count=5)
if not all_results.empty:
ticker = all_results.index[0]
result_recorder.append_results({
"company_name": company_name,
"ticker": ticker,
"search_type": "all_instruments",
"results_count": len(all_results),
"status": "success"
})
return ticker
result_recorder.append_results({
"company_name": company_name,
"error": "No ticker found",
"status": "not_found"
})
return None
except Exception as e:
result_recorder.append_results({
"company_name": company_name,
"error": str(e),
"status": "error"
})
print(f"Error searching for ticker: {e}")
return None
return self.helicone_logger.log_request(
provider=None,
request={
"_type": "tool",
"toolName": "find_ticker_symbol",
"input": {"company_name": company_name},
"metadata": {
"source": "yfinance_lookup",
"operation": "ticker_search"
}
},
operation=ticker_search_operation,
additional_headers={
**self.session_headers,
"Helicone-Session-Path": f"/stock-chat/search/{company_name.lower().replace(' ', '-')}"
}
)
def create_tool_definitions(self) -> List[Dict[str, Any]]:
"""Creates OpenAI function calling definitions for the tools."""
return [
{
"type": "function",
"function": {
"name": "get_stock_price",
"description": "Fetches the current stock price for the given ticker symbol",
"parameters": {
"type": "object",
"properties": {
"ticker_symbol": {
"type": "string",
"description": "The stock ticker symbol (e.g., 'AAPL', 'MSFT')"
}
},
"required": ["ticker_symbol"]
}
}
},
{
"type": "function",
"function": {
"name": "get_company_ceo",
"description": "Fetches the name of the CEO for the company associated with the ticker symbol",
"parameters": {
"type": "object",
"properties": {
"ticker_symbol": {
"type": "string",
"description": "The stock ticker symbol"
}
},
"required": ["ticker_symbol"]
}
}
},
{
"type": "function",
"function": {
"name": "find_ticker_symbol",
"description": "Tries to identify the stock ticker symbol for a given company name",
"parameters": {
"type": "object",
"properties": {
"company_name": {
"type": "string",
"description": "The name of the company"
}
},
"required": ["company_name"]
}
}
}
]
def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Executes the specified tool with given arguments."""
if tool_name == "get_stock_price":
return self.get_stock_price(arguments["ticker_symbol"])
elif tool_name == "get_company_ceo":
return self.get_company_ceo(arguments["ticker_symbol"])
elif tool_name == "find_ticker_symbol":
return self.find_ticker_symbol(arguments["company_name"])
else:
return None
def process_user_query(self, user_query: str) -> str:
"""Processes a user query using the OpenAI API with function calling and Helicone logging."""
# Add user message to conversation history
self.conversation_history.append({"role": "user", "content": user_query})
# System prompt to guide the agent's behavior
system_prompt = """You are a helpful stock information assistant. You have access to tools that can:
1. Get current stock prices
2. Find company CEOs
3. Find ticker symbols for company names
4. Ask users for clarification when needed
Use these tools one at a time to help answer user questions about stocks and companies. If information is ambiguous, ask for clarification."""
while True:
messages = [
{"role": "system", "content": system_prompt},
*self.conversation_history
]
def openai_operation(result_recorder):
# Call OpenAI API with function calling
response = self.client.chat.completions.create(
model="gpt-4o-mini-2024-07-18",
messages=messages,
tools=self.create_tool_definitions(),
tool_choice="auto"
)
# Log the response
result_recorder.append_results({
"model": "gpt-4o-mini-2024-07-18",
"response": response.choices[0].message.model_dump(),
"usage": response.usage.model_dump() if response.usage else None
})
return response
# Log the OpenAI call
response = self.helicone_logger.log_request(
provider="openai",
request={
"model": "gpt-4o-mini-2024-07-18",
"messages": messages,
"tools": self.create_tool_definitions(),
"tool_choice": "auto"
},
operation=openai_operation,
additional_headers={
**self.session_headers,
"Helicone-Prompt-Id": "stock-agent-reasoning"
}
)
response_message = response.choices[0].message
# If no tool calls, we're done
if not response_message.tool_calls:
self.conversation_history.append({"role": "assistant", "content": response_message.content})
return response_message.content
# Execute the first tool call
tool_call = response_message.tool_calls[0]
function_name = tool_call.function.name
function_args = json.loads(tool_call.function.arguments)
print(f"\nExecuting tool: {function_name} with args: {function_args}")
# Execute the tool (this will be logged separately by each tool method)
result = self.execute_tool(function_name, function_args)
# Add the assistant's message with tool calls to history
self.conversation_history.append({
"role": "assistant",
"content": None,
"tool_calls": [{
"id": tool_call.id,
"type": "function",
"function": {
"name": function_name,
"arguments": json.dumps(function_args)
}
}]
})
# Add tool result to history
self.conversation_history.append({
"tool_call_id": tool_call.id,
"role": "tool",
"name": function_name,
"content": str(result) if result is not None else "No result found"
})
def chat(self):
"""Interactive chat loop with session tracking."""
print("Stock Information Agent with Helicone Monitoring")
print("Ask me about stock prices, company CEOs, or any stock-related questions!")
print("Type 'quit' to exit.\n")
# Start a new session
self.start_new_session()
while True:
user_input = input("You: ")
if user_input.lower() in ['quit', 'exit', 'bye']:
print("Goodbye!")
break
try:
response = self.process_user_query(user_input)
print(f"\nAgent: {response}\n")
except Exception as e:
print(f"\nError: {e}\n")
if __name__ == "__main__":
agent = StockInfoAgent()
agent.chat()
