EDS-Chatbot / app.py
hayuh's picture
Update app.py
7628860 verified
import os
import glob
from pathlib import Path
import gradio as gr
import nest_asyncio
# Ensure async compatibility in Jupyter
nest_asyncio.apply()
# Import OpenAI key with helper function
from helper import get_openai_api_key
OPENAI_API_KEY = get_openai_api_key()
# Define the path to the directory containing the PDF files
folder_path = 'Ehlers-Danlos-1'
# Get the list of all PDF files in the directory
pdf_files = glob.glob(os.path.join(folder_path, '*.pdf'))
print(pdf_files)
# Extract just the filenames (optional)
pdf_filenames = [os.path.basename(pdf) for pdf in pdf_files]
print(pdf_filenames)
# Import utilities
from utils import get_doc_tools
# Truncate function names if necessary
def truncate_function_name(name, max_length=64):
return name if len(name) <= max_length else name[:max_length]
# Create tools for each PDF
paper_to_tools_dict = {}
for pdf in pdf_files:
print(f"Getting tools for paper: {pdf}")
vector_tool, summary_tool = get_doc_tools(pdf, Path(pdf).stem)
paper_to_tools_dict[pdf] = [vector_tool, summary_tool]
# Combine all tools into a single list
all_tools = [t for pdf in pdf_files for t in paper_to_tools_dict[pdf]]
# Define an object index and retriever over these tools
from llama_index.core import VectorStoreIndex
from llama_index.core.objects import ObjectIndex
obj_index = ObjectIndex.from_objects(
all_tools,
index_cls=VectorStoreIndex,
)
obj_retriever = obj_index.as_retriever(similarity_top_k=3)
# Initialize the OpenAI LLM
from llama_index.llms.openai import OpenAI
llm = OpenAI(model="gpt-3.5-turbo")
# Set up the agent
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner
agent_worker = FunctionCallingAgentWorker.from_tools(
tool_retriever=obj_retriever,
llm=llm,
verbose=True
)
agent = AgentRunner(agent_worker)
# Define the function to query the agent
def ask_agent(question):
response = agent.query(question)
return str(response)
# Create the Gradio interface
iface = gr.Interface(
fn=ask_agent,
inputs="text",
outputs="text",
title="EDS Diagnosis Helper",
description="Ask questions related to Ehlers-Danlos Syndrome diagnosis.",
)
# Launch the Gradio app
iface.launch()