Update main.py

main.py

@@ -381,6 +381,7 @@ async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verif
 from pydantic import BaseModel, Field
 import yaml
 import json
+from yaml.loader import SafeLoader
 
 class FollowupQueryModel(BaseModel):
     query: str = Field(..., description="User's query for the followup agent")