CLI runs with OpenAI endpoints

app.py

@@ -4,7 +4,7 @@ import os
 
 from smolagents import CodeAgent, OpenAIServerModel, FinalAnswerTool
 
-from guest_info import guest_info_tool_factory
+from guest_info_retriever import guest_info_retriever_factory
 
 
 MODEL_ID = os.environ["MODEL_ID"]
@@ -16,7 +16,7 @@ model = OpenAIServerModel(model_id=MODEL_ID, api_base=API_BASE, api_key=API_KEY)
 
 
 # Load the guest dataset and initialize the guest info tool
-guest_info_tool = guest_info_tool_factory()
+guest_info_tool = guest_info_retriever_factory()
 
 final_answer_tool = FinalAnswerTool()
 

guest_info.py → guest_info_retriever.py

@@ -30,50 +30,34 @@ class GuestInfoRetrieverTool(Tool):
         return "No matching guest information found."
 
 
-def guest_info_tool_factory():
-    """Load the dataset"""
+def guest_info_retriever_factory():
+    """Get Guest Information Retrieval Tool"""
 
-    DATASETS = [  # pylint: disable=invalid-name
+    hf_datasets = [
         "agents-course/unit3-invitees",
-        "ANT-TECH/Unit3-dataset3",
         "Data-Gem/agents-course-unit3-invitees-expanded",
     ]
-    NAME = "name"  # pylint: disable=invalid-name
-    RELATION = "relation"  # pylint: disable=invalid-name
-    DESCRIPTION = "description"  # pylint: disable=invalid-name
-    EMAIL = "email"  # pylint: disable=invalid-name
-    COLUMN_LIST = [NAME, RELATION, DESCRIPTION, EMAIL]  # pylint: disable=invalid-name
 
-    raw_guest_list = []
-    guest_list = []
-    for dataset_name in DATASETS:
-        guest_dataset = datasets.load_dataset(dataset_name, split="train")
-        raw_guest_list.append(guest_dataset)
+    guest_dataset = []
 
-    for raw_guest in raw_guest_list:
-        guest = {}
-        for col_name in COLUMN_LIST:
-            if col_name in raw_guest:
-                guest[col_name] = raw_guest[col_name]
-            else:
-                guest[col_name] = f"{col_name} not in data".capitalize()
-        guest_list.append(guest)
+    for dataset_name in hf_datasets:
+        new_dataset = datasets.load_dataset(dataset_name, split="train")
+        guest_dataset.append(new_dataset)
 
     # Convert dataset entries into Document objects
     docs = [
         Document(
             page_content="\n".join(
                 [
-                    f"Name: {guest[NAME]}",
-                    f"Relation: {guest[RELATION]}",
-                    f"Description: {guest[DESCRIPTION]}",
-                    f"Email: {guest[EMAIL]}",
+                    f"Name: {guest["name"]}",
+                    f"Relation: {guest["relation"]}",
+                    f"Description: {guest["description"]}",
+                    f"Email: {guest["email"]}",
                 ]
             ),
-            metadata={"name": guest[NAME]},
+            metadata={"name": guest["name"]},
         )
-        for guest in guest_list
+        for guest in guest_dataset
     ]
 
-    # Return the tool
     return GuestInfoRetrieverTool(docs)