awesome-llm-apps/local_llama3.1_tool_use/llama3_tool_use.py
2024-08-03 21:23:47 -05:00

90 lines
3.1 KiB
Python

import streamlit as st
import os
from phi.assistant import Assistant
from phi.llm.ollama import Ollama
from phi.tools.yfinance import YFinanceTools
from phi.tools.serpapi_tools import SerpApiTools
st.set_page_config(page_title="Llama-3 Tool Use", page_icon="🦙")
# Ensure SERPAPI_API_KEY is set
if 'SERPAPI_API_KEY' not in os.environ:
st.error("Please set the SERPAPI_API_KEY environment variable.")
st.stop()
def get_assistant(tools):
return Assistant(
name="llama3_assistant",
llm=Ollama(model="llama3"),
tools=tools,
description="You are a helpful assistant that can access specific tools based on user selection.",
show_tool_calls=True,
debug_mode=True,
# This setting adds the current datetime to the instructions
add_datetime_to_instructions=True,
)
st.title("🦙 Local Llama-3 Tool Use")
st.markdown("""
This app demonstrates function calling with the local Llama3 model using Ollama.
Select tools in the sidebar and ask relevant questions!
""")
# Sidebar for tool selection
st.sidebar.title("Tool Selection")
use_yfinance = st.sidebar.checkbox("YFinance (Stock Data)", value=True)
use_serpapi = st.sidebar.checkbox("SerpAPI (Web Search)", value=True)
# Initialize or update the assistant based on selected tools
tools = []
if use_yfinance:
tools.append(YFinanceTools(stock_price=True, company_info=True))
if use_serpapi:
tools.append(SerpApiTools())
if "assistant" not in st.session_state or st.session_state.get("tools") != tools:
st.session_state.assistant = get_assistant(tools)
st.session_state.tools = tools
st.session_state.messages = [] # Reset messages when tools change
# Display current tool status
st.sidebar.markdown("### Current Tools:")
st.sidebar.markdown(f"- YFinance: {'Enabled' if use_yfinance else 'Disabled'}")
st.sidebar.markdown(f"- SerpAPI: {'Enabled' if use_serpapi else 'Disabled'}")
# Chat interface
for message in st.session_state.get("messages", []):
with st.chat_message(message["role"]):
st.markdown(message["content"])
if prompt := st.chat_input("Ask a question based on the enabled tools"):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
response_container = st.empty()
response = ""
for chunk in st.session_state.assistant.run(prompt):
response += chunk
response_container.write(response + "")
response_container.write(response)
st.session_state.messages.append({"role": "assistant", "content": response})
# Sidebar instructions
st.sidebar.markdown("""
### How to use:
1. Select the tools you want to use in the sidebar
2. Ask questions related to the enabled tools
3. The assistant will use only the selected tools to answe
### Note:
Make sure you have set the SERPAPI_API_KEY environment variable to use the SerpAPI tool.
""")
st.sidebar.markdown("""
### Sample questions:
- YFinance: "What's the current price of AAPL?"
- SerpAPI: "What are the latest developments in AI?"
- Both: "Compare TSLA stock price with recent news about Tesla's performance"
""")