Fix failing PeasyAI tests: RAG import path, langchain import, and chunk test

- Fix sys.path in tests/rag/test_rag_index.py to point to src/rag/
- Add pytest.importorskip for faiss graceful skip when not installed
- Update create_rag_index.py import from langchain.text_splitter to
  langchain_text_splitters (current package layout)
- Fix test_create_chunks to use text longer than chunk_size (500 chars)
  so the splitter actually produces multiple chunks

Author: ankushdesai

Src/PeasyAI/scripts/mcp_e2e_protocols.py

@@ -188,10 +188,11 @@ def run_protocol_ensemble(
     return result
 
 
-def run_protocol(tools: Dict[str, Any], design_doc: str, out_root: Path, project_name: str) -> Dict[str, Any]:
+def run_protocol(tools: Dict[str, Any], design_doc: str, out_root: Path, project_name: str, ensemble_size: int = 3) -> Dict[str, Any]:
     """
-    Run a protocol end-to-end using the step-by-step flow (legacy).
-    Kept for backward compatibility; prefer ``run_protocol_ensemble``.
+    Run a protocol end-to-end using the step-by-step flow.
+    Now uses ensemble generation (best-of-N) for each file by default.
+    For the single-call ensemble flow, prefer ``run_protocol_ensemble``.
     """
     result: Dict[str, Any] = {
         "project_name": project_name,
@@ -242,6 +243,7 @@ def run_protocol(tools: Dict[str, Any], design_doc: str, out_root: Path, project
                 design_doc=design_doc,
                 project_path=project_path,
                 context_files=context_files,
+                ensemble_size=ensemble_size,
             )
         )
         generated[f"machine:{machine_name}"] = machine_resp
@@ -255,6 +257,7 @@ def run_protocol(tools: Dict[str, Any], design_doc: str, out_root: Path, project
             design_doc=design_doc,
             project_path=project_path,
             context_files=context_files,
+            ensemble_size=ensemble_size,
         )
     )
     generated["spec"] = spec_resp
@@ -266,6 +269,7 @@ def run_protocol(tools: Dict[str, Any], design_doc: str, out_root: Path, project
             design_doc=design_doc,
             project_path=project_path,
             context_files=context_files,
+            ensemble_size=ensemble_size,
         )
     )
     generated["test"] = test_resp

Src/PeasyAI/src/core/workflow/factory.py

