Skip to content

Commit

Permalink
feat: persistent valves
Browse files Browse the repository at this point in the history
  • Loading branch information
tjbck committed Jun 1, 2024
1 parent 8aa82f9 commit 3d758a9
Showing 1 changed file with 37 additions and 0 deletions.
37 changes: 37 additions & 0 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,8 +124,37 @@ async def load_modules_from_directory(directory):
if filename.endswith(".py"):
module_name = filename[:-3] # Remove the .py extension
module_path = os.path.join(directory, filename)

# Create subfolder matching the filename without the .py extension
subfolder_path = os.path.join(directory, module_name)
if not os.path.exists(subfolder_path):
os.makedirs(subfolder_path)
logging.info(f"Created subfolder: {subfolder_path}")

# Create a valves.json file if it doesn't exist
valves_json_path = os.path.join(subfolder_path, "valves.json")
if not os.path.exists(valves_json_path):
with open(valves_json_path, "w") as f:
json.dump({}, f)
logging.info(f"Created valves.json in: {subfolder_path}")

pipeline = await load_module_from_path(module_name, module_path)
if pipeline:
# Overwrite pipeline.valves with values from valves.json
if os.path.exists(valves_json_path):
with open(valves_json_path, "r") as f:
valves_json = json.load(f)
ValvesModel = pipeline.valves.__class__
# Create a ValvesModel instance using default values and overwrite with valves_json
combined_valves = {
**pipeline.valves.model_dump(),
**valves_json,
}
valves = ValvesModel(**combined_valves)
pipeline.valves = valves

logging.info(f"Updated valves for module: {module_name}")

pipeline_id = pipeline.id if hasattr(pipeline, "id") else module_name
PIPELINE_MODULES[pipeline_id] = pipeline
PIPELINE_NAMES[pipeline_id] = module_name
Expand Down Expand Up @@ -441,6 +470,14 @@ async def update_valves(pipeline_id: str, form_data: dict):
valves = ValvesModel(**form_data)
pipeline.valves = valves

# Determine the directory path for the valves.json file
subfolder_path = os.path.join(PIPELINES_DIR, PIPELINE_NAMES[pipeline_id])
valves_json_path = os.path.join(subfolder_path, "valves.json")

# Save the updated valves data back to the valves.json file
with open(valves_json_path, "w") as f:
json.dump(valves.model_dump(), f)

if hasattr(pipeline, "on_valves_updated"):
await pipeline.on_valves_updated()
except Exception as e:
Expand Down

0 comments on commit 3d758a9

Please sign in to comment.