Merge pull request #6 from harshalmore31/main

Refactor README and module documentation; and ensured type safe methods

.gitignore

@@ -2,6 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
+agent_workspace
 
 # C extensions
 *.so

README.md

@@ -1,72 +1,118 @@
-# Open-MAI-Dx-Orchestrator
+# MAI Diagnostic Orchestrator (MAI-DxO)
 
-> **An open-source implementation of the "Sequential Diagnosis with Language Models" paper by Microsoft Research, built with the Swarms AI framework.**
+> **An open-source implementation of Microsoft Research's "Sequential Diagnosis with Language Models" paper, built with the Swarms AI framework.**
+
+MAI-DxO (MAI Diagnostic Orchestrator) is a sophisticated AI-powered diagnostic system that simulates a virtual panel of physician-agents to perform iterative medical diagnosis with cost-effectiveness optimization. This implementation faithfully reproduces the methodology described in the Microsoft Research paper while providing additional features and flexibility.
 
 [![Paper](https://img.shields.io/badge/Paper-arXiv:2506.22405-red.svg)](https://arxiv.org/abs/2506.22405)
 [![License](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
 [![Python](https://img.shields.io/badge/Python-3.8+-green.svg)](https://python.org)
 
-MAI-DxO (MAI Diagnostic Orchestrator) is a sophisticated AI-powered diagnostic system that simulates a virtual panel of physician-agents to perform iterative medical diagnosis with cost-effectiveness optimization. This implementation faithfully reproduces the methodology described in the Microsoft Research paper while providing additional features and flexibility.
+## ✨ Key Features
+
+- **8 AI Physician Agents**: Specialized roles for comprehensive diagnosis.
+- **5 Operational Modes**: Instant, question-only, budgeted, no-budget, and ensemble modes.
+- **Cost Tracking**: Real-time budget monitoring with costs for 25+ medical tests.
+- **Clinical Evaluation**: 5-point accuracy scoring with detailed feedback.
+- **Model Agnostic**: Works with GPT, Gemini, Claude, and other leading LLMs.
 
 ## 🚀 Quick Start
 
+### 1. Installation
+
+Install the package directly via pip:
+
 ```bash
-# Install the package
 pip install mai-dx
+```
+
+Or, for development, clone the repository and install the requirements:
 
-# Or install from source
+```bash
 git clone https://github.com/The-Swarm-Corporation/Open-MAI-Dx-Orchestrator.git
 cd Open-MAI-Dx-Orchestrator
-pip install -e .
+pip install -r requirements.txt
 ```
 
----- 
+### 2. Environment Setup
 
-## Enivronement Configuration
-Configure your api keys and environment variables like the following. 
+Create a `.env` file in your project root and add your API keys:
 
 ```txt
-WORKSPACE_DIR="" # for the swarms library
-OPENAI_API_KEY="" # Your model api key
-GEMINI_API_KEY="" # your gemini api key
-ANTHROPIC_API_KEY=""
+OPENAI_API_KEY="Your OpenAI API key"
+GEMINI_API_KEY="Your Gemini API key"
+ANTHROPIC_API_KEY="Your Anthropic API key"
 ```
 
-----
-
-## Example
+### 3. Basic Usage
 
 ```python
 from mai_dx import MaiDxOrchestrator
 
-# Create orchestrator
-orchestrator = MaiDxOrchestrator(model_name="gemini/gemini-2.5-flash")
+# Create the orchestrator (defaults to a capable model)
+orchestrator = MaiDxOrchestrator()
 
-# Run diagnosis
+# Run a diagnosis
 result = orchestrator.run(
     initial_case_info="29-year-old woman with sore throat and peritonsillar swelling...",
     full_case_details="Patient: 29-year-old female. History: Onset of sore throat...",
     ground_truth_diagnosis="Embryonal rhabdomyosarcoma of the pharynx"
 )
 
-print(f"Diagnosis: {result.final_diagnosis}")
+# Print the results
+print(f"Final Diagnosis: {result.final_diagnosis}")
 print(f"Accuracy: {result.accuracy_score}/5.0")
-print(f"Cost: ${result.total_cost:,}")
+print(f"Total Cost: ${result.total_cost:,.2f}")
+```
+
+## ⚙️ Advanced Usage & Configuration
+
+Customize the orchestrator's model, budget, and operational mode.
+
+```python
+from mai_dx import MaiDxOrchestrator
+
+# Configure with a specific model and budget
+orchestrator = MaiDxOrchestrator(
+    model_name="gemini/gemini-2.5-flash",  # or "gpt-4", "claude-3-5-sonnet"
+    max_iterations=10,
+    initial_budget=3000,
+    mode="budgeted"  # Other modes: "instant", "question_only", "no_budget"
+)
+
+# Run the diagnosis
+# ...
 ```
 
---- 
+## 🏥 How It Works: The Virtual Physician Panel
+
+MAI-DxO employs a multi-agent system where each agent has a specific role:
+
+- **🧠 Dr. Hypothesis**: Maintains the differential diagnosis.
+- **🔬 Dr. Test-Chooser**: Selects the most cost-effective diagnostic tests.
+- **🤔 Dr. Challenger**: Prevents cognitive biases and diagnostic errors.
+- **💰 Dr. Stewardship**: Ensures cost-effective care.
+- **✅ Dr. Checklist**: Performs quality control checks.
+- **🤝 Consensus Coordinator**: Synthesizes panel decisions.
+- **🔑 Gatekeeper**: Acts as the clinical information oracle.
+- **⚖️ Judge**: Evaluates the final diagnostic accuracy.
+
 
 ## Documentation
 
 Learn more about this repository [with the docs](DOCS.md)
 
+## 🤝 Contributing
+
+We welcome contributions! Please feel free to open an issue or submit a pull request.
+
 ## 📄 License
 
-This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
 
 ## 📚 Citation
 
-If you use this implementation in your research, please cite both the original paper and this implementation:
+If you use this work in your research, please cite both the original paper and this software implementation.
 
 ```bibtex
 @misc{nori2025sequentialdiagnosislanguagemodels,
@@ -103,4 +149,4 @@ If you use this implementation in your research, please cite both the original p
 
 <p align="center">
   <strong>Built with Swarms for advancing AI-powered medical diagnosis</strong>
-</p>
+</p>

mai_dx/__init__.py

@@ -1,3 +1,70 @@
-from mai_dx.main import MaiDxOrchestrator, run_mai_dxo_demo
+from typing import Any, Optional
 
-__all__ = ["MaiDxOrchestrator", "run_mai_dxo_demo"]
+from mai_dx.main import (
+    MaiDxOrchestrator,
+    DiagnosisResult,
+    Action,
+    AgentRole,
+    run_mai_dxo_demo
+)
+
+__version__ = "1.0.0"
+__author__ = "The Swarm Corporation"
+__description__ = "AI-powered diagnostic system with virtual physician panels"
+
+# Main exports
+__all__ = [
+    "MaiDxOrchestrator",
+    "DiagnosisResult", 
+    "Action",
+    "AgentRole",
+    "run_mai_dxo_demo"
+]
+
+# Convenience imports for common usage patterns
+def create_orchestrator(
+    model_name: str = "gemini/gemini-2.5-flash",
+    mode: str = "no_budget",
+    **kwargs: Any
+) -> MaiDxOrchestrator:
+    """
+    Convenience function to create a MAI-DxO orchestrator.
+    
+    Args:
+        model_name: Language model to use
+        mode: Operational mode
+        **kwargs: Additional configuration parameters
+        
+    Returns:
+        Configured MaiDxOrchestrator instance
+    """
+    if mode in ["instant", "question_only", "budgeted", "no_budget", "ensemble"]:
+        return MaiDxOrchestrator.create_variant(mode, model_name=model_name, **kwargs)
+    else:
+        return MaiDxOrchestrator(model_name=model_name, mode=mode, **kwargs)
+
+
+def quick_diagnosis(
+    case_info: str,
+    case_details: str,
+    ground_truth: Optional[str] = None,
+    model_name: str = "gemini/gemini-2.5-flash"
+) -> DiagnosisResult:
+    """
+    Convenience function for quick diagnosis without configuration.
+    
+    Args:
+        case_info: Initial case presentation
+        case_details: Complete case information
+        ground_truth: Correct diagnosis (optional)
+        model_name: Model to use
+        
+    Returns:
+        DiagnosisResult with diagnosis and evaluation
+    """
+    orchestrator = MaiDxOrchestrator(model_name=model_name, max_iterations=5)
+    return orchestrator.run(
+        initial_case_info=case_info,
+        full_case_details=case_details,
+        ground_truth_diagnosis=ground_truth or "Unknown"
+    )

mai_dx/main.py

@@ -18,7 +18,7 @@ Key Features:
 
 Example Usage:
     # Standard MAI-DxO usage
-    orchestrator = MaiDxOrchestrator(model_name="gemini/gemini-2.5-flash")
+    orchestrator = MaiDxOrchestrator(model_name="gpt-4.1")
     result = orchestrator.run(initial_case_info, full_case_details, ground_truth)
 
     # Budget-constrained variant
@@ -131,7 +131,7 @@ class MaiDxOrchestrator:
 
     def __init__(
         self,
-        model_name: str = "gemini/gemini-2.5-flash",
+        model_name: str = "gpt-4.1",
         max_iterations: int = 10,
         initial_budget: int = 10000,
         mode: str = "no_budget",  # "instant", "question_only", "budgeted", "no_budget", "ensemble"
@@ -198,7 +198,7 @@ class MaiDxOrchestrator:
             f"🏥 MAI Diagnostic Orchestrator initialized successfully in '{mode}' mode with budget ${initial_budget:,}"
         )
 
-    def _init_agents(self):
+    def _init_agents(self) -> None:
         """Initializes all required agents with their specific roles and prompts."""
         self.agents = {
             role: Agent(
@@ -1387,11 +1387,11 @@ def run_mai_dxo_demo(
                 orchestrator = MaiDxOrchestrator.create_variant(
                     variant,
                     budget=3000,
-                    model_name="gemini/gemini-2.5-flash",
+                    model_name="gpt-4.1",
                 )
             else:
                 orchestrator = MaiDxOrchestrator.create_variant(
-                    variant, model_name="gemini/gemini-2.5-flash"
+                    variant, model_name="gpt-4.1"
                 )
 
             result = orchestrator.run(
@@ -1466,13 +1466,13 @@ def run_mai_dxo_demo(
 #                 orchestrator = MaiDxOrchestrator.create_variant(
 #                     variant_name,
 #                     budget=3000,
-#                     model_name="gemini/gemini-2.5-flash",
+#                     model_name="gpt-4.1",
 #                     max_iterations=5,
 #                 )
 #             else:
 #                 orchestrator = MaiDxOrchestrator.create_variant(
 #                     variant_name,
-#                     model_name="gemini/gemini-2.5-flash",
+#                     model_name="gpt-4.1",
 #                     max_iterations=5,
 #                 )
 
@@ -1504,7 +1504,7 @@ def run_mai_dxo_demo(
 
 #         ensemble_orchestrator = MaiDxOrchestrator.create_variant(
 #             "ensemble",
-#             model_name="gemini/gemini-2.5-flash",
+#             model_name="gpt-4.1",
 #             max_iterations=3,  # Shorter iterations for ensemble
 #         )