@@ -67,7 +67,8 @@ def create_full_generation_workflow(
         self,
         machine_names: List[str],
         spec_name: str = "Safety",
-        test_name: str = "TestDriver"
+        test_name: str = "TestDriver",
+        ensemble_size: int = 3,
     ) -> WorkflowDefinition:
         """
         Create a full project generation workflow.
@@ -76,6 +77,7 @@ def create_full_generation_workflow(
             machine_names: List of machine names to generate
             spec_name: Name for the specification file
             test_name: Name for the test file
+            ensemble_size: Number of candidates per file for ensemble selection
             
         Returns:
             WorkflowDefinition ready for execution
@@ -88,12 +90,12 @@ def create_full_generation_workflow(
         # Add machine generation steps
         for machine_name in machine_names:
             steps.append(
-                GenerateMachineStep(self.generation_service, machine_name)
+                GenerateMachineStep(self.generation_service, machine_name, ensemble_size=ensemble_size)
             )
 
         # Add spec and test
-        steps.append(GenerateSpecStep(self.generation_service, spec_name))
-        steps.append(GenerateTestStep(self.generation_service, test_name))
+        steps.append(GenerateSpecStep(self.generation_service, spec_name, ensemble_size=ensemble_size))
+        steps.append(GenerateTestStep(self.generation_service, test_name, ensemble_size=ensemble_size))
 
         # Save files
         steps.append(SaveGeneratedFilesStep())
@@ -109,27 +111,27 @@ def create_full_generation_workflow(
             continue_on_failure=False
         )
 
-    def create_add_machine_workflow(self, machine_name: str) -> WorkflowDefinition:
+    def create_add_machine_workflow(self, machine_name: str, ensemble_size: int = 3) -> WorkflowDefinition:
         """Create workflow to add a single machine to existing project."""
         return WorkflowDefinition(
             name=f"add_machine_{machine_name}",
             description=f"Add {machine_name} machine to project",
             steps=[
-                GenerateMachineStep(self.generation_service, machine_name),
+                GenerateMachineStep(self.generation_service, machine_name, ensemble_size=ensemble_size),
                 SaveGeneratedFilesStep(),
                 CompileProjectStep(self.compilation_service),
                 FixCompilationErrorsStep(self.fixer_service),
             ],
             continue_on_failure=False
         )
 
-    def create_add_spec_workflow(self, spec_name: str) -> WorkflowDefinition:
+    def create_add_spec_workflow(self, spec_name: str, ensemble_size: int = 3) -> WorkflowDefinition:
         """Create workflow to add a specification to existing project."""
         return WorkflowDefinition(
             name=f"add_spec_{spec_name}",
             description=f"Add {spec_name} specification to project",
             steps=[
-                GenerateSpecStep(self.generation_service, spec_name),
+                GenerateSpecStep(self.generation_service, spec_name, ensemble_size=ensemble_size),
                 SaveGeneratedFilesStep(),
                 CompileProjectStep(self.compilation_service),
                 FixCompilationErrorsStep(self.fixer_service),

Src/PeasyAI/src/core/workflow/p_steps.py

@@ -120,14 +120,16 @@ class GenerateMachineStep(WorkflowStep):
     description = "Generate a P state machine implementation"
     max_retries = 3
 
-    def __init__(self, generation_service: GenerationService, machine_name: str):
+    def __init__(self, generation_service: GenerationService, machine_name: str, ensemble_size: int = 3):
         self.service = generation_service
         self.machine_name = machine_name
+        self.ensemble_size = ensemble_size
         self.name = f"generate_machine_{machine_name}"
 
     def execute(self, context: Dict[str, Any]) -> StepResult:
         design_doc = context.get("design_doc")
         project_path = context.get("project_path")
+        ensemble_size = context.get("ensemble_size", self.ensemble_size)
 
         if not design_doc:
             return StepResult.failure("design_doc is required")
@@ -148,13 +150,23 @@ def execute(self, context: Dict[str, Any]) -> StepResult:
                 context_files[machine_file] = value
 
         try:
-            result = self.service.generate_machine(
-                machine_name=self.machine_name,
-                design_doc=design_doc,
-                project_path=project_path,
-                context_files=context_files,
-                save_to_disk=False  # Preview mode
-            )
+            if ensemble_size > 1:
+                result = self.service.generate_machine_ensemble(
+                    machine_name=self.machine_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    ensemble_size=ensemble_size,
+                    save_to_disk=False  # Preview mode
+                )
+            else:
+                result = self.service.generate_machine(
+                    machine_name=self.machine_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    save_to_disk=False  # Preview mode
+                )
 
             if result.success:
                 return StepResult.success(
@@ -185,14 +197,16 @@ class GenerateSpecStep(WorkflowStep):
     description = "Generate a P specification/monitor file"
     max_retries = 3
 
-    def __init__(self, generation_service: GenerationService, spec_name: str = "Safety"):
+    def __init__(self, generation_service: GenerationService, spec_name: str = "Safety", ensemble_size: int = 3):
         self.service = generation_service
         self.spec_name = spec_name
+        self.ensemble_size = ensemble_size
         self.name = f"generate_spec_{spec_name}"
 
     def execute(self, context: Dict[str, Any]) -> StepResult:
         design_doc = context.get("design_doc")
         project_path = context.get("project_path")
+        ensemble_size = context.get("ensemble_size", self.ensemble_size)
 
         if not design_doc:
             return StepResult.failure("design_doc is required")
@@ -203,13 +217,23 @@ def execute(self, context: Dict[str, Any]) -> StepResult:
         context_files = self._collect_context_files(context, project_path)
 
         try:
-            result = self.service.generate_spec(
-                spec_name=self.spec_name,
-                design_doc=design_doc,
-                project_path=project_path,
-                context_files=context_files,
-                save_to_disk=False
-            )
+            if ensemble_size > 1:
+                result = self.service.generate_spec_ensemble(
+                    spec_name=self.spec_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    ensemble_size=ensemble_size,
+                    save_to_disk=False
+                )
+            else:
+                result = self.service.generate_spec(
+                    spec_name=self.spec_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    save_to_disk=False
+                )
 
             if result.success:
                 return StepResult.success(
@@ -265,14 +289,16 @@ class GenerateTestStep(WorkflowStep):
     description = "Generate a P test driver file"
     max_retries = 3
 
-    def __init__(self, generation_service: GenerationService, test_name: str = "TestDriver"):
+    def __init__(self, generation_service: GenerationService, test_name: str = "TestDriver", ensemble_size: int = 3):
         self.service = generation_service
         self.test_name = test_name
+        self.ensemble_size = ensemble_size
         self.name = f"generate_test_{test_name}"
 
     def execute(self, context: Dict[str, Any]) -> StepResult:
         design_doc = context.get("design_doc")
         project_path = context.get("project_path")
+        ensemble_size = context.get("ensemble_size", self.ensemble_size)
 
         if not design_doc:
             return StepResult.failure("design_doc is required")
@@ -283,13 +309,23 @@ def execute(self, context: Dict[str, Any]) -> StepResult:
         context_files = self._collect_all_context(context, project_path)
 
         try:
-            result = self.service.generate_test(
-                test_name=self.test_name,
-                design_doc=design_doc,
-                project_path=project_path,
-                context_files=context_files,
-                save_to_disk=False
-            )
+            if ensemble_size > 1:
+                result = self.service.generate_test_ensemble(
+                    test_name=self.test_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    ensemble_size=ensemble_size,
+                    save_to_disk=False
+                )
+            else:
+                result = self.service.generate_test(
+                    test_name=self.test_name,
+                    design_doc=design_doc,
+                    project_path=project_path,
+                    context_files=context_files,
+                    save_to_disk=False
+                )
 
             if result.success:
                 return StepResult.success(

Src/PeasyAI/src/rag/create_rag_index.py

@@ -9,7 +9,7 @@
 import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 import pickle
 import sys
 from datetime import datetime