From 96b9ed1bca7f27a696b795db4372884e33e9fda7 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Wed, 5 Jul 2023 17:20:05 -0400
Subject: [PATCH] clean up

---
 swarms/swarms.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/swarms/swarms.py b/swarms/swarms.py
index f32810b1..2146f687 100644
--- a/swarms/swarms.py
+++ b/swarms/swarms.py
@@ -2,7 +2,7 @@
 import os
 import asyncio
 import faiss
-from typing import Optional
+from typing import Any, Optional, List
 from contextlib import contextmanager
 
 from pydantic import BaseModel, Field
@@ -59,12 +59,12 @@ vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {
 
 # ---------- Worker Node ----------
 
+# Define the input schema for the WorkerNode
 # Define the input schema for the WorkerNode
 class WorkerNodeInput(BaseModel):
-    ai_name: str = Field(description="Name of the AI")
-    ai_role: str = Field(description="Role of the AI")
-    human_in_the_loop: bool = Field(description="Whether there is a human in the loop")
-    search_kwargs: dict = Field(description="Search parameters")
+    llm: Any = Field(description="Language model")
+    tools: List[Tool] = Field(description="List of tools")
+    vectorstore: VectorStore = Field(description="Vector store")
 
 
 @tool("WorkerNode", args_schema=WorkerNodeInput)