diff --git a/docs/swarms/agents/omni_agent.md b/docs/swarms/agents/omni_agent.md
index dcbf4bbe..cf65b380 100644
--- a/docs/swarms/agents/omni_agent.md
+++ b/docs/swarms/agents/omni_agent.md
@@ -32,6 +32,7 @@ During initialization, various tools like "document-question-answering", "image-
 
 #### Agent Runner (`run`):
 `run(self, input: str) -> str`:
+
 This method represents the primary operation of the OmniModalAgent. It takes an input, devises a plan using the chat planner, executes the plan with the task executor, and finally, the response generator crafts a response based on the tasks executed.
 
 Parameters:
diff --git a/tests/agents/omni_modal.py b/tests/agents/omni_modal.py
new file mode 100644
index 00000000..e8a602ba
--- /dev/null
+++ b/tests/agents/omni_modal.py
@@ -0,0 +1,40 @@
+import pytest
+from langchain.base_language import BaseLanguageModel
+from langchain_experimental.autonomous_agents.hugginggpt.repsonse_generator import (
+    load_response_generator,
+)
+from langchain_experimental.autonomous_agents.hugginggpt.task_executor import (
+    TaskExecutor,
+)
+from langchain_experimental.autonomous_agents.hugginggpt.task_planner import (
+    load_chat_planner,
+)
+from transformers import load_tool
+from swarms.agents import OmniModalAgent  # Replace `your_module_name` with the appropriate module name
+
+# Mock objects or set up fixtures for dependent classes or external methods
+@pytest.fixture
+def mock_llm():
+    # For this mock, we are assuming the BaseLanguageModel has a method named "process"
+    class MockLLM(BaseLanguageModel):
+        def process(self, input):
+            return "mock response"
+    
+    return MockLLM()
+
+@pytest.fixture
+def omni_agent(mock_llm):
+    return OmniModalAgent(mock_llm)
+
+def test_omnimodalagent_initialization(omni_agent):
+    assert omni_agent.llm is not None, "LLM initialization failed"
+    assert len(omni_agent.tools) > 0, "Tools initialization failed"
+
+def test_omnimodalagent_run(omni_agent):
+    input_string = "Hello, how are you?"
+    response = omni_agent.run(input_string)
+    assert response is not None, "Response generation failed"
+    assert isinstance(response, str), "Response should be a string"
+
+def test_task_executor_initialization(omni_agent):
+    assert omni_agent.task_executor is not None, "TaskExecutor initialization failed"
diff --git a/tests/agents/memory/main.py b/tests/memory/main.py
similarity index 100%
rename from tests/agents/memory/main.py
rename to tests/memory/main.py
diff --git a/tests/agents/models/LLM.py b/tests/models/LLM.py
similarity index 100%
rename from tests/agents/models/LLM.py
rename to tests/models/LLM.py
diff --git a/tests/agents/models/hf.py b/tests/models/hf.py
similarity index 100%
rename from tests/agents/models/hf.py
rename to tests/models/hf.py
diff --git a/tests/agents/workers/multi_model_worker.py b/tests/workers/multi_model_worker.py
similarity index 100%
rename from tests/agents/workers/multi_model_worker.py
rename to tests/workers/multi_model_worker.py
diff --git a/tests/agents/workers/omni_worker.py b/tests/workers/omni_worker.py
similarity index 100%
rename from tests/agents/workers/omni_worker.py
rename to tests/workers/omni_worker.py
diff --git a/tests/agents/workers/worker_agent_ultra.py b/tests/workers/worker_agent_ultra.py
similarity index 100%
rename from tests/agents/workers/worker_agent_ultra.py
rename to tests/workers/worker_agent_ultra.py
diff --git a/tests/agents/workers/worker_node.py b/tests/workers/worker_node.py
similarity index 100%
rename from tests/agents/workers/worker_node.py
rename to tests/workers/worker_node.py
diff --git a/tests/agents/workers/worker_ultra.py b/tests/workers/worker_ultra.py
similarity index 100%
rename from tests/agents/workers/worker_ultra.py
rename to tests/workers/worker_ultra.py