virtual-data-analyst / functions /chat_functions.py
Nolan Zandi
refactor functions and improve llm accuracy
3d660e2
raw
history blame
4.52 kB
from utils import TEMP_DIR, message_dict
from haystack.dataclasses import ChatMessage
from haystack.components.generators.chat import OpenAIChatGenerator
chat_generator = OpenAIChatGenerator(model="gpt-4o")
response = None
def example_question_generator(session_hash):
import sqlite3
example_response = None
example_messages = [
ChatMessage.from_system(
"You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'."
)
]
dir_path = TEMP_DIR / str(session_hash)
connection = sqlite3.connect(f'{dir_path}/data_source.db')
print("Querying questions");
cur=connection.execute('select * from data_source')
columns = [i[0] for i in cur.description]
print("QUESTION COLUMNS")
print(columns)
cur.close()
connection.close()
example_messages.append(ChatMessage.from_user(text=f"""We have a SQLite database with the following {columns}.
We also have an AI agent with access to the same database that will be performing data analysis.
Please return an array of seven strings, each one being a question for our data analysis agent
that we can suggest that you believe will be insightful or helpful to a data analysis looking for
data insights. Return nothing more than the array of questions because I need that specific data structure
to process your response. No other response type or data structure will work."""))
example_response = chat_generator.run(messages=example_messages)
return example_response["replies"][0].text
def chatbot_with_fc(message, history, session_hash):
from functions import sqlite_query_func, chart_generation_func, table_generation_func
import tools
available_functions = {"sql_query_func": sqlite_query_func, "chart_generation_func": chart_generation_func, "table_generation_func":table_generation_func }
if message_dict[session_hash] != None:
message_dict[session_hash].append(ChatMessage.from_user(message))
else:
messages = [
ChatMessage.from_system(
"""You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'.
You also have access to a chart function that uses plotly dictionaries to generate charts and graphs.
You also have access to a function, called table_generation_func, that builds table formatted html and generates a link to download as CSV."""
)
]
messages.append(ChatMessage.from_user(message))
message_dict[session_hash] = messages
response = chat_generator.run(messages=message_dict[session_hash], generation_kwargs={"tools": tools.tools_call(session_hash)})
while True:
# if OpenAI response is a tool call
if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
function_calls = response["replies"][0].tool_calls
for function_call in function_calls:
message_dict[session_hash].append(ChatMessage.from_assistant(tool_calls=[function_call]))
## Parse function calling information
function_name = function_call.tool_name
function_args = function_call.arguments
## Find the corresponding function and call it with the given arguments
function_to_call = available_functions[function_name]
function_response = function_to_call(**function_args, session_hash=session_hash)
print(function_name)
## Append function response to the messages list using `ChatMessage.from_tool`
message_dict[session_hash].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
response = chat_generator.run(messages=message_dict[session_hash], generation_kwargs={"tools": tools.tools_call(session_hash)})
# Regular Conversation
else:
message_dict[session_hash].append(response["replies"][0])
break
return response["replies"][0].text